diff --git a/docs/source/examples/04_additional/15_decorator_subcommands.rst b/docs/source/examples/04_additional/15_decorator_subcommands.rst new file mode 100644 index 00000000..5b8b5f50 --- /dev/null +++ b/docs/source/examples/04_additional/15_decorator_subcommands.rst @@ -0,0 +1,85 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Decorator-based Subcommands +========================================== + +:func:`tyro.extras.app.command()` and :func:`tyro.extras.app.cli()` provide a +decorator-based API for subcommands, which is inspired by `click +`_. + + +.. code-block:: python + :linenos: + + + from tyro.extras import SubcommandApp + + app = SubcommandApp() + + + @app.command + def greet(name: str, loud: bool = False) -> None: + """Greet someone.""" + greeting = f"Hello, {name}!" + if loud: + greeting = greeting.upper() + print(greeting) + + + @app.command(name="addition") + def add(a: int, b: int) -> None: + """Add two numbers.""" + print(f"{a} + {b} = {a + b}") + + + if __name__ == "__main__": + app.cli() + +------------ + +.. raw:: html + + python 04_additional/15_decorator_subcommands.py --help + +.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py --help + +------------ + +.. raw:: html + + python 04_additional/15_decorator_subcommands.py greet --help + +.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --help + +------------ + +.. raw:: html + + python 04_additional/15_decorator_subcommands.py greet --name Alice + +.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --name Alice + +------------ + +.. raw:: html + + python 04_additional/15_decorator_subcommands.py greet --name Bob --loud + +.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --name Bob --loud + +------------ + +.. raw:: html + + python 04_additional/15_decorator_subcommands.py addition --help + +.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py addition --help + +------------ + +.. raw:: html + + python 04_additional/15_decorator_subcommands.py addition --a 5 --b 3 + +.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py addition --a 5 --b 3 diff --git a/examples/04_additional/15_decorator_subcommands.py b/examples/04_additional/15_decorator_subcommands.py new file mode 100644 index 00000000..466b9a78 --- /dev/null +++ b/examples/04_additional/15_decorator_subcommands.py @@ -0,0 +1,37 @@ +"""Decorator-based Subcommands + +:func:`tyro.extras.app.command()` and :func:`tyro.extras.app.cli()` provide a +decorator-based API for subcommands, which is inspired by `click +`_. + +Usage: +`python my_script.py --help` +`python my_script.py greet --help` +`python my_script.py greet --name Alice` +`python my_script.py greet --name Bob --loud` +`python my_script.py addition --help` +`python my_script.py addition --a 5 --b 3` +""" + +from tyro.extras import SubcommandApp + +app = SubcommandApp() + + +@app.command +def greet(name: str, loud: bool = False) -> None: + """Greet someone.""" + greeting = f"Hello, {name}!" + if loud: + greeting = greeting.upper() + print(greeting) + + +@app.command(name="addition") +def add(a: int, b: int) -> None: + """Add two numbers.""" + print(f"{a} + {b} = {a + b}") + + +if __name__ == "__main__": + app.cli() diff --git a/src/tyro/extras/__init__.py b/src/tyro/extras/__init__.py index 0afb6b85..39878d58 100644 --- a/src/tyro/extras/__init__.py +++ b/src/tyro/extras/__init__.py @@ -11,6 +11,7 @@ from ._choices_type import literal_type_from_choices as literal_type_from_choices from ._serialization import from_yaml as from_yaml from ._serialization import to_yaml as to_yaml +from ._subcommand_app import SubcommandApp as SubcommandApp from ._subcommand_cli_from_dict import ( subcommand_cli_from_dict as subcommand_cli_from_dict, ) diff --git a/src/tyro/extras/_subcommand_app.py b/src/tyro/extras/_subcommand_app.py new file mode 100644 index 00000000..0a8a906f --- /dev/null +++ b/src/tyro/extras/_subcommand_app.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, overload + +import tyro + +CallableT = TypeVar("CallableT", bound=Callable) + + +class SubcommandApp: + """This module provides a decorator-based API for subcommands in `tyro`, inspired by click. + + Example: + + ```python + from tyro.extras import SubcommandApp + + app = SubcommandApp() + + @app.command + def greet(name: str, loud: bool = False): + '''Greet someone.''' + greeting = f"Hello, {name}!" + if loud: + greeting = greeting.upper() + print(greeting) + + @app.command(name="addition") + def add(a: int, b: int): + '''Add two numbers.''' + print(f"{a} + {b} = {a + b}") + + if __name__ == "__main__": + app.cli() + ``` + + Usage: + `python my_script.py greet Alice` + `python my_script.py greet Bob --loud` + `python my_script.py addition 5 3` + """ + + def __init__(self) -> None: + self._subcommands: Dict[str, Callable] = {} + + @overload + def command(self, func: CallableT) -> CallableT: ... + + @overload + def command( + self, + func: None = None, + *, + name: str | None = None, + ) -> Callable[[CallableT], CallableT]: ... + + def command( + self, + func: CallableT | None = None, + *, + name: str | None = None, + ) -> CallableT | Callable[[CallableT], CallableT]: + """A decorator to register a function as a subcommand. + + This method is inspired by Click's @cli.command() decorator. + It adds the decorated function to the list of subcommands. + + Args: + func: The function to register as a subcommand. If None, returns a + function to use as a decorator. + name: The name of the subcommand. If None, the name of the function is used. + """ + + def inner(func: CallableT) -> CallableT: + nonlocal name + if name is None: + name = func.__name__ + + self._subcommands[name] = func + return func + + if func is not None: + return inner(func) + else: + return inner + + def cli( + self, + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + use_underscores: bool = False, + sort_subcommands: bool = True, + ) -> Any: + """Run the command-line interface. + + This method creates a CLI using tyro, with all subcommands registered using + :func:`command()`. + + Args: + prog: The name of the program printed in helptext. Mirrors argument from + `argparse.ArgumentParser()`. + description: Description text for the parser, displayed when the --help flag is + passed in. If not specified, the class docstring is used. Mirrors argument from + `argparse.ArgumentParser()`. + args: If set, parse arguments from a sequence of strings instead of the + commandline. Mirrors argument from `argparse.ArgumentParser.parse_args()`. + use_underscores: If True, use underscores as a word delimiter instead of hyphens. + This primarily impacts helptext; underscores and hyphens are treated equivalently + when parsing happens. We default helptext to hyphens to follow the GNU style guide. + https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html + sort_subcommands: If True, sort the subcommands alphabetically by name. + """ + assert self._subcommands is not None + + # Sort subcommands by name. + if sort_subcommands: + sorted_subcommands = dict( + sorted(self._subcommands.items(), key=lambda x: x[0]) + ) + else: + sorted_subcommands = self._subcommands + + if len(sorted_subcommands) == 1: + return tyro.cli( + next(iter(sorted_subcommands.values())), + prog=prog, + description=description, + args=args, + use_underscores=use_underscores, + ) + else: + return tyro.extras.subcommand_cli_from_dict( + sorted_subcommands, + prog=prog, + description=description, + args=args, + use_underscores=use_underscores, + ) diff --git a/tests/test_decorator_subcommands.py b/tests/test_decorator_subcommands.py new file mode 100644 index 00000000..422819d1 --- /dev/null +++ b/tests/test_decorator_subcommands.py @@ -0,0 +1,72 @@ +import pytest + +from tyro.extras import SubcommandApp + +app = SubcommandApp() +app_just_one = SubcommandApp() + + +@app_just_one.command +@app.command +def greet(name: str, loud: bool = False) -> None: + """Greet someone.""" + greeting = f"Hello, {name}!" + if loud: + greeting = greeting.upper() + print(greeting) + + +@app.command(name="addition") +def add(a: int, b: int) -> None: + """Add two numbers.""" + print(f"{a} + {b} = {a + b}") + + +def test_app_just_one_cli(capsys): + # Test: `python my_script.py --help` + with pytest.raises(SystemExit): + app_just_one.cli(args=["--help"], sort_subcommands=False) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "greet" not in captured.out + assert "addition" not in captured.out + assert "--name" in captured.out + + +def test_app_cli(capsys): + # Test: `python my_script.py --help` + with pytest.raises(SystemExit): + app.cli(args=["--help"]) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "greet" in captured.out + assert "addition" in captured.out + + # Test: `python my_script.py greet --help` + with pytest.raises(SystemExit): + app.cli(args=["greet", "--help"]) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "Greet someone." in captured.out + + # Test: `python my_script.py greet --name Alice` + app.cli(args=["greet", "--name", "Alice"]) + captured = capsys.readouterr() + assert captured.out.strip() == "Hello, Alice!" + + # Test: `python my_script.py greet --name Bob --loud` + app.cli(args=["greet", "--name", "Bob", "--loud"]) + captured = capsys.readouterr() + assert captured.out.strip() == "HELLO, BOB!" + + # Test: `python my_script.py addition --help` + with pytest.raises(SystemExit): + app.cli(args=["addition", "--help"]) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "Add two numbers." in captured.out + + # Test: `python my_script.py addition 5 3` + app.cli(args=["addition", "--a", "5", "--b", "3"]) + captured = capsys.readouterr() + assert captured.out.strip() == "5 + 3 = 8" diff --git a/tests/test_py311_generated/test_decorator_subcommands_generated.py b/tests/test_py311_generated/test_decorator_subcommands_generated.py new file mode 100644 index 00000000..8a623bc3 --- /dev/null +++ b/tests/test_py311_generated/test_decorator_subcommands_generated.py @@ -0,0 +1,59 @@ +import pytest + +from tyro.extras import SubcommandApp + +app = SubcommandApp() + + +@app.command +def greet(name: str, loud: bool = False) -> None: + """Greet someone.""" + greeting = f"Hello, {name}!" + if loud: + greeting = greeting.upper() + print(greeting) + + +@app.command(name="addition") +def add(a: int, b: int) -> None: + """Add two numbers.""" + print(f"{a} + {b} = {a + b}") + + +def test_app_cli(capsys): + # Test: `python my_script.py --help` + with pytest.raises(SystemExit): + app.cli(args=["--help"]) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "greet" in captured.out + assert "addition" in captured.out + + # Test: `python my_script.py greet --help` + with pytest.raises(SystemExit): + app.cli(args=["greet", "--help"]) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "Greet someone." in captured.out + + # Test: `python my_script.py greet --name Alice` + app.cli(args=["greet", "--name", "Alice"]) + captured = capsys.readouterr() + assert captured.out.strip() == "Hello, Alice!" + + # Test: `python my_script.py greet --name Bob --loud` + app.cli(args=["greet", "--name", "Bob", "--loud"]) + captured = capsys.readouterr() + assert captured.out.strip() == "HELLO, BOB!" + + # Test: `python my_script.py addition --help` + with pytest.raises(SystemExit): + app.cli(args=["addition", "--help"]) + captured = capsys.readouterr() + assert "usage: " in captured.out + assert "Add two numbers." in captured.out + + # Test: `python my_script.py addition 5 3` + app.cli(args=["addition", "--a", "5", "--b", "3"]) + captured = capsys.readouterr() + assert captured.out.strip() == "5 + 3 = 8"