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 for will be generated by parsing the signature of the constructor function.

In this example, we use this pattern to define custom behavior for instantiating a NumPy array.

# 01_simple_constructors.py
from typing import Literal

import numpy as np
from typing_extensions import Annotated

import tyro

def construct_array(
    values: tuple[float, ...], dtype: Literal["float32", "float64"] = "float64"
) -> np.ndarray:
    """A custom constructor for 1D NumPy arrays."""
    return np.array(
        values,
        dtype={"float32": np.float32, "float64": np.float64}[dtype],
    )

def main(
    # We can specify a custom constructor for an argument in `tyro.conf.arg()`.
    array: Annotated[np.ndarray, tyro.conf.arg(constructor=construct_array)],
) -> None:
    print(f"{array=}")

if __name__ == "__main__":
    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.

# 02_primitive_annotation.py
import json

from typing_extensions import Annotated

import tyro

# A dictionary type, but `tyro` will expect a JSON string from the CLI.
JsonDict = Annotated[
    dict,
    tyro.constructors.PrimitiveConstructorSpec(
        # Number of arguments to consume.
        nargs=1,
        # Argument name in usage messages.
        metavar="JSON",
        # Convert a list of strings to an instance. The length of the list
        # should match `nargs`.
        instance_from_str=lambda args: json.loads(args[0]),
        # Check if an instance is of the expected type. This is only used for
        # helptext formatting in the presence of union types.
        is_instance=lambda instance: isinstance(instance, dict),
        # Convert an instance to a list of strings. This is used for handling
        # default values that are set in Python. The length of the list should
        # match `nargs`.
        str_from_instance=lambda instance: [json.dumps(instance)],
    ),
]

def main(
    dict1: JsonDict,
    dict2: JsonDict = {"default": None},
) -> None:
    print(f"{dict1=}")
    print(f"{dict2=}")

if __name__ == "__main__":
    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].

# 03_primitive_registry.py
import json
from typing import Any

import tyro

# Create a custom registry, which stores constructor rules.
custom_registry = tyro.constructors.ConstructorRegistry()

# Define a rule that applies to all types that match `dict[str, Any]`.
@custom_registry.primitive_rule
def _(
    type_info: tyro.constructors.PrimitiveTypeInfo,
) -> tyro.constructors.PrimitiveConstructorSpec | None:
    # We return `None` if the rule does not apply.
    if type_info.type != dict[str, Any]:
        return None

    # If the rule applies, we return the constructor spec.
    return tyro.constructors.PrimitiveConstructorSpec(
        nargs=1,
        metavar="JSON",
        instance_from_str=lambda args: json.loads(args[0]),
        is_instance=lambda instance: isinstance(instance, dict),
        str_from_instance=lambda instance: [json.dumps(instance)],
    )

def main(
    dict1: dict[str, Any],
    dict2: dict[str, Any] = {"default": None},
) -> None:
    """A function with two arguments, which can be populated from the CLI via JSON."""
    print(f"{dict1=}")
    print(f"{dict2=}")

if __name__ == "__main__":
    # To activate a custom registry, we should use it as a context manager.
    with custom_registry:
        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 type.

Warning

This will be complicated!

# 04_struct_registry.py
import tyro

# A custom type that we'll add support for to tyro.
class Bounds:
    def __init__(self, lower: int, upper: int):
        self.bounds = (lower, upper)

    def __repr__(self) -> str:
        return f"(lower={self.bounds[0]}, upper={self.bounds[1]})"

# Create a custom registry, which stores constructor rules.
custom_registry = tyro.constructors.ConstructorRegistry()

# Define a rule that applies to all types that match `Bounds`.
@custom_registry.struct_rule
def _(
    type_info: tyro.constructors.StructTypeInfo,
) -> tyro.constructors.StructConstructorSpec | None:
    # We return `None` if the rule does not apply.
    if type_info.type != Bounds:
        return None

    # We can extract the default value of the field from `type_info`.
    if isinstance(type_info.default, Bounds):
        # If the default value is a `Bounds` instance, we don't need to generate a constructor.
        default = (type_info.default.bounds[0], type_info.default.bounds[1])
        is_default_overridden = True
    else:
        # Otherwise, the default value is missing. We'll mark the child defaults as missing as well.
        assert type_info.default in (
            tyro.constructors.MISSING,
            tyro.constructors.MISSING_NONPROP,
        )
        default = (tyro.MISSING, tyro.MISSING)
        is_default_overridden = False

    # If the rule applies, we return the constructor spec.
    return tyro.constructors.StructConstructorSpec(
        # The instantiate function will be called with the fields as keyword arguments.
        instantiate=Bounds,
        fields=(
            tyro.constructors.StructFieldSpec(
                name="lower",
                type=int,
                default=default[0],
                is_default_overridden=is_default_overridden,
                helptext="Lower bound." "",
            ),
            tyro.constructors.StructFieldSpec(
                name="upper",
                type=int,
                default=default[1],
                is_default_overridden=is_default_overridden,
                helptext="Upper bound." "",
            ),
        ),
    )

def main(
    bounds: Bounds,
    bounds_with_default: Bounds = Bounds(0, 100),
) -> None:
    """A function with two `Bounds` instances as input."""
    print(f"{bounds=}")
    print(f"{bounds_with_default=}")

if __name__ == "__main__":
    # To activate a custom registry, we should use it as a context manager.
    with custom_registry:
        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)