Update download_and_prepare to use build_croissant CLI command.

PiperOrigin-RevId: 621454738

Author: fineguy

tensorflow_datasets/scripts/cli/build_test.py

@@ -83,7 +83,7 @@ def mock_cwd(path: epath.PathLike) -> Iterator[None]:
 def _build(cmd_flags: str, mock_download_and_prepare: bool = True) -> List[str]:
   """Executes `tfds build {cmd_flags}` and returns the list of generated ds."""
   # Execute the command
-  args = main.parse_flags(f'tfds build {cmd_flags}'.split())
+  args = main._parse_flags(f'tfds build {cmd_flags}'.split())
 
   original_dl_and_prepare = tfds.core.DatasetBuilder.download_and_prepare
 
@@ -325,7 +325,7 @@ def test_download_only():
     ),
 )
 def test_make_download_config(args: str, download_config_kwargs):
-  args = main.parse_flags(f'tfds build x {download_config_kwargs}'.split())
+  args = main._parse_flags(f'tfds build x {download_config_kwargs}'.split())
   actual = build_lib._make_download_config(args, dataset_name='x')
   # Ignore the beam runner
   actual.replace(beam_runner=None)

tensorflow_datasets/scripts/cli/main.py

@@ -42,7 +42,7 @@
 FLAGS = flags.FLAGS
 
 
-def parse_flags(argv: List[str]) -> argparse.Namespace:
+def _parse_flags(argv: List[str]) -> argparse.Namespace:
   """Command lines flag parsing."""
   argv = flag_utils.normalize_flags(argv)  # See b/174043007 for context.
 
@@ -102,7 +102,7 @@ def main(args: argparse.Namespace) -> None:
 
 def launch_cli() -> None:
   """Parse arguments and launch the CLI main function."""
-  app.run(main, flags_parser=parse_flags)
+  app.run(main, flags_parser=_parse_flags)
 
 
 if __name__ == '__main__':

tensorflow_datasets/scripts/cli/main_test.py

@@ -28,4 +28,4 @@ def _check_exit(status=0, message=None):
   # Argparse call `sys.exit(0)` when `--version` is passed.
   with mock.patch('sys.exit', _check_exit):
     version_flag = '--version'
-    main.main(main.parse_flags(['', version_flag]))
+    main.main(main._parse_flags(['', version_flag]))

tensorflow_datasets/scripts/cli/new_test.py

@@ -23,7 +23,7 @@
 
 
 def _run_cli(cmd: str) -> None:
-  main.main(main.parse_flags([''] + cmd.split()))
+  main.main(main._parse_flags([''] + cmd.split()))
 
 
 def test_new_without_args(capsys):

tensorflow_datasets/scripts/download_and_prepare.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-r"""Wrapper around `tfds build` and `tfds build_croissant`."""
+r"""Wrapper around `tfds build`."""
 
 import argparse
 from typing import List
@@ -35,18 +35,7 @@
 
 def _parse_flags(argv: List[str]) -> argparse.Namespace:
   """Command lines flag parsing."""
-  parser = argparse.ArgumentParser(
-      description='Runs either `tfds build` or `tfds build_croissant`.'
-  )
-  parser.add_argument('--croissant', action='store_true')
-  args, other_argv = parser.parse_known_args(argv)
-
-  if args.croissant:
-    build_cmd = 'build_croissant'
-  else:
-    build_cmd = 'build'
-
-  return main_cli.parse_flags([other_argv[0], build_cmd] + other_argv[1:])
+  return main_cli._parse_flags([argv[0], 'build'] + argv[1:])  # pylint: disable=protected-access
 
 
 _display_warning = True