Refactor TFDS CLI commands.

PiperOrigin-RevId: 793983080

Author: fineguy

tensorflow_datasets/scripts/cli/build.py

@@ -15,7 +15,6 @@
 
 """`tfds build` command."""
 
-import argparse
 from collections.abc import Iterator
 import dataclasses
 import functools
@@ -24,7 +23,6 @@
 import json
 import multiprocessing
 import os
-import typing
 from typing import Any, Type
 
 from absl import logging
@@ -34,8 +32,8 @@
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class Args:
-  """CLI arguments for building datasets.
+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
@@ -120,16 +118,6 @@ def execute(self) -> None:
         pool.map(process_builder_fn, builders)
 
 
-def register_subparser(parsers: argparse._SubParsersAction) -> None:  # pylint: disable=protected-access
-  """Add subparser for `build` command."""
-  parser = parsers.add_parser(
-      'build', help='Commands for downloading and preparing datasets.'
-  )
-  parser = typing.cast(simple_parsing.ArgumentParser, parser)
-  parser.add_arguments(Args, dest='args')
-  parser.set_defaults(subparser_fn=lambda args: args.args.execute())
-
-
 def _make_builders(
     args: Args,
     builder_cls: Type[tfds.core.DatasetBuilder],

tensorflow_datasets/scripts/cli/build_test.py

@@ -19,6 +19,7 @@
 import functools
 import multiprocessing
 import os
+import typing
 from unittest import mock
 
 from etils import epath
@@ -311,7 +312,7 @@ def test_download_only(build):
 )
 def test_make_download_config(args: str, download_config_kwargs):
   args = main._parse_flags(f'tfds build x {args}'.split())
-  cmd_args: build_lib.Args = args.args
+  cmd_args = typing.cast(build_lib.Args, args.command)
   actual = build_lib._make_download_config(cmd_args, dataset_name='x')
   # Ignore the beam runner
   actual = actual.replace(beam_runner=None)

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -15,11 +15,13 @@
 
 """Utility functions for TFDS CLI."""
 
+import abc
 import argparse
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 import dataclasses
 import itertools
 import pathlib
+from typing import TypeVar
 
 from absl import logging
 from absl.flags import argparse_flags
@@ -33,6 +35,8 @@
 from tensorflow_datasets.core.utils import file_utils
 from tensorflow_datasets.scripts.utils import flag_utils
 
+_DataclassT = TypeVar('_DataclassT')
+
 
 class ArgumentParser(
     argparse_flags.ArgumentParser, simple_parsing.ArgumentParser
@@ -77,6 +81,33 @@ def parse_known_args(
     return super().parse_known_args(args, namespace)
 
 
+def make_flags_parser(
+    args_dataclass: type[_DataclassT], description: str
+) -> Callable[[list[str]], _DataclassT]:
+  """Returns a function that parses flags and returns the dataclass instance."""
+
+  def _parse_flags(argv: list[str]) -> _DataclassT:
+    """Command lines flag parsing."""
+    parser = ArgumentParser(
+        description=description,
+        allow_abbrev=False,
+    )
+    parser.add_arguments(args_dataclass, dest='args')
+    return parser.parse_args(argv[1:]).args
+
+  return _parse_flags
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Args(abc.ABC):
+  """CLI arguments for TFDS CLI commands."""
+
+  @abc.abstractmethod
+  def execute(self) -> None:
+    """Execute the CLI command."""
+    ...
+
+
 @dataclasses.dataclass
 class DatasetInfo:
   """Structure for common string used for formatting.

tensorflow_datasets/scripts/cli/convert_format.py

@@ -25,19 +25,18 @@
 ```
 """
 
-import argparse
 import dataclasses
-import typing
 
 from etils import epath
 import simple_parsing
 from tensorflow_datasets.core import file_adapters
+from tensorflow_datasets.scripts.cli import cli_utils
 from tensorflow_datasets.scripts.cli import convert_format_utils
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class Args:
-  """CLI arguments for converting datasets from one file format to another.
+class Args(cli_utils.Args):
+  """Converts a dataset from one file format to another format.
 
   Attributes:
     root_data_dir: Root data dir that contains all datasets. All datasets and
@@ -94,14 +93,3 @@ def execute(self) -> None:
         num_workers=self.num_workers,
         fail_on_error=not self.only_log_errors,
     )
-
-
-def register_subparser(parsers: argparse._SubParsersAction) -> None:
-  """Add subparser for `convert_format` command."""
-  parser = parsers.add_parser(
-      'convert_format',
-      help='Converts a dataset from one file format to another format.',
-  )
-  parser = typing.cast(simple_parsing.ArgumentParser, parser)
-  parser.add_arguments(Args, dest='args')
-  parser.set_defaults(subparser_fn=lambda args: args.args.execute())

tensorflow_datasets/scripts/cli/croissant.py

@@ -26,11 +26,9 @@
 ```
 """
 
-import argparse
 import dataclasses
 import functools
 import json
-import typing
 
 from etils import epath
 import mlcroissant as mlc
@@ -43,8 +41,8 @@
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class CmdArgs(simple_parsing.helpers.FrozenSerializable):
-  """CLI arguments for preparing a Croissant dataset.
+class CmdArgs(simple_parsing.helpers.FrozenSerializable, cli_utils.Args):
+  """Prepares a Croissant dataset.
 
   Attributes:
     jsonld: Path to the JSONLD file.
@@ -122,18 +120,10 @@ def version(self) -> version_lib.Version:
         self.overwrite_version or self.dataset.metadata.version or '1.0.0'
     )
 
-
-def register_subparser(parsers: argparse._SubParsersAction):
-  """Add subparser for `convert_format` command."""
-  parser = parsers.add_parser(
-      'build_croissant',
-      help='Prepares a croissant dataset',
-  )
-  parser = typing.cast(simple_parsing.ArgumentParser, parser)
-  parser.add_arguments(CmdArgs, dest='args')
-  parser.set_defaults(
-      subparser_fn=lambda args: prepare_croissant_builders(args.args)
-  )
+  def execute(self) -> None:
+    """Creates Croissant Builders and prepares them."""
+    for record_set_id in self.record_set_ids:
+      prepare_croissant_builder(args=self, record_set_id=record_set_id)
 
 
 def prepare_croissant_builder(
@@ -163,14 +153,3 @@ def prepare_croissant_builder(
       beam_pipeline_options=None,
   )
   return builder
-
-
-def prepare_croissant_builders(args: CmdArgs):
-  """Creates Croissant Builders and prepares them.
-
-  Args:
-    args: CLI arguments.
-  """
-  # Generate each config sequentially.
-  for record_set_id in args.record_set_ids:
-    prepare_croissant_builder(args=args, record_set_id=record_set_id)

tensorflow_datasets/scripts/cli/main.py

@@ -21,13 +21,13 @@
 See: https://www.tensorflow.org/datasets/cli
 """
 
-import argparse
+import dataclasses
 import logging as python_logging
-from typing import List
 
 from absl import app
 from absl import flags
 from absl import logging
+import simple_parsing
 
 import tensorflow_datasets.public_api as tfds
 
@@ -41,33 +41,60 @@
 FLAGS = flags.FLAGS
 
 
-def _parse_flags(argv: List[str]) -> argparse.Namespace:
-  """Command lines flag parsing."""
-  parser = cli_utils.ArgumentParser(
-      description='Tensorflow Datasets CLI tool',
-      allow_abbrev=False,
-  )
-  parser.add_argument(
-      '--version',
-      action='version',
-      version='TensorFlow Datasets: ' + tfds.__version__,
-  )
-  parser.add_argument(
-      '--dry_run',
-      action='store_true',
-      help='If True, print the parsed arguments.',
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class _DummyCommand:
+  """Dummy command to avoid `command is MISSING` error."""
+
+  pass
+
+
+version_field = simple_parsing.field(
+    action='version',
+    version='TensorFlow Datasets: ' + tfds.__version__,
+    help='The version of the TensorFlow Datasets package.',
+)
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Args(cli_utils.Args):
+  """Tensorflow Datasets CLI tool."""
+
+  version: str = version_field
+  """The version of the TensorFlow Datasets package."""
+
+  dry_run: bool = simple_parsing.flag(default=False)
+  """If True, print the parsed arguments and exit."""
+
+  command: build.Args | new.Args | convert_format.Args | croissant.CmdArgs = (
+      simple_parsing.subparsers(
+          {
+              'build': build.Args,
+              'new': new.Args,
+              'convert_format': convert_format.Args,
+              'build_croissant': croissant.CmdArgs,
+          },
+          default_factory=_DummyCommand,
+      )
   )
-  parser.set_defaults(subparser_fn=lambda _: parser.print_help())
-  # Register sub-commands
-  subparser = parser.add_subparsers(title='command')
-  build.register_subparser(subparser)
-  new.register_subparser(subparser)
-  convert_format.register_subparser(subparser)
-  croissant.register_subparser(subparser)
-  return parser.parse_args(argv[1:])
+  """The command to execute."""
+
+  def execute(self) -> None:
+    """Run the command."""
+    if self.dry_run:
+      print(self)
+    # When no command is given, print the help message.
+    elif isinstance(self.command, _DummyCommand):
+      _parse_flags(['', '--help'])
+    else:
+      self.command.execute()
+
+
+_parse_flags = cli_utils.make_flags_parser(
+    Args, description='Tensorflow Datasets CLI tool'
+)
 
 
-def main(args: argparse.Namespace) -> None:
+def main(args: Args) -> None:
 
   # From the CLI, all datasets are visible
   tfds.core.visibility.set_availables([
@@ -98,11 +125,7 @@ def main(args: argparse.Namespace) -> None:
     new_stream = tfds.core.utils.tqdm_utils.TqdmStream()
     python_handler.setStream(new_stream)
 
-  if args.dry_run:
-    print(args)
-  else:
-    # Launch the subcommand defined in the subparser (or default to print help)
-    args.subparser_fn(args)
+  args.execute()
 
 
 def launch_cli() -> None:

tensorflow_datasets/scripts/cli/new.py

@@ -15,13 +15,11 @@
 
 """`tfds new` command."""
 
-import argparse
 import dataclasses
 import os
 import pathlib
 import subprocess
 import textwrap
-import typing
 
 import simple_parsing
 from tensorflow_datasets.core import constants
@@ -33,8 +31,8 @@
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class Args:
-  """CLI arguments for creating a new dataset directory.
+class Args(utils.Args):
+  """Creates a new dataset directory from the template.
 
   Attributes:
     dataset_name: Name of the dataset to be created (in snake_case).
@@ -71,17 +69,6 @@ def execute(self) -> None:
     )
 
 
-def register_subparser(parsers: argparse._SubParsersAction) -> None:
-  """Add subparser for `new` command."""
-  parser = parsers.add_parser(
-      'new',
-      help='Creates a new dataset directory from the template.',
-  )
-  parser = typing.cast(simple_parsing.ArgumentParser, parser)
-  parser.add_arguments(Args, dest='args')
-  parser.set_defaults(subparser_fn=lambda args: args.args.execute())
-
-
 def create_dataset_files(
     dataset_name: str,
     dataset_dir: pathlib.Path,