Skip to content

Commit 80b1316

Browse files
brentyibrentyi
authored
Decorator-based subcommand API (#169)
* Decorator-based subcommands * Fix example * Switch to class, tests * Add missing * Coverage * Sync docs * Coverage
1 parent 3b7e37d commit 80b1316

File tree

6 files changed

+394
-0
lines changed

6 files changed

+394
-0
lines changed

‎docs/source/examples/04_additional/15_decorator_subcommands.rst

+85
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
.. Comment: this file is automatically generated by `update_example_docs.py`.
2+
It should not be modified manually.
3+
4+
Decorator-based Subcommands
5+
==========================================
6+
7+
:func:`tyro.extras.app.command()` and :func:`tyro.extras.app.cli()` provide a
8+
decorator-based API for subcommands, which is inspired by `click
9+
<https://click.palletsprojects.com/>`_.
10+
11+
12+
.. code-block:: python
13+
:linenos:
14+
15+
16+
from tyro.extras import SubcommandApp
17+
18+
app = SubcommandApp()
19+
20+
21+
@app.command
22+
def greet(name: str, loud: bool = False) -> None:
23+
"""Greet someone."""
24+
greeting = f"Hello, {name}!"
25+
if loud:
26+
greeting = greeting.upper()
27+
print(greeting)
28+
29+
30+
@app.command(name="addition")
31+
def add(a: int, b: int) -> None:
32+
"""Add two numbers."""
33+
print(f"{a} + {b} = {a + b}")
34+
35+
36+
if __name__ == "__main__":
37+
app.cli()
38+
39+
------------
40+
41+
.. raw:: html
42+
43+
<kbd>python 04_additional/15_decorator_subcommands.py --help</kbd>
44+
45+
.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py --help
46+
47+
------------
48+
49+
.. raw:: html
50+
51+
<kbd>python 04_additional/15_decorator_subcommands.py greet --help</kbd>
52+
53+
.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --help
54+
55+
------------
56+
57+
.. raw:: html
58+
59+
<kbd>python 04_additional/15_decorator_subcommands.py greet --name Alice</kbd>
60+
61+
.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --name Alice
62+
63+
------------
64+
65+
.. raw:: html
66+
67+
<kbd>python 04_additional/15_decorator_subcommands.py greet --name Bob --loud</kbd>
68+
69+
.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py greet --name Bob --loud
70+
71+
------------
72+
73+
.. raw:: html
74+
75+
<kbd>python 04_additional/15_decorator_subcommands.py addition --help</kbd>
76+
77+
.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py addition --help
78+
79+
------------
80+
81+
.. raw:: html
82+
83+
<kbd>python 04_additional/15_decorator_subcommands.py addition --a 5 --b 3</kbd>
84+
85+
.. program-output:: python ../../examples/04_additional/15_decorator_subcommands.py addition --a 5 --b 3

‎examples/04_additional/15_decorator_subcommands.py

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
"""Decorator-based Subcommands
2+
3+
:func:`tyro.extras.app.command()` and :func:`tyro.extras.app.cli()` provide a
4+
decorator-based API for subcommands, which is inspired by `click
5+
<https://click.palletsprojects.com/>`_.
6+
7+
Usage:
8+
`python my_script.py --help`
9+
`python my_script.py greet --help`
10+
`python my_script.py greet --name Alice`
11+
`python my_script.py greet --name Bob --loud`
12+
`python my_script.py addition --help`
13+
`python my_script.py addition --a 5 --b 3`
14+
"""
15+
16+
from tyro.extras import SubcommandApp
17+
18+
app = SubcommandApp()
19+
20+
21+
@app.command
22+
def greet(name: str, loud: bool = False) -> None:
23+
"""Greet someone."""
24+
greeting = f"Hello, {name}!"
25+
if loud:
26+
greeting = greeting.upper()
27+
print(greeting)
28+
29+
30+
@app.command(name="addition")
31+
def add(a: int, b: int) -> None:
32+
"""Add two numbers."""
33+
print(f"{a} + {b} = {a + b}")
34+
35+
36+
if __name__ == "__main__":
37+
app.cli()

‎src/tyro/extras/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from ._choices_type import literal_type_from_choices as literal_type_from_choices
1212
from ._serialization import from_yaml as from_yaml
1313
from ._serialization import to_yaml as to_yaml
14+
from ._subcommand_app import SubcommandApp as SubcommandApp
1415
from ._subcommand_cli_from_dict import (
1516
subcommand_cli_from_dict as subcommand_cli_from_dict,
1617
)

‎src/tyro/extras/_subcommand_app.py

+140
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
from __future__ import annotations
2+
3+
from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, overload
4+
5+
import tyro
6+
7+
CallableT = TypeVar("CallableT", bound=Callable)
8+
9+
10+
class SubcommandApp:
11+
"""This module provides a decorator-based API for subcommands in `tyro`, inspired by click.
12+
13+
Example:
14+
15+
```python
16+
from tyro.extras import SubcommandApp
17+
18+
app = SubcommandApp()
19+
20+
@app.command
21+
def greet(name: str, loud: bool = False):
22+
'''Greet someone.'''
23+
greeting = f"Hello, {name}!"
24+
if loud:
25+
greeting = greeting.upper()
26+
print(greeting)
27+
28+
@app.command(name="addition")
29+
def add(a: int, b: int):
30+
'''Add two numbers.'''
31+
print(f"{a} + {b} = {a + b}")
32+
33+
if __name__ == "__main__":
34+
app.cli()
35+
```
36+
37+
Usage:
38+
`python my_script.py greet Alice`
39+
`python my_script.py greet Bob --loud`
40+
`python my_script.py addition 5 3`
41+
"""
42+
43+
def __init__(self) -> None:
44+
self._subcommands: Dict[str, Callable] = {}
45+
46+
@overload
47+
def command(self, func: CallableT) -> CallableT: ...
48+
49+
@overload
50+
def command(
51+
self,
52+
func: None = None,
53+
*,
54+
name: str | None = None,
55+
) -> Callable[[CallableT], CallableT]: ...
56+
57+
def command(
58+
self,
59+
func: CallableT | None = None,
60+
*,
61+
name: str | None = None,
62+
) -> CallableT | Callable[[CallableT], CallableT]:
63+
"""A decorator to register a function as a subcommand.
64+
65+
This method is inspired by Click's @cli.command() decorator.
66+
It adds the decorated function to the list of subcommands.
67+
68+
Args:
69+
func: The function to register as a subcommand. If None, returns a
70+
function to use as a decorator.
71+
name: The name of the subcommand. If None, the name of the function is used.
72+
"""
73+
74+
def inner(func: CallableT) -> CallableT:
75+
nonlocal name
76+
if name is None:
77+
name = func.__name__
78+
79+
self._subcommands[name] = func
80+
return func
81+
82+
if func is not None:
83+
return inner(func)
84+
else:
85+
return inner
86+
87+
def cli(
88+
self,
89+
*,
90+
prog: Optional[str] = None,
91+
description: Optional[str] = None,
92+
args: Optional[Sequence[str]] = None,
93+
use_underscores: bool = False,
94+
sort_subcommands: bool = True,
95+
) -> Any:
96+
"""Run the command-line interface.
97+
98+
This method creates a CLI using tyro, with all subcommands registered using
99+
:func:`command()`.
100+
101+
Args:
102+
prog: The name of the program printed in helptext. Mirrors argument from
103+
`argparse.ArgumentParser()`.
104+
description: Description text for the parser, displayed when the --help flag is
105+
passed in. If not specified, the class docstring is used. Mirrors argument from
106+
`argparse.ArgumentParser()`.
107+
args: If set, parse arguments from a sequence of strings instead of the
108+
commandline. Mirrors argument from `argparse.ArgumentParser.parse_args()`.
109+
use_underscores: If True, use underscores as a word delimiter instead of hyphens.
110+
This primarily impacts helptext; underscores and hyphens are treated equivalently
111+
when parsing happens. We default helptext to hyphens to follow the GNU style guide.
112+
https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html
113+
sort_subcommands: If True, sort the subcommands alphabetically by name.
114+
"""
115+
assert self._subcommands is not None
116+
117+
# Sort subcommands by name.
118+
if sort_subcommands:
119+
sorted_subcommands = dict(
120+
sorted(self._subcommands.items(), key=lambda x: x[0])
121+
)
122+
else:
123+
sorted_subcommands = self._subcommands
124+
125+
if len(sorted_subcommands) == 1:
126+
return tyro.cli(
127+
next(iter(sorted_subcommands.values())),
128+
prog=prog,
129+
description=description,
130+
args=args,
131+
use_underscores=use_underscores,
132+
)
133+
else:
134+
return tyro.extras.subcommand_cli_from_dict(
135+
sorted_subcommands,
136+
prog=prog,
137+
description=description,
138+
args=args,
139+
use_underscores=use_underscores,
140+
)

‎tests/test_decorator_subcommands.py

+72
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import pytest
2+
3+
from tyro.extras import SubcommandApp
4+
5+
app = SubcommandApp()
6+
app_just_one = SubcommandApp()
7+
8+
9+
@app_just_one.command
10+
@app.command
11+
def greet(name: str, loud: bool = False) -> None:
12+
"""Greet someone."""
13+
greeting = f"Hello, {name}!"
14+
if loud:
15+
greeting = greeting.upper()
16+
print(greeting)
17+
18+
19+
@app.command(name="addition")
20+
def add(a: int, b: int) -> None:
21+
"""Add two numbers."""
22+
print(f"{a} + {b} = {a + b}")
23+
24+
25+
def test_app_just_one_cli(capsys):
26+
# Test: `python my_script.py --help`
27+
with pytest.raises(SystemExit):
28+
app_just_one.cli(args=["--help"], sort_subcommands=False)
29+
captured = capsys.readouterr()
30+
assert "usage: " in captured.out
31+
assert "greet" not in captured.out
32+
assert "addition" not in captured.out
33+
assert "--name" in captured.out
34+
35+
36+
def test_app_cli(capsys):
37+
# Test: `python my_script.py --help`
38+
with pytest.raises(SystemExit):
39+
app.cli(args=["--help"])
40+
captured = capsys.readouterr()
41+
assert "usage: " in captured.out
42+
assert "greet" in captured.out
43+
assert "addition" in captured.out
44+
45+
# Test: `python my_script.py greet --help`
46+
with pytest.raises(SystemExit):
47+
app.cli(args=["greet", "--help"])
48+
captured = capsys.readouterr()
49+
assert "usage: " in captured.out
50+
assert "Greet someone." in captured.out
51+
52+
# Test: `python my_script.py greet --name Alice`
53+
app.cli(args=["greet", "--name", "Alice"])
54+
captured = capsys.readouterr()
55+
assert captured.out.strip() == "Hello, Alice!"
56+
57+
# Test: `python my_script.py greet --name Bob --loud`
58+
app.cli(args=["greet", "--name", "Bob", "--loud"])
59+
captured = capsys.readouterr()
60+
assert captured.out.strip() == "HELLO, BOB!"
61+
62+
# Test: `python my_script.py addition --help`
63+
with pytest.raises(SystemExit):
64+
app.cli(args=["addition", "--help"])
65+
captured = capsys.readouterr()
66+
assert "usage: " in captured.out
67+
assert "Add two numbers." in captured.out
68+
69+
# Test: `python my_script.py addition 5 3`
70+
app.cli(args=["addition", "--a", "5", "--b", "3"])
71+
captured = capsys.readouterr()
72+
assert captured.out.strip() == "5 + 3 = 8"

0 commit comments

Comments
 (0)