Custom constructors¶
tyro.cli()
aims for comprehensive support of standard Python type
constructs. It can still, however, be useful to extend the set of suported
types.
We provide two complementary approaches for doing so:
tyro.conf
provides a simple API for specifying custom constructor functions.tyro.constructors
provides a more flexible API for defining behavior for different types. There are two categories of types: primitive types are instantiated from a single commandline argument, while struct types are broken down into multiple arguments.
Warning
Custom constructors are useful, but can be verbose and require care. We recommend using them sparingly.
Simple Constructors¶
For simple custom constructors, we can pass a constructor function into
tyro.conf.arg()
or tyro.conf.subcommand()
. Arguments will be
generated by parsing the signature of the constructor function.
In this example, we define custom behavior for instantiating a NumPy array.
1# 01_simple_constructors.py
2from typing import Literal
3
4import numpy as np
5from typing_extensions import Annotated
6
7import tyro
8
9def construct_array(
10 values: tuple[float, ...], dtype: Literal["float32", "float64"] = "float64"
11) -> np.ndarray:
12 """A custom constructor for 1D NumPy arrays."""
13 return np.array(
14 values,
15 dtype={"float32": np.float32, "float64": np.float64}[dtype],
16 )
17
18def main(
19 # We can specify a custom constructor for an argument in `tyro.conf.arg()`.
20 array: Annotated[np.ndarray, tyro.conf.arg(constructor=construct_array)],
21) -> None:
22 print(f"{array=}")
23
24if __name__ == "__main__":
25 tyro.cli(main)
$ python ./01_simple_constructors.py --help usage: 01_simple_constructors.py [-h] --array.values [FLOAT [FLOAT ...]] [--array.dtype {float32,float64}] ╭─ options ─────────────────────────────────────────╮ │ -h, --help show this help message and exit │ ╰───────────────────────────────────────────────────╯ ╭─ array options ───────────────────────────────────╮ │ A custom constructor for 1D NumPy arrays. │ │ ───────────────────────────────────────── │ │ --array.values [FLOAT [FLOAT ...]] │ │ (required) │ │ --array.dtype {float32,float64} │ │ (default: float64) │ ╰───────────────────────────────────────────────────╯
$ python ./01_simple_constructors.py --array.values 1 2 3
array=array([1., 2., 3.])
$ python ./01_simple_constructors.py --array.values 1 2 3 4 5 --array.dtype float32
array=array([1., 2., 3., 4., 5.], dtype=float32)
Custom Primitive¶
In this example, we use tyro.constructors
to attach a primitive
constructor via a runtime annotation.
1# 02_primitive_annotation.py
2import json
3
4from typing_extensions import Annotated
5
6import tyro
7
8# A dictionary type, but `tyro` will expect a JSON string from the CLI.
9JsonDict = Annotated[
10 dict,
11 tyro.constructors.PrimitiveConstructorSpec(
12 # Number of arguments to consume.
13 nargs=1,
14 # Argument name in usage messages.
15 metavar="JSON",
16 # Convert a list of strings to an instance. The length of the list
17 # should match `nargs`.
18 instance_from_str=lambda args: json.loads(args[0]),
19 # Check if an instance is of the expected type. This is only used for
20 # helptext formatting in the presence of union types.
21 is_instance=lambda instance: isinstance(instance, dict),
22 # Convert an instance to a list of strings. This is used for handling
23 # default values that are set in Python. The length of the list should
24 # match `nargs`.
25 str_from_instance=lambda instance: [json.dumps(instance)],
26 ),
27]
28
29def main(
30 dict1: JsonDict,
31 dict2: JsonDict = {"default": None},
32) -> None:
33 print(f"{dict1=}")
34 print(f"{dict2=}")
35
36if __name__ == "__main__":
37 tyro.cli(main)
$ python ./02_primitive_annotation.py --help usage: 02_primitive_annotation.py [-h] --dict1 JSON [--dict2 JSON] ╭─ options ───────────────────────────────────────────╮ │ -h, --help show this help message and exit │ │ --dict1 JSON (required) │ │ --dict2 JSON (default: '{"default": null}') │ ╰─────────────────────────────────────────────────────╯
$ python ./02_primitive_annotation.py --dict1 '{"hello": "world"}'
dict1={'hello': 'world'}
dict2={'default': None}
$ python ./02_primitive_annotation.py --dict1 '{"hello": "world"}' --dict2 '{"hello": "world"}'
dict1={'hello': 'world'}
dict2={'hello': 'world'}
Custom Primitive (Registry)¶
In this example, we use a tyro.constructors.ConstructorRegistry
to
define a rule that applies to all types that match dict[str, Any]
.
1# 03_primitive_registry.py
2import json
3from typing import Any
4
5import tyro
6
7# Create a custom registry, which stores constructor rules.
8custom_registry = tyro.constructors.ConstructorRegistry()
9
10# Define a rule that applies to all types that match `dict[str, Any]`.
11@custom_registry.primitive_rule
12def _(
13 type_info: tyro.constructors.PrimitiveTypeInfo,
14) -> tyro.constructors.PrimitiveConstructorSpec | None:
15 # We return `None` if the rule does not apply.
16 if type_info.type != dict[str, Any]:
17 return None
18
19 # If the rule applies, we return the constructor spec.
20 return tyro.constructors.PrimitiveConstructorSpec(
21 nargs=1,
22 metavar="JSON",
23 instance_from_str=lambda args: json.loads(args[0]),
24 is_instance=lambda instance: isinstance(instance, dict),
25 str_from_instance=lambda instance: [json.dumps(instance)],
26 )
27
28def main(
29 dict1: dict[str, Any],
30 dict2: dict[str, Any] = {"default": None},
31) -> None:
32 """A function with two arguments, which can be populated from the CLI via JSON."""
33 print(f"{dict1=}")
34 print(f"{dict2=}")
35
36if __name__ == "__main__":
37 # To activate a custom registry, we should use it as a context manager.
38 with custom_registry:
39 tyro.cli(main)
$ python ./03_primitive_registry.py --help usage: 03_primitive_registry.py [-h] --dict1 JSON [--dict2 JSON] A function with two arguments, which can be populated from the CLI via JSON. ╭─ options ───────────────────────────────────────────╮ │ -h, --help show this help message and exit │ │ --dict1 JSON (required) │ │ --dict2 JSON (default: '{"default": null}') │ ╰─────────────────────────────────────────────────────╯
$ python ./03_primitive_registry.py --dict1 '{"hello": "world"}'
dict1={'hello': 'world'}
dict2={'default': None}
$ python ./03_primitive_registry.py --dict1 '{"hello": "world"}' --dict2 '{"hello": "world"}'
dict1={'hello': 'world'}
dict2={'hello': 'world'}
Custom Structs (Registry)¶
In this example, we use a tyro.constructors.ConstructorRegistry
to
add support for a custom struct type. Each struct type is broken down into
multiple fields, which themselves can be either primitive types or nested
structs.
Warning
This will be complicated!
1# 04_struct_registry.py
2import tyro
3
4# A custom type that we'll add support for to tyro.
5class Bounds:
6 def __init__(self, lower: int, upper: int):
7 self.bounds = (lower, upper)
8
9 def __repr__(self) -> str:
10 return f"(lower={self.bounds[0]}, upper={self.bounds[1]})"
11
12# Create a custom registry, which stores constructor rules.
13custom_registry = tyro.constructors.ConstructorRegistry()
14
15# Define a rule that applies to all types that match `Bounds`.
16@custom_registry.struct_rule
17def _(
18 type_info: tyro.constructors.StructTypeInfo,
19) -> tyro.constructors.StructConstructorSpec | None:
20 # We return `None` if the rule does not apply.
21 if type_info.type != Bounds:
22 return None
23
24 # We can extract the default value of the field from `type_info`.
25 if isinstance(type_info.default, Bounds):
26 # If the default value is a `Bounds` instance, we don't need to generate a constructor.
27 default = (type_info.default.bounds[0], type_info.default.bounds[1])
28 is_default_overridden = True
29 else:
30 # Otherwise, the default value is missing. We'll mark the child defaults as missing as well.
31 assert type_info.default in (
32 tyro.constructors.MISSING,
33 tyro.constructors.MISSING_NONPROP,
34 )
35 default = (tyro.MISSING, tyro.MISSING)
36 is_default_overridden = False
37
38 # If the rule applies, we return the constructor spec.
39 return tyro.constructors.StructConstructorSpec(
40 # The instantiate function will be called with the fields as keyword arguments.
41 instantiate=Bounds,
42 fields=(
43 tyro.constructors.StructFieldSpec(
44 name="lower",
45 type=int,
46 default=default[0],
47 is_default_overridden=is_default_overridden,
48 helptext="Lower bound." "",
49 ),
50 tyro.constructors.StructFieldSpec(
51 name="upper",
52 type=int,
53 default=default[1],
54 is_default_overridden=is_default_overridden,
55 helptext="Upper bound." "",
56 ),
57 ),
58 )
59
60def main(
61 bounds: Bounds,
62 bounds_with_default: Bounds = Bounds(0, 100),
63) -> None:
64 """A function with two `Bounds` instances as input."""
65 print(f"{bounds=}")
66 print(f"{bounds_with_default=}")
67
68if __name__ == "__main__":
69 # To activate a custom registry, we should use it as a context manager.
70 with custom_registry:
71 tyro.cli(main)
$ python ./04_struct_registry.py --help usage: 04_struct_registry.py [-h] [OPTIONS] A function with two `Bounds` instances as input. ╭─ options ───────────────────────────────────────────────╮ │ -h, --help show this help message and exit │ ╰─────────────────────────────────────────────────────────╯ ╭─ bounds options ────────────────────────────────────────╮ │ --bounds.lower INT Lower bound. (required) │ │ --bounds.upper INT Upper bound. (required) │ ╰─────────────────────────────────────────────────────────╯ ╭─ bounds-with-default options ───────────────────────────╮ │ --bounds-with-default.lower INT │ │ Lower bound. (default: 0) │ │ --bounds-with-default.upper INT │ │ Upper bound. (default: 100) │ ╰─────────────────────────────────────────────────────────╯
$ python ./04_struct_registry.py --bounds.lower 5 --bounds.upper 10
bounds=(lower=5, upper=10)
bounds_with_default=(lower=0, upper=100)