Use dataclasses for CLI arguments in scripts/cli/croissant.py.

fineguy · The TensorFlow Datasets Authors · commit a586b8e80af0 · 2024-05-14T04:58:27.000-07:00
PiperOrigin-RevId: 633542690
diff --git a/setup.py b/setup.py
@@ -73,6 +73,7 @@
     'psutil',
     'pyarrow',
     'requests>=2.19.0',
+    'simple_parsing',
     'tensorflow-metadata',
     'termcolor',
     'toml',
diff --git a/tensorflow_datasets/scripts/cli/croissant.py b/tensorflow_datasets/scripts/cli/croissant.py
@@ -21,101 +21,38 @@
   --jsonld=/tmp/croissant.json \
   --data_dir=/tmp/foo \
   --file_format=array_record \
-  --record_sets=record1 --record_sets=record2 \
+  --record_sets=record1,record2 \
   --mapping='{"document.csv": "~/Downloads/document.csv"}"'
 ```
 """
 
 import argparse
-from collections.abc import Sequence
+import dataclasses
 import json
+import typing
 
 from etils import epath
+import simple_parsing
+from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.core.dataset_builders import croissant_builder
 from tensorflow_datasets.scripts.cli import cli_utils
 
 
-def add_parser_arguments(parser: argparse.ArgumentParser):
-  """Add arguments for `build_croissant` subparser."""
-  parser.add_argument(
-      '--jsonld',
-      type=str,
-      help='The Croissant config file for the given dataset.',
-      required=True,
-  )
-  parser.add_argument(
-      '--record_sets',
-      nargs='*',
-      help=(
-          'The names of the record sets to generate. Each record set will'
-          ' correspond to a separate config. If not specified, it will use all'
-          ' the record sets'
-      ),
-  )
-  parser.add_argument(
-      '--mapping',
-      type=str,
-      help=(
-          'Mapping filename->filepath as a Python dict[str, str] to handle'
-          ' manual downloads. If `document.csv` is the FileObject and you'
-          ' downloaded it to `~/Downloads/document.csv`, you can'
-          ' specify`--mapping=\'{"document.csv": "~/Downloads/document.csv"}\''
-      ),
-  )
-
-  cli_utils.add_debug_argument_group(parser)
-  cli_utils.add_path_argument_group(parser)
-  cli_utils.add_generation_argument_group(parser)
-  cli_utils.add_publish_argument_group(parser)
+@dataclasses.dataclass
+class CmdArgs:
+  """CLI arguments for preparing a Croissant dataset.
 
-
-def register_subparser(parsers: argparse._SubParsersAction):
-  """Add subparser for `convert_format` command."""
-  parser = parsers.add_parser(
-      'build_croissant',
-      help='Prepares a croissant dataset',
-  )
-  add_parser_arguments(parser)
-  parser.set_defaults(
-      subparser_fn=lambda args: prepare_croissant_builder(
-          jsonld=args.jsonld,
-          data_dir=args.data_dir,
-          file_format=args.file_format,
-          record_sets=args.record_sets,
-          mapping=args.mapping,
-          download_dir=args.download_dir,
-          publish_dir=args.publish_dir,
-          skip_if_published=args.skip_if_published,
-          overwrite=args.overwrite,
-      )
-  )
-
-
-def prepare_croissant_builder(
-    jsonld: epath.PathLike,
-    data_dir: epath.PathLike,
-    file_format: str,
-    record_sets: Sequence[str],
-    mapping: str | None,
-    download_dir: epath.PathLike | None,
-    publish_dir: epath.PathLike | None,
-    skip_if_published: bool,
-    overwrite: bool,
-) -> None:
-  # pyformat: disable
-  """Creates a Croissant Builder and runs the preparation.
-
-  Args:
-    jsonld: The Croissant config file for the given dataset
+  Attributes:
+    jsonld: Path to the JSONLD file.
     data_dir: Path where the converted dataset will be stored.
     file_format: File format to convert the dataset to.
-    record_sets: The `@id`s of the record sets to generate. Each record set will
+    record_sets: The names of the record sets to generate. Each record set will
       correspond to a separate config. If not specified, it will use all the
-      record sets
+      record sets.
     mapping: Mapping filename->filepath as a Python dict[str, str] to handle
       manual downloads. If `document.csv` is the FileObject and you downloaded
       it to `~/Downloads/document.csv`, you can specify
-      `mapping={"document.csv": "~/Downloads/document.csv"}`.,
+      `--mapping='{"document.csv": "~/Downloads/document.csv"}'`
     download_dir: Where to place downloads. Default to `<data_dir>/downloads/`.
     publish_dir: Where to optionally publish the dataset after it has been
       generated successfully. Should be the root data dir under which datasets
@@ -124,29 +61,74 @@ def prepare_croissant_builder(
       already published, then it will not be regenerated.
     overwrite: Delete pre-existing dataset if it exists.
   """
-  # pyformat: enable
-  if not record_sets:
-    record_sets = None
 
-  if mapping:
+  jsonld: epath.PathLike
+  data_dir: epath.PathLike
+  # Need to override the default use of `Enum.name` for choice options.
+  file_format: str = simple_parsing.choice(
+      *(file_format.value for file_format in file_adapters.FileFormat),
+      default=file_adapters.FileFormat.ARRAY_RECORD.value,
+  )
+  # Need to manually parse comma-separated list of values, see:
+  # https://github.com/lebrice/SimpleParsing/issues/142.
+  record_sets: list[str] = simple_parsing.field(
+      default_factory=list,
+      type=lambda record_sets_str: record_sets_str.split(','),
+      nargs='?',
+  )
+  mapping: str | None = None
+  download_dir: epath.PathLike | None = None
+  publish_dir: epath.PathLike | None = None
+  skip_if_published: bool = False
+  overwrite: bool = False
+
+
+def register_subparser(parsers: argparse._SubParsersAction):
+  """Add subparser for `convert_format` command."""
+  orig_parser_class = parsers._parser_class  # pylint: disable=protected-access
+  try:
+    parsers._parser_class = simple_parsing.ArgumentParser  # pylint: disable=protected-access
+    parser = parsers.add_parser(
+        'build_croissant',
+        help='Prepares a croissant dataset',
+    )
+    parser = typing.cast(simple_parsing.ArgumentParser, parser)
+  finally:
+    parsers._parser_class = orig_parser_class  # pylint: disable=protected-access
+  parser.add_arguments(CmdArgs, dest='args')
+  parser.set_defaults(
+      subparser_fn=lambda args: prepare_croissant_builder(args.args)
+  )
+
+
+def prepare_croissant_builder(args: CmdArgs) -> None:
+  """Creates a Croissant Builder and runs the preparation.
+
+  Args:
+    args: CLI arguments.
+  """
+  if args.mapping:
     try:
-      mapping = json.loads(mapping)
+      mapping = json.loads(args.mapping)
     except json.JSONDecodeError as e:
-      raise ValueError(f'Error parsing mapping parameter: {mapping}') from e
+      raise ValueError(
+          f'Error parsing mapping parameter: {args.mapping}'
+      ) from e
+  else:
+    mapping = None
 
   builder = croissant_builder.CroissantBuilder(
-      jsonld=jsonld,
-      record_set_ids=record_sets,
-      file_format=file_format,
-      data_dir=data_dir,
+      jsonld=args.jsonld,
+      record_set_ids=args.record_sets or None,
+      file_format=args.file_format,
+      data_dir=args.data_dir,
       mapping=mapping,
   )
   cli_utils.download_and_prepare(
       builder=builder,
       download_config=None,
-      download_dir=epath.Path(download_dir) if download_dir else None,
-      publish_dir=epath.Path(publish_dir) if publish_dir else None,
-      skip_if_published=skip_if_published,
-      freeze_files=freeze_files,
-      overwrite=overwrite,
+      download_dir=epath.Path(args.download_dir) if args.download_dir else None,
+      publish_dir=epath.Path(args.publish_dir) if args.publish_dir else None,
+      skip_if_published=args.skip_if_published,
+      overwrite=args.overwrite,
   )
diff --git a/tensorflow_datasets/scripts/prepare_croissant.py b/tensorflow_datasets/scripts/prepare_croissant.py
@@ -26,76 +26,13 @@
 ```
 """
 
-import dataclasses
-
 from absl import app
 from etils import eapp
-from etils import epath
-import simple_parsing
-from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.scripts.cli import croissant
 
 
-@dataclasses.dataclass
-class CmdArgs:
-  """CLI arguments for preparing a Croissant dataset.
-
-  Attributes:
-    jsonld: Path to the JSONLD file.
-    data_dir: Path where the converted dataset will be stored.
-    file_format: File format to convert the dataset to.
-    record_sets: The names of the record sets to generate. Each record set will
-      correspond to a separate config. If not specified, it will use all the
-      record sets.
-    mapping: Mapping filename->filepath as a Python dict[str, str] to handle
-      manual downloads. If `document.csv` is the FileObject and you downloaded
-      it to `~/Downloads/document.csv`, you can specify
-      `--mapping='{"document.csv": "~/Downloads/document.csv"}'`
-    download_dir: Where to place downloads. Default to `<data_dir>/downloads/`.
-    publish_dir: Where to optionally publish the dataset after it has been
-      generated successfully. Should be the root data dir under which datasets
-      are stored. If unspecified, dataset will not be published.
-    skip_if_published: If the dataset with the same version and config is
-      already published, then it will not be regenerated.
-    overwrite: Delete pre-existing dataset if it exists.
-  """
-
-  jsonld: epath.PathLike
-  data_dir: epath.PathLike
-  # Need to override the default use of `Enum.name` for choice options.
-  file_format: str = simple_parsing.choice(
-      *(file_format.value for file_format in file_adapters.FileFormat),
-      default=file_adapters.FileFormat.ARRAY_RECORD.value,
-  )
-  # Need to manually parse comma-separated list of values, see:
-  # https://github.com/lebrice/SimpleParsing/issues/142.
-  record_sets: list[str] = simple_parsing.field(
-      default_factory=list,
-      type=lambda record_sets_str: record_sets_str.split(','),
-      nargs='?',
-  )
-  mapping: str | None = None
-  download_dir: epath.PathLike | None = None
-  publish_dir: epath.PathLike | None = None
-  skip_if_published: bool = False
-  overwrite: bool = False
-
-parse_flags = eapp.make_flags_parser(CmdArgs)
-
-
-def main(args: CmdArgs):
-  croissant.prepare_croissant_builder(
-      jsonld=args.jsonld,
-      data_dir=args.data_dir,
-      file_format=args.file_format,
-      record_sets=args.record_sets,
-      mapping=args.mapping,
-      download_dir=args.download_dir,
-      publish_dir=args.publish_dir,
-      skip_if_published=args.skip_if_published,
-      overwrite=args.overwrite,
-  )
+parse_flags = eapp.make_flags_parser(croissant.CmdArgs)
 
 
 if __name__ == '__main__':
-  app.run(main, flags_parser=parse_flags)
+  app.run(croissant.prepare_croissant_builder, flags_parser=parse_flags)
