Refactor download_and_prepare script.

PiperOrigin-RevId: 795187236

Author: fineguy

tensorflow_datasets/scripts/cli/build.py

@@ -33,42 +33,42 @@
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class Args(cli_utils.Args):
-  """Commands for downloading and preparing datasets.
-
-  Attributes:
-    positional_datasets: Name(s) of the dataset(s) to build. Default to current
-      dir. See https://www.tensorflow.org/datasets/cli for accepted values.
-    datasets: Datasets can also be provided as keyword argument.
-    debug: Debug & tests options. Use --pdb to enter post-mortem debugging mode
-      if an exception is raised.
-    paths: Path options.
-    generation: Generation options.
-    publishing: Publishing options.
-    automation: Automation options.
-  """
+  """Commands for downloading and preparing datasets."""
 
-  positional_datasets: list[str] = simple_parsing.field(
+  # Name(s) of the dataset(s) to build. Default to current dir. See
+  # https://www.tensorflow.org/datasets/cli for accepted values.
+  positional_datasets: list[str] = simple_parsing.list_field(
       positional=True,
-      nargs='*',
-      default_factory=list,
       # Need to explicitly set metavar for command-line help.
       metavar='datasets',
   )
-  datasets: list[str] = simple_parsing.field(nargs='*', default_factory=list)
+
+  datasets: list[str] = simple_parsing.list_field(alias='--dataset')
+  """Datasets can also be provided as keyword argument."""
 
   debug: cli_utils.DebugOptions = cli_utils.DebugOptions()
+  """Debug & tests options. Use --pdb to enter post-mortem debugging mode if an
+  exception is raised."""
+
   paths: cli_utils.PathOptions = simple_parsing.field(
       default_factory=cli_utils.PathOptions
   )
+  """Path options."""
+
   generation: cli_utils.GenerationOptions = simple_parsing.field(
       default_factory=cli_utils.GenerationOptions
   )
+  """Generation options."""
+
   publishing: cli_utils.PublishingOptions = simple_parsing.field(
       default_factory=cli_utils.PublishingOptions
   )
+  """Publishing options."""
+
   automation: cli_utils.AutomationOptions = simple_parsing.field(
       default_factory=cli_utils.AutomationOptions
   )
+  """Automation options."""
 
   def execute(self) -> None:
     """Build the given datasets."""

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -243,10 +243,14 @@ class GenerationOptions:
 
   download_only: bool = simple_parsing.flag(default=False)
   config: str | None = simple_parsing.field(default=None, alias='-c')
-  config_idx: int | None = None
+  config_idx: int | None = simple_parsing.field(
+      default=None, alias='--builder_config_id'
+  )
   update_metadata_only: bool = simple_parsing.flag(default=False)
   download_config: str | None = None
-  imports: str | None = simple_parsing.field(default=None, alias='-i')
+  imports: str | None = simple_parsing.field(
+      default=None, alias=['-i', '--module_import']
+  )
   register_checksums: bool = simple_parsing.flag(default=False)
   force_checksums_validation: bool = simple_parsing.flag(default=False)
   beam_pipeline_options: str | None = None

tensorflow_datasets/scripts/download_and_prepare.py

@@ -15,43 +15,21 @@
 
 r"""Wrapper around `tfds build`."""
 
-import typing
-
 from absl import app
 from absl import flags
-from absl import logging
-from tensorflow_datasets.scripts.cli import build
 from tensorflow_datasets.scripts.cli import main as main_cli
 
-module_import = flags.DEFINE_string('module_import', None, '`--imports` flag.')
-dataset = flags.DEFINE_string('dataset', None, 'singleton `--datasets` flag.')
-
-builder_config_id = flags.DEFINE_integer(
-    'builder_config_id', None, '`--config_idx` flag'
-)
-
-
 
 def _parse_flags(argv: list[str]) -> main_cli.Args:
   """Command lines flag parsing."""
   return main_cli._parse_flags([argv[0], 'build'] + argv[1:])  # pylint: disable=protected-access
 
 
-_display_warning = True
-
-
 def main(args: main_cli.Args) -> None:
-  if _display_warning:
-    logging.warning(
-        '***`tfds build` should be used instead of `download_and_prepare`.***'
-    )
-  cmd_args = typing.cast(build.Args, args.command)
-  if module_import.value:
-    cmd_args.generation.imports = module_import.value
-  if dataset.value:
-    cmd_args.datasets = [dataset.value]
-  if builder_config_id.value is not None:
-    cmd_args.generation.config_idx = builder_config_id.value
+  from absl import logging
+  logging.warning(
+      '***`tfds build` should be used instead of `download_and_prepare`.***'
+  )
   main_cli.main(args)