Decorator-based subcommand API (#169)

* Decorator-based subcommands * Fix example * Switch to class, tests * Add missing * Coverage * Sync docs * Coverage
brentyi · Oct 10, 2024 · 80b1316 · 80b1316
1 parent 3b7e37d
commit 80b1316
Show file tree

Hide file tree

Showing 6 changed files with 394 additions and 0 deletions.
diff --git a/docs/source/examples/04_additional/15_decorator_subcommands.rst 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
+<https://click.palletsprojects.com/>`_.
+
+
+.. 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
+
+        <kbd>python 04_additional/15_decorator_subcommands.py --help</kbd>
+
+.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py --help
+
+------------
+
+.. raw:: html
+
+        <kbd>python 04_additional/15_decorator_subcommands.py greet --help</kbd>
+
+.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --help
+
+------------
+
+.. raw:: html
+
+        <kbd>python 04_additional/15_decorator_subcommands.py greet --name Alice</kbd>
+
+.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --name Alice
+
+------------
+
+.. raw:: html
+
+        <kbd>python 04_additional/15_decorator_subcommands.py greet --name Bob --loud</kbd>
+
+.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --name Bob --loud
+
+------------
+
+.. raw:: html
+
+        <kbd>python 04_additional/15_decorator_subcommands.py addition --help</kbd>
+
+.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py addition --help
+
+------------
+
+.. raw:: html
+
+        <kbd>python 04_additional/15_decorator_subcommands.py addition --a 5 --b 3</kbd>
+
+.. 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
@@ -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
+<https://click.palletsprojects.com/>`_.
+
+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
@@ -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
@@ -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
@@ -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"