Fix parsing CLI arguments and ABSL flags.

PiperOrigin-RevId: 792554983

Author: fineguy

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -16,18 +16,65 @@
 """Utility functions for TFDS CLI."""
 
 import argparse
+from collections.abc import Sequence
 import dataclasses
 import itertools
 import os
 import pathlib
 
 from absl import logging
+from absl.flags import argparse_flags
 from etils import epath
+import simple_parsing
 from tensorflow_datasets.core import dataset_builder
 from tensorflow_datasets.core import download
 from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.core import naming
 from tensorflow_datasets.core.utils import file_utils
+from tensorflow_datasets.scripts.utils import flag_utils
+
+
+class ArgumentParser(
+    argparse_flags.ArgumentParser, simple_parsing.ArgumentParser
+):
+  """An `ArgumentParser` that handles both `simple_parsing` and `absl` flags.
+
+  This class is a workaround for the fact that `simple_parsing.ArgumentParser`
+  does not natively handle `absl.flags`. Without this, `absl` flags are not
+  correctly parsed, especially when they are mixed with positional arguments,
+  leading to errors.
+
+  The `absl.flags.argparse_flags.ArgumentParser` is designed to integrate `absl`
+  flags into an `argparse` setup. It does this by dynamically adding all
+  defined `absl` flags to the parser instance upon initialization.
+
+  By inheriting from both, we get the features of both:
+  - `simple_parsing.ArgumentParser`: Allows defining arguments from typed
+    dataclasses.
+  - `argparse_flags.ArgumentParser`: Adds support for `absl` flags.
+
+  The Method Resolution Order (MRO) is:
+  `ArgumentParser` -> `argparse_flags.ArgumentParser` ->
+  `simple_parsing.ArgumentParser` -> `argparse.ArgumentParser` -> `object`.
+
+  This order is important. `argparse_flags.ArgumentParser` is first so that it
+  can intercept arguments and handle `absl` flags before they are passed to
+  `simple_parsing.ArgumentParser`.
+  """
+
+  def parse_known_args(
+      self,
+      args: Sequence[str] | None = None,
+      namespace: argparse.Namespace | None = None,
+      attempt_to_reorder: bool = False,
+  ):
+    # `argparse_flags.ArgumentParser` does not support `attempt_to_reorder` that
+    # is used by `simple_parsing.ArgumentParser`. Since we don't need it, we can
+    # just ignore it.
+    del attempt_to_reorder
+    if args:
+      args = flag_utils.normalize_flags(args)
+    return super().parse_known_args(args, namespace)
 
 
 @dataclasses.dataclass

tensorflow_datasets/scripts/cli/conftest.py

@@ -15,10 +15,11 @@
 
 """CLI Fixtures."""
 
-import argparse
 from unittest import mock
 
 import pytest
+import simple_parsing
+from tensorflow_datasets.scripts.cli import cli_utils
 
 
 @pytest.fixture(scope='session', autouse=True)
@@ -29,7 +30,7 @@ def _mock_argparse_flags():
   # another test):
   # `flags.DEFINE_string('data_dir')` with `parser.add_argument('--data_dir')`
   # We patch argparse_flags during test, so absl flags are ignored.
-  with mock.patch(
-      'absl.flags.argparse_flags.ArgumentParser', argparse.ArgumentParser
+  with mock.patch.object(
+      cli_utils, 'ArgumentParser', simple_parsing.ArgumentParser
   ):
     yield

tensorflow_datasets/scripts/cli/main.py

@@ -29,24 +29,21 @@
 from absl import flags
 from absl import logging
 
-import simple_parsing
 import tensorflow_datasets.public_api as tfds
 
 # Import commands
 from tensorflow_datasets.scripts.cli import build
+from tensorflow_datasets.scripts.cli import cli_utils
 from tensorflow_datasets.scripts.cli import convert_format
 from tensorflow_datasets.scripts.cli import croissant
 from tensorflow_datasets.scripts.cli import new
-from tensorflow_datasets.scripts.utils import flag_utils
 
 FLAGS = flags.FLAGS
 
 
 def _parse_flags(argv: List[str]) -> argparse.Namespace:
   """Command lines flag parsing."""
-  argv = flag_utils.normalize_flags(argv)  # See b/174043007 for context.
-
-  parser = simple_parsing.ArgumentParser(
+  parser = cli_utils.ArgumentParser(
       description='Tensorflow Datasets CLI tool',
       allow_abbrev=False,
   )
@@ -67,22 +64,7 @@ def _parse_flags(argv: List[str]) -> argparse.Namespace:
   new.register_subparser(subparser)
   convert_format.register_subparser(subparser)
   croissant.register_subparser(subparser)
-
-  namespace, remaining_argv = parser.parse_known_args(argv[1:])
-
-  # Manually parse absl flags from the remaining arguments.
-  try:
-    # FLAGS requires the program name as the first argument.
-    positionals = FLAGS(argv[:1] + remaining_argv)
-  except flags.Error as e:
-    parser.error(str(e))
-
-  # There should be no positional arguments left, as they should have been
-  # handled by the sub-commands.
-  if len(positionals) > 1:
-    parser.error(f"unrecognized arguments: {' '.join(positionals[1:])}")
-
-  return namespace
+  return parser.parse_args(argv[1:])
 
 
 def main(args: argparse.Namespace) -> None:

tensorflow_datasets/scripts/utils/flag_utils.py

@@ -15,11 +15,11 @@
 
 """Utility for handling flags."""
 
+from collections.abc import Sequence
 import re
-from typing import List
 
 
-def normalize_flags(argv: List[str]) -> List[str]:
+def normalize_flags(argv: Sequence[str]) -> list[str]:
   """Returns normalized explicit bolean flags for `absl.flags` compatibility.
 
   Note: Boolean flags in `absl.flags` can be specified with --bool, --nobool,