diff --git a/.flake8 b/.flake8 deleted file mode 100644 index a0bebef2b..000000000 --- a/.flake8 +++ /dev/null @@ -1,6 +0,0 @@ -[flake8] -# E203: whitespace before : -# E501: line too long ( characters) -# W503: line break before binary operator -; ignore = E203,E501,D100,D101,D102,D103,W503 -ignore = E203,E501,W503 diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index be5c679c0..94c1a7d4d 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -2,27 +2,33 @@ name: build on: push: - branches: [master] + branches: [main] pull_request: - branches: [master] + branches: [main] jobs: build: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10"] + python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - name: Install dependencies + if: matrix.python-version == '3.7' run: | - python -m pip install --upgrade pip - pip install -e ".[testing]" + pip install --upgrade pip + pip install ".[dev]" + - name: Install dependencies + if: matrix.python-version != '3.7' + run: | + pip install uv + uv pip install --system ".[dev]" - name: Test with pytest run: | pytest diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml new file mode 100644 index 000000000..b00e9eef8 --- /dev/null +++ b/.github/workflows/coverage.yml @@ -0,0 +1,33 @@ +name: coverage + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + coverage: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - name: Set up Python 3.12 + uses: actions/setup-python@v4 + with: + python-version: "3.12" + - name: Install dependencies + run: | + pip install uv + uv pip install --system -e ".[dev]" + - name: Generate coverage report + run: | + pytest --cov=tyro --cov-report=xml + - name: Upload to Codecov + uses: codecov/codecov-action@v3 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: ./coverage.xml + flags: unittests + name: codecov-umbrella + fail_ci_if_error: true + verbose: true diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 000000000..4d339befb --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,33 @@ +name: docs + +on: + push: + branches: [main] + +jobs: + docs: + runs-on: ubuntu-latest + steps: + # Check out source + - uses: actions/checkout@v2 + + # Set up Python + - name: "Set up Python 3.12" + uses: actions/setup-python@v4 + with: + python-version: "3.12" + + # Build documentation + - name: Building documentation + run: | + pip install uv + uv pip install --system -e ".[dev]" + uv pip install --system -r docs/requirements.txt + sphinx-build docs/source docs/build -b dirhtml + + # Deploy + - name: Deploy to GitHub Pages + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./docs/build diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml deleted file mode 100644 index 5df6615b9..000000000 --- a/.github/workflows/lint.yml +++ /dev/null @@ -1,17 +0,0 @@ -name: lint - -on: - push: - branches: [master] - pull_request: - branches: [master] - -jobs: - black-check: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v1 - - name: Black Code Formatter - uses: lgeiger/black-action@master - with: - args: ". --check" diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml index 7ea44da67..1ed0331d9 100644 --- a/.github/workflows/mypy.yml +++ b/.github/workflows/mypy.yml @@ -2,28 +2,23 @@ name: mypy on: push: - branches: [master] + branches: [main] pull_request: - branches: [master] + branches: [main] jobs: mypy: runs-on: ubuntu-latest - strategy: - matrix: - python-version: ["3.8"] - steps: - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v1 + - name: "Set up Python 3.12" + uses: actions/setup-python@v4 with: - python-version: ${{ matrix.python-version }} + python-version: "3.12" - name: Install dependencies run: | - sudo apt update - python -m pip install --upgrade pip - pip install -e .[type-checking] + pip install uv + uv pip install --system -e ".[dev]" - name: Test with mypy run: | - mypy . + mypy --install-types --non-interactive . diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index d5f3859d1..deca26a88 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -1,4 +1,4 @@ -# This workflows will upload a Python Package using Twine when a release is created +# This workflows will upload a Python Package using Poetry when a release is created # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries name: Upload Python Package @@ -9,23 +9,26 @@ on: jobs: deploy: - runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - name: Set up Python - uses: actions/setup-python@v1 - with: - python-version: '3.x' - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install setuptools wheel twine - - name: Build and publish - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} - run: | - python setup.py sdist bdist_wheel - twine upload dist/* + - uses: actions/checkout@v2 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install dependencies + run: | + pip install uv + uv pip install --system -e ".[dev]" + uv pip install --system build twine + - name: Strip unsupported tags in README + run: | + sed -i '//,//d' README.md + - name: Build and publish + env: + PYPI_USERNAME: ${{ secrets.PYPI_USERNAME }} + PYPI_PASSWORD: ${{ secrets.PYPI_PASSWORD }} + run: | + python -m build + twine upload --username $PYPI_USERNAME --password $PYPI_PASSWORD dist/* diff --git a/.github/workflows/pyright.yml b/.github/workflows/pyright.yml new file mode 100644 index 000000000..7da8fd2ce --- /dev/null +++ b/.github/workflows/pyright.yml @@ -0,0 +1,29 @@ +name: pyright + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + pyright: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v1 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + pip install uv + uv pip install --system -e ".[dev]" + uv pip install --system -r docs/requirements.txt + - name: Run pyright + run: | + pyright . diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml new file mode 100644 index 000000000..657d36a7d --- /dev/null +++ b/.github/workflows/ruff.yml @@ -0,0 +1,27 @@ +name: ruff + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + ruff: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - name: "Set up Python 3.10" + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install dependencies + run: | + pip install uv + uv pip install --system -e ".[dev]" + - name: Ruff check + run: | + ruff check --output-format github + - name: Ruff format + run: | + ruff format --diff diff --git a/.gitignore b/.gitignore index d2104f166..c0ac32cb7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +# pyenv and direnv alternative +.rtx.toml *.swp *.pyc *.egg-info @@ -8,5 +10,6 @@ __pycache__ .hypothesis .ipynb_checkpoints .cache +.coverage build/ dist/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 000000000..83cd2bb52 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Brent Yi + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index ca1eb20b1..ae307195a 100644 --- a/README.md +++ b/README.md @@ -1,201 +1,222 @@ -# dcargs - -![build](https://github.com/brentyi/dcargs/workflows/build/badge.svg) -![mypy](https://github.com/brentyi/dcargs/workflows/mypy/badge.svg?branch=master) -![lint](https://github.com/brentyi/dcargs/workflows/lint/badge.svg) - - - -* [Feature list](#feature-list) -* [Comparisons to alternative tools](#comparisons-to-alternative-tools) -* [Example usage](#example-usage) - - - -`dcargs` is a tool for generating portable, reusable, and strongly typed CLI -interfaces from dataclass definitions. - -We expose one function, `parse(Type[T]) -> T`, which takes a dataclass type and -instantiates it via an argparse-style CLI interface. If we create a script -called `simple.py`: +
+

+ + + + + + tyro logo + + + + +

+ +

+ Documentation +   •   + pip install tyro +

+ +

+ build + mypy + pyright + ruff + + codecov + + + codecov + +

+ +
+ +tyro.cli() is a tool for generating CLI +interfaces. + +We can define configurable scripts using functions: ```python -import dataclasses - -import dcargs +"""A command-line interface defined using a function signature. +Usage: python script_name.py --foo INT [--bar STR] +""" -@dataclasses.dataclass -class Args: - field1: str # A string field. - field2: int # A numeric field. +import tyro +def main( + foo: int, + bar: str = "default", +) -> None: + ... # Main body of a script. if __name__ == "__main__": - args = dcargs.parse(Args) - print(args) + # Generate a CLI and call `main` with its two arguments: `foo` and `bar`. + tyro.cli(main) ``` -Running `python simple.py --help` would print: - -``` -usage: simple.py [-h] --field1 STR --field2 INT - -optional arguments: - -h, --help show this help message and exit - -required arguments: - --field1 STR A string field. - --field2 INT A numeric field. -``` - -And, from `python simple.py --field1 string --field2 4`: - -``` -Args(field1='string', field2=4) -``` - -### Feature list - -The parse function supports a wide range of dataclass definitions, while -automatically generating helptext from comments/docstrings. Some of the basic -features are shown in the [example below](#example-usage). - -Our unit tests cover many more complex type annotations, including classes -containing: - -- Types natively accepted by `argparse`: str, int, float, pathlib.Path, etc -- Default values for optional parameters -- Booleans, which are automatically converted to flags when provided a default - value (eg `action="store_true"` or `action="store_false"`; in the latter case, - we prefix names with `no-`) -- Enums (via `enum.Enum`; argparse's `choices` is populated and arguments are - converted automatically) -- Various container types. Some examples: - - `typing.ClassVar` types (omitted from parser) - - `typing.Optional` types - - `typing.Literal` types (populates argparse's `choices`) - - `typing.Sequence` types (populates argparse's `nargs`) - - `typing.List` types (populates argparse's `nargs`) - - `typing.Tuple` types, such as `typing.Tuple[T, T, T]` or - `typing.Tuple[T, ...]` (populates argparse's `nargs`, and converts - automatically) - - `typing.Final` types and `typing.Annotated` (for parsing, these are - effectively no-ops) - - Nested combinations of the above: `Optional[Literal[T]]`, - `Final[Optional[Sequence[T]]]`, etc -- Nested dataclasses - - Simple nesting (see `OptimizerConfig` example below) - - Unions over nested dataclasses (subparsers) - - Optional unions over nested dataclasses (optional subparsers) -- Generic dataclasses (including nested generics, see - [./examples/generics.py](./examples/generics.py)) - -A usage example is available below. Examples of additional features can be found -in the [tests](./tests/). - -### Comparisons to alternative tools - -There are several alternative libraries to `dcargs`; here's a rough summary of -some of them: - -| | Parsers from dataclasses | Parsers from attrs | Nested dataclasses | Subparsers (via Unions) | Containers | Choices from literals | Docstrings as helptext | -| ----------------------------------------------------------------------------------------------- | ------------------------ | ------------------ | ------------------ | ----------------------- | ---------- | -------------------------------------------------------- | ---------------------- | -| **dcargs** | ✓ | | ✓ | ✓ | ✓ | ✓ | ✓ | -| **[datargs](https://github.com/roee30/datargs)** | ✓ | ✓ | | ✓ | ✓ | ✓ | | -| **[simple-parsing](https://github.com/lebrice/SimpleParsing)** | ✓ | | ✓ | ✓ | ✓ | [soon](https://github.com/lebrice/SimpleParsing/pull/86) | ✓ | -| **[argparse-dataclass](https://pypi.org/project/argparse-dataclass/)** | ✓ | | | | | | | -| **[argparse-dataclasses](https://pypi.org/project/argparse-dataclasses/)** | ✓ | | | | | | | -| **[dataclass-cli](https://github.com/malte-soe/dataclass-cli)** | ✓ | | | | | | | -| **[hf_argparser](https://huggingface.co/transformers/_modules/transformers/hf_argparser.html)** | ✓ | | | | ✓ | | | - -Some other distinguishing factors that `dcargs` has put effort into: - -- Robust handling of forward references -- Support for nested containers and generics -- Strong typing: we actively avoid relying on strings or dynamic namespace - objects (eg `argparse.Namespace`) -- Simplicity + strict abstractions: we're focused on a single function API, and - don't leak any argparse implementation details to the user level. We also - intentionally don't offer any way to add argument parsing-specific logic to - dataclass definitions. (in contrast, some of the libaries above rely heavily - on dataclass field metadata, or on the more extreme end inheritance+decorators - to make parsing-specific dataclasses) - -### Example usage - -This code: +Or instantiate config objects defined using tools like `dataclasses`, `pydantic`, and `attrs`: ```python -"""An argument parsing example. +"""A command-line interface defined using a class signature. -Note that there are multiple possible ways to document dataclass attributes, all -of which are supported by the automatic helptext generator. +Usage: python script_name.py --foo INT [--bar STR] """ -import dataclasses -import enum - -import dcargs - - -class OptimizerType(enum.Enum): - ADAM = enum.auto() - SGD = enum.auto() - - -@dataclasses.dataclass -class OptimizerConfig: - # Variant of SGD to use. - type: OptimizerType - - # Learning rate to use. - learning_rate: float = 3e-4 - - # Coefficient for L2 regularization. - weight_decay: float = 1e-2 - - -@dataclasses.dataclass -class ExperimentConfig: - experiment_name: str # Experiment name to use. - - optimizer: OptimizerConfig - - seed: int = 0 - """Random seed. This is helpful for making sure that our experiments are - all reproducible!""" +from dataclasses import dataclass +import tyro +@dataclass +class Config: + foo: int + bar: str = "default" if __name__ == "__main__": - config = dcargs.parse(ExperimentConfig, description=__doc__) - print(config) -``` - -Generates the following argument parser: + # Generate a CLI and instantiate `Config` with its two arguments: `foo` and `bar`. + config = tyro.cli(Config) + # Rest of script. + assert isinstance(config, Config) # Should pass. ``` -$ python example.py --help -usage: example.py [-h] --experiment-name STR --optimizer.type {ADAM,SGD} [--optimizer.learning-rate FLOAT] - [--optimizer.weight-decay FLOAT] [--seed INT] - -An argument parsing example. - -Note that there are multiple possible ways to document dataclass attributes, all -of which are supported by the automatic helptext generator. - -optional arguments: - -h, --help show this help message and exit - --optimizer.learning-rate FLOAT - Learning rate to use. (default: 0.0003) - --optimizer.weight-decay FLOAT - Coefficient for L2 regularization. (default: 0.01) - --seed INT Random seed. This is helpful for making sure that our experiments are - all reproducible! (default: 0) - -required arguments: - --experiment-name STR - Experiment name to use. - --optimizer.type {ADAM,SGD} - Variant of SGD to use. -``` + +Other features include helptext generation, nested structures, subcommands, and +shell completion. For examples and the API reference, see our +[documentation](https://brentyi.github.io/tyro). + +### In the wild + +`tyro` is designed to be lightweight enough for throwaway scripts, while +facilitating type safety and modularity for larger projects. Examples: + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+ + nerfstudio-project/nerfstudio +
GitHub star count + + 		+ Open-source tools for neural radiance fields. +
+ + Sea-Snell/JAXSeq +
GitHub star count + + 		Train very large language models in Jax.
+ + kevinzakka/obj2mjcf +
GitHub star count + + 		Interface for processing OBJ files for Mujoco.
+ + blurgyy/jaxngp +
GitHub star count + + 		+ CUDA-accelerated implementation of + instant-ngp, in JAX. +
+ + NVIDIAGameWorks/kaolin-wisp +
GitHub star count + + 		PyTorch library for neural fields.
+ + autonomousvision/sdfstudio +
GitHub star count + + 		Unified framework for surface reconstruction.
+ + openrlbenchmark/openrlbenchmark +
GitHub star count + + 		Collection of tracked experiments for reinforcement learning.
+ + vwxyzjn/cleanrl +
GitHub star count + + 		Single-file implementation of deep RL algorithms.
+ +### Alternatives + +`tyro` is an opinionated library. If any design decisions don't make sense, +feel free to file an issue! + +You might also consider one of many alternative libraries. Some that we +particularly like: + +- [simple-parsing](https://github.com/lebrice/SimpleParsing) and + [jsonargparse](https://github.com/omni-us/jsonargparse), which provide deeper + integration with configuration file formats like YAML and JSON. +- [clipstick](https://github.com/sander76/clipstick), which focuses on + generating CLIs from Pydantic models. +- [datargs](https://github.com/roee30/datargs), which provides a minimal API for + dataclasses. +- [fire](https://github.com/google/python-fire) and + [clize](https://github.com/epsy/clize), which support arguments without type + annotations. + +We also have some notes on `tyro`'s design goals and other alternatives in the +docs [here](https://brentyi.github.io/tyro/goals_and_alternatives/). diff --git a/dcargs/__init__.py b/dcargs/__init__.py deleted file mode 100644 index a784b4a67..000000000 --- a/dcargs/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from ._parse import parse - -__all__ = ["parse"] diff --git a/dcargs/_arguments.py b/dcargs/_arguments.py deleted file mode 100644 index 6896caa68..000000000 --- a/dcargs/_arguments.py +++ /dev/null @@ -1,363 +0,0 @@ -import argparse -import collections.abc -import dataclasses -import enum -from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, TypeVar, Union - -from typing_extensions import Final, Literal, _AnnotatedAlias # Backward compatibility. - -from . import _construction, _docstrings, _strings - - -@dataclasses.dataclass(frozen=True) -class ArgumentDefinition: - """Options for defining arguments. Contains all necessary arguments for argparse's - add_argument() method.""" - - # Fields that will be populated initially. - name: str - field: dataclasses.Field - parent_class: Type - type: Optional[Union[Type, TypeVar]] - - # Fields that will be handled by argument transformations. - required: Optional[bool] = None - action: Optional[str] = None - nargs: Optional[Union[int, str]] = None - default: Optional[Any] = None - choices: Optional[Set[Any]] = None - metavar: Optional[str] = None - help: Optional[str] = None - dest: Optional[str] = None - - def add_argument( - self, parser: Union[argparse.ArgumentParser, argparse._ArgumentGroup] - ) -> None: - """Add a defined argument to a parser.""" - kwargs = {k: v for k, v in vars(self).items() if v is not None} - name = "--" + kwargs.pop("name").replace("_", "-") - kwargs.pop("field") - kwargs.pop("parent_class") - parser.add_argument(name, **kwargs) - - def prefix(self, prefix: str) -> "ArgumentDefinition": - """Prefix an argument's name and destination. Used for nested dataclasses.""" - _strings.NESTED_DATACLASS_DELIMETER - arg = self - arg = dataclasses.replace(arg, name=prefix + arg.name) - if arg.dest is not None: - arg = dataclasses.replace(arg, dest=prefix + arg.dest) - return arg - - @staticmethod - def make_from_field( - parent_class: Type, - field: dataclasses.Field, - type_from_typevar: Dict[TypeVar, Type], - ) -> Tuple["ArgumentDefinition", _construction.FieldRole]: - """Create an argument definition from a field. Also returns a field role, which - specifies special instructions for reconstruction.""" - - assert field.init, "Field must be in class constructor" - - # Create initial argument. - arg = ArgumentDefinition( - name=field.name, - field=field, - parent_class=parent_class, - type=field.type, - ) - - # Propagate argument through transforms until stable. - prev_arg = arg - role: _construction.FieldRole = _construction.FieldRole.VANILLA_FIELD - - def _handle_generics(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - if isinstance(arg.type, TypeVar): - assert arg.type in type_from_typevar, "TypeVar not bounded" - return ( - dataclasses.replace( - arg, type=type_from_typevar[arg.type] # type:ignore - ), - None, - ) - else: - return arg, None - - while True: - for transform in [_handle_generics] + _argument_transforms: # type: ignore - # Apply transform. - arg, new_role = transform(arg) - - # Update field role. - if new_role is not None: - assert ( - role == _construction.FieldRole.VANILLA_FIELD - ), "Something went wrong -- only one field role can be specified per argument!" - role = new_role - - # Stability check. - if arg == prev_arg: - break - prev_arg = arg - return arg, role - - -# Argument transformations. -# Each transform returns an argument definition and (optionall) a special role for -# reconstruction -- note that a field can only ever have one role. - -_ArgumentTransformOutput = Tuple[ArgumentDefinition, Optional[_construction.FieldRole]] - - -def _unwrap_final(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Treat Final[T] as just T.""" - if hasattr(arg.type, "__origin__") and arg.type.__origin__ is Final: # type: ignore - (typ,) = arg.type.__args__ # type: ignore - return ( - dataclasses.replace( - arg, - type=typ, - ), - None, - ) - else: - return arg, None - - -def _unwrap_annotated(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Treat Annotated[T, annotation] as just T.""" - if hasattr(arg.type, "__class__") and arg.type.__class__ == _AnnotatedAlias: - typ = arg.type.__origin__ # type: ignore - return ( - dataclasses.replace( - arg, - type=typ, - ), - None, - ) - else: - return arg, None - - -def _handle_optionals(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Transform for handling Optional[T] types. Sets default to None and marks arg as - not required.""" - if hasattr(arg.type, "__origin__") and arg.type.__origin__ is Union: # type: ignore - options = set(arg.type.__args__) # type: ignore - assert ( - len(options) == 2 and type(None) in options - ), "Union must be either over dataclasses (for subparsers) or Optional" - (typ,) = options - {type(None)} - required = False - return ( - dataclasses.replace( - arg, - type=typ, - required=required, - ), - None, - ) - else: - return arg, None - - -def _populate_defaults(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Populate default values.""" - if arg.default is not None: - # Skip if another handler has already populated the default. - return arg, None - - default = None - required = True - if arg.field.default is not dataclasses.MISSING: - default = arg.field.default - required = False - elif arg.field.default_factory is not dataclasses.MISSING: # type: ignore - default = arg.field.default_factory() # type: ignore - required = False - - if arg.required is not None: - required = arg.required - - return dataclasses.replace(arg, default=default, required=required), None - - -def _bool_flags(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """For booleans, we use a `store_true` action.""" - if arg.type != bool: - return arg, None - - if arg.default is None: - return ( - dataclasses.replace( - arg, - type=_strings.bool_from_string, # type: ignore - metavar="{True,False}", - ), - None, - ) - elif arg.default is False: - return ( - dataclasses.replace( - arg, - action="store_true", - type=None, - ), - None, - ) - elif arg.default is True: - return ( - dataclasses.replace( - arg, - dest=arg.name, - name="no_" + arg.name, - action="store_false", - type=None, - ), - None, - ) - else: - assert False, "Invalid default" - - -def _nargs_from_sequences_and_lists( - arg: ArgumentDefinition, -) -> _ArgumentTransformOutput: - """Transform for handling Sequence[T] and list types.""" - if hasattr(arg.type, "__origin__") and arg.type.__origin__ in ( # type: ignore - collections.abc.Sequence, # different from typing.Sequence! - list, # different from typing.List! - ): - assert len(arg.type.__args__) == 1 # type: ignore - (typ,) = arg.type.__args__ # type: ignore - - return ( - dataclasses.replace( - arg, - type=typ, - # `*` is >=0 values, `+` is >=1 values - # We're going to require at least 1 value; if a user wants to accept no - # input, they can use Optional[Tuple[...]] - nargs="+", - ), - None, - ) - else: - return arg, None - - -def _nargs_from_tuples(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Transform for handling Tuple[T, T, ...] types.""" - if hasattr(arg.type, "__origin__") and arg.type.__origin__ == tuple: # type: ignore - argset = set(arg.type.__args__) # type: ignore - argset_no_ellipsis = argset - {Ellipsis} # - assert len(argset_no_ellipsis) == 1, "Tuples must be of a single type!" - - if argset != argset_no_ellipsis: - # `*` is >=0 values, `+` is >=1 values. - # We're going to require at least 1 value; if a user wants to accept no - # input, they can use Optional[Tuple[...]]. - nargs = "+" - else: - nargs = len(arg.type.__args__) # type: ignore - (typ,) = argset_no_ellipsis - - return ( - dataclasses.replace( - arg, - nargs=nargs, - type=typ, - ), - _construction.FieldRole.TUPLE, - ) - else: - return arg, None - - -def _choices_from_literals(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """For literal types, set choices.""" - if hasattr(arg.type, "__origin__") and arg.type.__origin__ is Literal: # type: ignore - choices = set(arg.type.__args__) # type: ignore - assert ( - len(set(map(type, choices))) == 1 - ), "All choices in literal must have the same type!" - return ( - dataclasses.replace( - arg, - type=type(arg.type.__args__[0]), # type: ignore - choices=choices, - ), - None, - ) - else: - return arg, None - - -def _enums_as_strings(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """For enums, use string representations.""" - if isinstance(arg.type, type) and issubclass(arg.type, enum.Enum): - if arg.choices is None: - choices = set(x.name for x in arg.type) - else: - choices = set(x.name for x in arg.choices) - - return ( - dataclasses.replace( - arg, - choices=choices, - type=str, - default=None if arg.default is None else arg.default.name, - ), - _construction.FieldRole.ENUM, - ) - else: - return arg, None - - -def _generate_helptext(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Generate helptext from docstring and argument name.""" - if arg.help is None: - help_parts = [] - docstring_help = _docstrings.get_field_docstring( - arg.parent_class, arg.field.name - ) - if docstring_help is not None: - help_parts.append(docstring_help) - - if arg.default is not None and hasattr(arg.default, "name"): - # Special case for enums. - help_parts.append(f"(default: {arg.default.name})") - elif arg.default is not None: - # General case. - help_parts.append("(default: %(default)s)") - - return dataclasses.replace(arg, help=" ".join(help_parts)), None - else: - return arg, None - - -def _use_type_as_metavar(arg: ArgumentDefinition) -> _ArgumentTransformOutput: - """Communicate the argument type using the metavar.""" - if hasattr(arg.type, "__name__") and arg.choices is None and arg.metavar is None: - return ( - dataclasses.replace(arg, metavar=arg.type.__name__.upper()), # type: ignore - None, - ) - else: - return arg, None - - -_argument_transforms: List[Callable[[ArgumentDefinition], _ArgumentTransformOutput]] = [ - _unwrap_final, - _unwrap_annotated, - _handle_optionals, - _populate_defaults, - _bool_flags, - _nargs_from_sequences_and_lists, - _nargs_from_tuples, - _choices_from_literals, - _enums_as_strings, - _generate_helptext, - _use_type_as_metavar, -] diff --git a/dcargs/_construction.py b/dcargs/_construction.py deleted file mode 100644 index 1c264914f..000000000 --- a/dcargs/_construction.py +++ /dev/null @@ -1,122 +0,0 @@ -import dataclasses -import enum -from typing import Any, Dict, Set, Tuple, Type, TypeVar, Union - -from typing_extensions import _GenericAlias # type: ignore - -from . import _resolver, _strings - -DataclassType = TypeVar("DataclassType", bound=Union[Type, _GenericAlias]) - - -class FieldRole(enum.Enum): - """Enum for specifying special behaviors for .""" - - VANILLA_FIELD = enum.auto() - TUPLE = enum.auto() - ENUM = enum.auto() - NESTED_DATACLASS = enum.auto() # Singular nested dataclass. - SUBPARSERS = enum.auto() # Unions over dataclasses. - - -@dataclasses.dataclass -class ConstructionMetadata: - """Metadata recorded during parsing that's needed for reconstructing dataclasses.""" - - role_from_field: Dict[dataclasses.Field, FieldRole] = dataclasses.field( - default_factory=dict - ) - subparser_name_from_type: Dict[Type, str] = dataclasses.field(default_factory=dict) - - def update(self, other: "ConstructionMetadata") -> None: - self.role_from_field.update(other.role_from_field) - self.subparser_name_from_type.update(other.subparser_name_from_type) - - -def construct_dataclass( - cls: Type[DataclassType], - value_from_arg: Dict[str, Any], - metadata: ConstructionMetadata, - field_name_prefix: str = "", -) -> Tuple[DataclassType, Set[str]]: - """Construct a dataclass object from a dictionary of values from argparse. - - Returns dataclass object and set of used arguments.""" - - assert _resolver.is_dataclass(cls) - - cls, _type_from_typevar = _resolver.resolve_generic_dataclasses(cls) - - kwargs: Dict[str, Any] = {} - consumed_keywords: Set[str] = set() - - def get_value_from_arg(arg: str) -> Any: - """Helper for getting values from `value_from_arg` + doing some extra - asserts.""" - assert arg in value_from_arg - assert arg not in consumed_keywords - consumed_keywords.add(arg) - return value_from_arg[arg] - - for field in _resolver.resolved_fields(cls): # type: ignore - if not field.init: - continue - - value: Any - role = metadata.role_from_field[field] - - prefixed_field_name = field_name_prefix + field.name - - if role is FieldRole.ENUM: - # Handle enums. - value = field.type[get_value_from_arg(prefixed_field_name)] - elif role is FieldRole.NESTED_DATACLASS: - # Nested dataclasses. - value, consumed_keywords_child = construct_dataclass( - field.type, - value_from_arg, - metadata, - field_name_prefix=prefixed_field_name - + _strings.NESTED_DATACLASS_DELIMETER, - ) - consumed_keywords |= consumed_keywords_child - elif role is FieldRole.SUBPARSERS: - # Unions over dataclasses (subparsers). - subparser_dest = _strings.SUBPARSER_DEST_FMT.format( - name=prefixed_field_name - ) - subparser_name = get_value_from_arg(subparser_dest) - if subparser_name is None: - # No subparser selected -- this should only happen when we do either - # Optional[Union[A, B, ...]] or Union[A, B, None]. - assert type(None) in field.type.__args__ - value = None - else: - options = field.type.__args__ - chosen_cls = None - for option in options: - if metadata.subparser_name_from_type[option] == subparser_name: - chosen_cls = option - break - assert chosen_cls is not None - value, consumed_keywords_child = construct_dataclass( - chosen_cls, - value_from_arg, - metadata, - ) - consumed_keywords |= consumed_keywords_child - elif role is FieldRole.TUPLE: - # For sequences, argparse always gives us lists. Sometimes we want tuples. - value = get_value_from_arg(prefixed_field_name) - if value is not None: - value = tuple(value) - - elif role is FieldRole.VANILLA_FIELD: - # General case. - value = get_value_from_arg(prefixed_field_name) - else: - assert False - - kwargs[field.name] = value - - return cls(**kwargs), consumed_keywords # type: ignore diff --git a/dcargs/_docstrings.py b/dcargs/_docstrings.py deleted file mode 100644 index b90baecb2..000000000 --- a/dcargs/_docstrings.py +++ /dev/null @@ -1,135 +0,0 @@ -import dataclasses -import functools -import inspect -import io -import tokenize -from typing import Dict, List, Optional, Type - -from typing_extensions import _GenericAlias # type: ignore - -from . import _strings - - -@dataclasses.dataclass -class _Token: - token_type: int - token: str - line_number: int - - -@dataclasses.dataclass -class _FieldData: - index: int - line_number: int - prev_field_line_number: int - - -@dataclasses.dataclass -class _Tokenization: - tokens: List[_Token] - tokens_from_line: Dict[int, List[_Token]] - field_data_from_name: Dict[str, _FieldData] - - @staticmethod - @functools.lru_cache(maxsize=4) - def make(cls) -> "_Tokenization": - """Parse the source code of a class, and cache some tokenization information.""" - readline = io.BytesIO(inspect.getsource(cls).encode("utf-8")).readline - - tokens: List[_Token] = [] - tokens_from_line: Dict[int, List[_Token]] = {1: []} - field_data_from_name: Dict[str, _FieldData] = {} - - line_number: int = 1 - for toktype, tok, start, end, line in tokenize.tokenize(readline): - if toktype in (tokenize.NEWLINE, tokenize.NL): - line_number += 1 - tokens_from_line[line_number] = [] - elif toktype is not tokenize.INDENT: - token = _Token(token_type=toktype, token=tok, line_number=line_number) - tokens.append(token) - tokens_from_line[line_number].append(token) - - prev_field_line_number: int = 1 - for i, token in enumerate(tokens[:-1]): - if token.token_type == tokenize.NAME: - # Naive heuristic for field names - # This will currently catch variable/argument annotations as well - if ( - tokens[i + 1].token == ":" - and token.token not in field_data_from_name - ): - field_data_from_name[token.token] = _FieldData( - index=i, - line_number=token.line_number, - prev_field_line_number=prev_field_line_number, - ) - prev_field_line_number = token.line_number - - return _Tokenization( - tokens=tokens, - tokens_from_line=tokens_from_line, - field_data_from_name=field_data_from_name, - ) - - -def get_field_docstring(cls: Type, field_name: str) -> Optional[str]: - """Get docstring for a field in a class.""" - - if isinstance(cls, _GenericAlias): - cls = cls.__origin__ - - assert dataclasses.is_dataclass(cls) - try: - tokenization = _Tokenization.make(cls) # type: ignore - except OSError as e: - # Dynamic dataclasses will result in an OSError -- this is fine, we just assume - # there's no docstring. - assert "could not find class definition" in e.args[0] - return None - - # Grab field-specific tokenization data. - assert ( - field_name in tokenization.field_data_from_name - ), "Docstring parsing error -- this usually means that there are multiple \ - dataclasses in the same file with the same name but different scopes." - field_data = tokenization.field_data_from_name[field_name] - - # Check for docstring-style comment. - if ( - field_data.line_number + 1 in tokenization.tokens_from_line - and len(tokenization.tokens_from_line[field_data.line_number + 1]) > 0 - ): - first_token_on_next_line = tokenization.tokens_from_line[ - field_data.line_number + 1 - ][0] - if first_token_on_next_line.token_type == tokenize.STRING: - docstring = first_token_on_next_line.token.strip() - assert docstring.endswith('"""') and docstring.startswith('"""') - return _strings.dedent(docstring[3:-3]) - - # Check for comment on the same line as the field. - final_token_on_line = tokenization.tokens_from_line[field_data.line_number][-1] - if final_token_on_line.token_type == tokenize.COMMENT: - comment: str = final_token_on_line.token - assert comment.startswith("#") - return comment[1:].strip() - - # Check for comment on the line before the field. - comment_index = field_data.index - comments: List[str] = [] - while True: - comment_index -= 1 - comment_token = tokenization.tokens[comment_index] - if ( - comment_token.token_type == tokenize.COMMENT - and comment_token.line_number > field_data.prev_field_line_number - ): - assert comment_token.token.startswith("#") - comments.append(comment_token.token[1:].strip()) - else: - break - if len(comments) > 0: - return "\n".join(comments[::-1]) - - return None diff --git a/dcargs/_parse.py b/dcargs/_parse.py deleted file mode 100644 index 0e0894ae1..000000000 --- a/dcargs/_parse.py +++ /dev/null @@ -1,41 +0,0 @@ -import argparse -from typing import Optional, Sequence, Type, TypeVar, Union - -from typing_extensions import _GenericAlias # type: ignore - -from . import _construction, _parsers, _strings - -DataclassType = TypeVar("DataclassType", bound=Union[Type, _GenericAlias]) - - -def parse( - cls: Type[DataclassType], - description: str = "", - args: Optional[Sequence[str]] = None, -) -> DataclassType: - """Populate a dataclass via CLI args.""" - - parser_definition, construction_metadata = _parsers.ParserDefinition.from_dataclass( - cls, - parent_dataclasses=set(), # Used for recursive calls. - subparser_name_from_type={}, # Aliases for subparsers; this is working, but not yet exposed. - parent_type_from_typevar=None, # Used for recursive calls. - ) - - parser = argparse.ArgumentParser( - description=_strings.dedent(description), - formatter_class=argparse.RawTextHelpFormatter, - ) - parser_definition.apply(parser) - - namespace = parser.parse_args(args) - - value_from_arg = vars(namespace) - out, consumed_keywords = _construction.construct_dataclass( - cls, value_from_arg, construction_metadata - ) - assert ( - consumed_keywords == value_from_arg.keys() - ), "Not all arguments were consumed!" - - return out diff --git a/dcargs/_parsers.py b/dcargs/_parsers.py deleted file mode 100644 index f830c3ef8..000000000 --- a/dcargs/_parsers.py +++ /dev/null @@ -1,173 +0,0 @@ -import argparse -import dataclasses -from typing import Any, Dict, List, Optional, Set, Tuple, Type, TypeVar, Union - -from typing_extensions import _GenericAlias # type: ignore - -from . import _arguments, _construction, _docstrings, _resolver, _strings - - -@dataclasses.dataclass -class ParserDefinition: - """Each parser contains a list of arguments and optionally a subparser.""" - - args: List["_arguments.ArgumentDefinition"] - subparsers: Optional["SubparsersDefinition"] - - def apply(self, parser: argparse.ArgumentParser) -> None: - """Create defined arguments and subparsers.""" - - # Put required group at start of group list. - required_group = parser.add_argument_group("required arguments") - parser._action_groups = parser._action_groups[::-1] - - # Add each argument. - for arg in self.args: - if arg.required: - arg.add_argument(required_group) - else: - arg.add_argument(parser) - - # Add subparsers. - if self.subparsers is not None: - subparsers = parser.add_subparsers( - dest=_strings.SUBPARSER_DEST_FMT.format(name=self.subparsers.name), - description=self.subparsers.description, - required=self.subparsers.required, - ) - for name, subparser_def in self.subparsers.parsers.items(): - subparser = subparsers.add_parser(name) - subparser_def.apply(subparser) - - @staticmethod - def from_dataclass( - cls: Union[Type[Any], _GenericAlias], - parent_dataclasses: Optional[Set[Type]], - subparser_name_from_type: Dict[Type, str], - parent_type_from_typevar: Optional[Dict[TypeVar, Type]], - ) -> Tuple["ParserDefinition", _construction.ConstructionMetadata]: - """Create a parser definition from a dataclass.""" - - if parent_dataclasses is None: - parent_dataclasses = set() - - assert _resolver.is_dataclass(cls) - - cls, type_from_typevar = _resolver.resolve_generic_dataclasses(cls) - - if parent_type_from_typevar is not None: - for typevar, typ in type_from_typevar.items(): - if typ in parent_type_from_typevar: - type_from_typevar[typevar] = parent_type_from_typevar[typ] # type: ignore - - assert ( - cls not in parent_dataclasses - ), f"Found a cyclic dataclass dependency with type {cls}" - - args = [] - subparsers = None - metadata = _construction.ConstructionMetadata() - for field in _resolver.resolved_fields(cls): # type: ignore - if not field.init: - continue - - # If set to False, we don't directly create an argument from this field. - arg_from_field: bool = True - - # Add arguments for nested dataclasses. - if _resolver.is_dataclass(field.type): - child_definition, child_metadata = ParserDefinition.from_dataclass( - field.type, - parent_dataclasses | {cls}, - subparser_name_from_type=subparser_name_from_type, - parent_type_from_typevar=type_from_typevar, - ) - metadata.update(child_metadata) - - child_args = child_definition.args - for i, arg in enumerate(child_args): - child_args[i] = arg.prefix( - field.name + _strings.NESTED_DATACLASS_DELIMETER - ) - - args.extend(child_args) - - if child_definition.subparsers is not None: - assert subparsers is None - subparsers = child_definition.subparsers - - metadata.role_from_field[ - field - ] = _construction.FieldRole.NESTED_DATACLASS - arg_from_field = False - - # Union of dataclasses should create subparsers. - if hasattr(field.type, "__origin__") and field.type.__origin__ is Union: - # We don't use sets here to retain order of subcommands. - options = field.type.__args__ - options_no_none = [o for o in options if o != type(None)] # noqa - if len(options_no_none) >= 2 and all( - map(_resolver.is_dataclass, options_no_none) - ): - assert ( - subparsers is None - ), "Only one subparser group is supported per dataclass" - - parsers: Dict[str, ParserDefinition] = {} - for option in options_no_none: - subparser_name = ( - subparser_name_from_type[option] - if option in subparser_name_from_type - else _strings.hyphen_separated_from_camel_case( - option.__name__ - ) - ) - metadata.subparser_name_from_type[option] = subparser_name - - ( - parsers[subparser_name], - child_metadata, - ) = ParserDefinition.from_dataclass( - option, - parent_dataclasses | {cls}, - subparser_name_from_type, - parent_type_from_typevar=type_from_typevar, - ) - metadata.update(child_metadata) - - subparsers = SubparsersDefinition( - name=field.name, - description=_docstrings.get_field_docstring(cls, field.name), - parsers=parsers, - required=( - options == options_no_none - ), # Not required if no options. - ) - metadata.role_from_field[field] = _construction.FieldRole.SUBPARSERS - arg_from_field = False - - # Make an argument! - if arg_from_field: - arg, role = _arguments.ArgumentDefinition.make_from_field( - cls, field, type_from_typevar - ) - args.append(arg) - metadata.role_from_field[field] = role - - return ( - ParserDefinition( - args=args, - subparsers=subparsers, - ), - metadata, - ) - - -@dataclasses.dataclass -class SubparsersDefinition: - """Structure for containing subparsers. Each subparser is a parser with a name.""" - - name: str - description: Optional[str] - parsers: Dict[str, ParserDefinition] - required: bool diff --git a/dcargs/_resolver.py b/dcargs/_resolver.py deleted file mode 100644 index 9687ba0a1..000000000 --- a/dcargs/_resolver.py +++ /dev/null @@ -1,48 +0,0 @@ -import copy -import dataclasses -import functools -from typing import Dict, List, Tuple, Type, TypeVar, Union - -from typing_extensions import _GenericAlias, get_type_hints # type: ignore - - -def is_dataclass(cls: Union[Type, _GenericAlias]) -> bool: - """Same as `dataclasses.is_dataclass`, but also handles generic aliases.""" - return dataclasses.is_dataclass(cls) or ( - isinstance(cls, _GenericAlias) and dataclasses.is_dataclass(cls.__origin__) - ) - - -def resolve_generic_dataclasses( - cls: Union[Type, _GenericAlias], -) -> Tuple[Type, Dict[TypeVar, Type]]: - """If the input is a dataclass: no-op. If it's a generic alias: returns the root - dataclass, and a mapping from typevars to concrete types.""" - - if isinstance(cls, _GenericAlias): - typevars = cls.__origin__.__parameters__ - typevar_values = cls.__args__ - assert len(typevars) == len(typevar_values) - cls = cls.__origin__ - return cls, dict(zip(typevars, typevar_values)) - else: - return cls, {} - - -@functools.lru_cache(maxsize=16) -def resolved_fields(cls: Union[Type, _GenericAlias]) -> List[dataclasses.Field]: - """Similar to dataclasses.fields, but resolves forward references.""" - - assert dataclasses.is_dataclass(cls) - fields = [] - annotations = get_type_hints(cls) - for field in dataclasses.fields(cls): - # Avoid mutating original field. - field = copy.copy(field) - - # Resolve forward references. - field.type = annotations[field.name] - - fields.append(field) - - return fields diff --git a/dcargs/_strings.py b/dcargs/_strings.py deleted file mode 100644 index 69bb2292c..000000000 --- a/dcargs/_strings.py +++ /dev/null @@ -1,33 +0,0 @@ -import functools -import re -import textwrap - -NESTED_DATACLASS_DELIMETER: str = "." -SUBPARSER_DEST_FMT: str = "{name} (positional)" - - -def dedent(text: str) -> str: - """Same as textwrap.dedent, but ignores the first line.""" - first_line, line_break, rest = text.partition("\n") - if line_break == "": - return textwrap.dedent(text) - return f"{first_line.strip()}\n{textwrap.dedent(rest)}" - - -_camel_separator_pattern = functools.lru_cache(maxsize=1)( - lambda: re.compile("((?<=[a-z0-9])[A-Z]|(?!^)[A-Z](?=[a-z]))") -) - - -def hyphen_separated_from_camel_case(name: str) -> str: - return _camel_separator_pattern().sub(r"-\1", name).lower() - - -def bool_from_string(text: str) -> bool: - text = text.lower() - if text in ("true", "1"): - return True - elif text in ("false", "0"): - return False - else: - raise ValueError(f"Boolean value expected, but got {text}.") diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 000000000..567609b12 --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1 @@ +build/ diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 000000000..201f60130 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line. +SPHINXOPTS = +SPHINXBUILD = sphinx-build +SPHINXPROJ = fannypack +SOURCEDIR = source +BUILDDIR = ./build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 000000000..7260ab5aa --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,7 @@ +sphinx==7.2.6 +furo==2023.9.10 +docutils==0.20.1 +sphinx-autoapi==3.0.0 +m2r2==0.3.3.post2 +git+https://github.com/brentyi/sphinxcontrib-programoutput.git +git+https://github.com/brentyi/ansi.git diff --git a/docs/source/_static/css/compact_table_header.css b/docs/source/_static/css/compact_table_header.css new file mode 100644 index 000000000..2861b5888 --- /dev/null +++ b/docs/source/_static/css/compact_table_header.css @@ -0,0 +1,30 @@ +.sidebar-logo { + max-width: 150px; + margin: 0; +} +.sidebar-brand-text { + display: None; +} + +/* This stylesheet is currently overfit to the comparison table. Should be + * revisited if we add more tables to the documentation. */ +table.docutils thead, +table.docutils tr, +table.docutils th { + background: none; +} +table.docutils thead th { + position: relative; + height: 12em; + text-align: left; + border: 0; +} +table.docutils thead th p { + transform: rotate(300deg); + width: 1em; + overflow: visible; + white-space: nowrap; + position: absolute; + bottom: 0; + z-index: 1; +} diff --git a/docs/source/_static/logo-dark.svg b/docs/source/_static/logo-dark.svg new file mode 100644 index 000000000..5143f41ac --- /dev/null +++ b/docs/source/_static/logo-dark.svg @@ -0,0 +1,55 @@ + + + + + + + + + + diff --git a/docs/source/_static/logo-light.svg b/docs/source/_static/logo-light.svg new file mode 100644 index 000000000..7ca6345f1 --- /dev/null +++ b/docs/source/_static/logo-light.svg @@ -0,0 +1,56 @@ + + + + + + + + + + diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 000000000..53f889d5d --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,398 @@ +# -*- coding: utf-8 -*- +# +# Configuration file for the Sphinx documentation builder. +# +# This file does only contain a selection of the most common options. For a +# full list see the documentation: +# http://www.sphinx-doc.org/en/stable/config + +from typing import Dict, List + +import m2r2 + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +# + +# -- Project information ----------------------------------------------------- + +project = "tyro" +copyright = "2022" +author = "brentyi" + +# The short X.Y version +version = "" +# The full version, including alpha/beta/rc tags +release = "" + + +# -- General configuration --------------------------------------------------- + +napoleon_numpy_docstring = False # Force consistency, leave only Google +napoleon_use_rtype = False # More legible + +# If your documentation needs a minimal Sphinx version, state it here. +# +# needs_sphinx = '1.0' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.todo", + "sphinx.ext.coverage", + # "sphinx.ext.mathjax", + "sphinx.ext.githubpages", + "sphinx.ext.napoleon", + # "sphinx.ext.inheritance_diagram", + "autoapi.extension", + "sphinx.ext.viewcode", + "m2r2", + "sphinxcontrib.programoutput", + "sphinxcontrib.ansi", +] +programoutput_use_ansi = True +html_ansi_stylesheet = "black-on-white.css" +html_static_path = ["_static"] +html_theme_options = { + "light_css_variables": { + "color-code-background": "#f4f4f4", + "color-code-foreground": "#000", + }, + "footer_icons": [ + { + "name": "GitHub", + "url": "https://github.com/brentyi/tyro", + "html": """ + + + + """, + "class": "", + }, + ], + "light_logo": "logo-light.svg", + "dark_logo": "logo-dark.svg", +} + +# Pull documentation types from hints +autodoc_typehints = "both" + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# +source_suffix = [".rst", ".md"] +# source_suffix = ".rst" + +# The master toctree document. +master_doc = "index" + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language: str = "en" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path . +exclude_patterns: List[str] = [] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "monokai" + + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "furo" +html_title = "tyro" + + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +# +# html_theme_options = {} + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ["_static"] + +# Custom sidebar templates, must be a dictionary that maps document names +# to template names. +# +# The default sidebars (for documents that don't match any pattern) are +# defined by theme itself. Builtin themes are using these templates by +# default: ``['localtoc.html', 'relations.html', 'sourcelink.html', +# 'searchbox.html']``. +# +# html_sidebars = {} + + +# -- Options for HTMLHelp output --------------------------------------------- + +# Output file base name for HTML help builder. +htmlhelp_basename = "tyro_doc" + + +# -- Options for Github output ------------------------------------------------ + +sphinx_to_github = True +sphinx_to_github_verbose = True +sphinx_to_github_encoding = "utf-8" + + +# -- Options for LaTeX output ------------------------------------------------ + +latex_elements: Dict[str, str] = { + # The paper size ('letterpaper' or 'a4paper'). + # + # 'papersize': 'letterpaper', + # The font size ('10pt', '11pt' or '12pt'). + # + # 'pointsize': '10pt', + # Additional stuff for the LaTeX preamble. + # + # 'preamble': '', + # Latex figure (float) alignment + # + # 'figure_align': 'htbp', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [ + ( + master_doc, + "tyro.tex", + "tyro", + "brentyi", + "manual", + ), +] + + +# -- Options for manual page output ------------------------------------------ + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [(master_doc, "tyro", "tyro documentation", [author], 1)] + + +# -- Options for Texinfo output ---------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + ( + master_doc, + "tyro", + "tyro", + author, + "tyro", + "tyro", + "Miscellaneous", + ), +] + + +# -- Extension configuration -------------------------------------------------- + +# -- Options for autoapi extension -------------------------------------------- +autoapi_dirs = ["../../src/tyro"] +autoapi_root = "api" +autoapi_options = [ + "members", + "undoc-members", + "imported-members", + "show-inheritance", + # "show-inheritance-diagram", + "special-members", + # "show-module-summary", +] +autoapi_add_toctree_entry = False + +# # Generate name aliases +# def _gen_name_aliases(): +# """Generate a name alias dictionary, which maps private names to ones in the public +# API. A little bit hardcoded/hacky.""" +# +# name_alias = {} +# +# def recurse(module, prefixes): +# if hasattr(module, "__name__") and module.__name__.startswith("tyro"): +# MAX_DEPTH = 5 +# if len(prefixes) > MAX_DEPTH: +# # Prevent infinite loops from cyclic imports +# return +# else: +# return +# +# for member_name in dir(module): +# if member_name == "tyro": +# continue +# +# member = getattr(module, member_name) +# if callable(member): +# full_name = ".".join(["tyro"] + prefixes + [member_name]) +# +# shortened_name = "tyro" +# current = tyro +# success = True +# for p in prefixes + [member_name]: +# if p.startswith("_"): +# continue +# if not hasattr(current, p): +# success = False +# break +# current = getattr(current, p) +# shortened_name += "." + p +# +# if success and shortened_name != full_name: +# if full_name in name_alias: +# assert full_name == name_alias[shortened_name], full_name +# else: +# name_alias[full_name] = shortened_name +# elif not member_name.startswith("__"): +# recurse(member, prefixes + [member_name]) +# +# import tyro +# +# recurse(tyro, prefixes=[]) +# return name_alias +# +# +# _name_aliases = _gen_name_aliases() +# +# # Set inheritance_alias setting for inheritance diagrams +# inheritance_alias = _name_aliases +# +# +# def _apply_name_aliases(name: Optional[str]) -> Optional[str]: +# if name is None: +# return None +# +# name = name.strip() +# +# if "[" in name: +# # Generics. +# name, _, suffix = name.partition("[") +# assert suffix[-1] == "]" +# suffix = suffix[:-1] +# if "," in suffix: +# suffix = ", ".join(_apply_name_aliases(x) for x in suffix.split(",")) # type: ignore +# else: +# suffix = _apply_name_aliases(suffix) # type: ignore +# suffix = "[" + suffix + "]" +# else: +# suffix = "" +# +# if name in _name_aliases: +# name = _name_aliases[name] +# +# return name + suffix # type: ignore +# +# +# # Apply our inheritance alias to autoapi base classes +# def _override_class_documenter(): +# import autoapi +# +# orig_init = autoapi.mappers.python.PythonClass.__init__ +# +# def __init__(self, obj, **kwargs): +# bases = obj["bases"] +# for i in range(len(bases)): +# bases[i] = _apply_name_aliases(bases[i]) +# +# args = obj["args"] +# if args is not None: +# for i in range(len(args)): +# assert isinstance(args[i], tuple) and len(args[i]) == 4 +# args[i] = ( +# args[i][0], +# args[i][1], +# _apply_name_aliases(args[i][2]), +# args[i][3], +# ) +# orig_init(self, obj, **kwargs) +# +# autoapi.mappers.python.PythonClass.__init__ = __init__ +# +# +# _override_class_documenter() +# +# +# # Apply our inheritance alias to autoapi type annotations +# def _override_function_documenter(): +# import autoapi +# +# orig_init = autoapi.mappers.python.PythonFunction.__init__ +# +# def __init__(self, obj, **kwargs): +# args = obj["args"] +# if args is not None: +# for i in range(len(args)): +# assert isinstance(args[i], tuple) and len(args[i]) == 4 +# args[i] = ( +# args[i][0], +# args[i][1], +# _apply_name_aliases(args[i][2]), +# args[i][3], +# ) +# +# obj["return_annotation"] = _apply_name_aliases(obj["return_annotation"]) +# orig_init(self, obj, **kwargs) +# +# autoapi.mappers.python.PythonFunction.__init__ = __init__ +# +# +# _override_function_documenter() +# +# # Apply our inheritance alias to autoapi attribute annotations +# def _override_attribute_documenter(): +# import autoapi +# +# orig_init = autoapi.mappers.python.PythonAttribute.__init__ +# +# def __init__(self, obj, **kwargs): +# obj["annotation"] = _apply_name_aliases(obj["annotation"]) +# orig_init(self, obj, **kwargs) +# +# autoapi.mappers.python.PythonAttribute.__init__ = __init__ +# +# +# _override_attribute_documenter() + + +# -- Options for todo extension ---------------------------------------------- + +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = True + +# -- Enable Markdown -> RST conversion ---------------------------------------- + + +def docstring(app, what, name, obj, options, lines): + md = "\n".join(lines) + rst = m2r2.convert(md) + lines.clear() + lines += rst.splitlines() # type: ignore + + +def setup(app): + app.connect("autodoc-process-docstring", docstring) + app.add_css_file("css/compact_table_header.css") diff --git a/docs/source/examples/01_basics/01_functions.rst b/docs/source/examples/01_basics/01_functions.rst new file mode 100644 index 000000000..e5729e93b --- /dev/null +++ b/docs/source/examples/01_basics/01_functions.rst @@ -0,0 +1,56 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Functions +========================================== + +In the simplest case, :func:`tyro.cli()` can be used to run a function with +arguments populated from the CLI. + + +.. code-block:: python + :linenos: + + + import tyro + + + def main( + field1: str, + field2: int = 3, + ) -> None: + """Function, whose arguments will be populated from a CLI interface. + + Args: + field1: A string field. + field2: A numeric field, with a default value. + """ + print(field1, field2) + + + if __name__ == "__main__": + tyro.cli(main) + +------------ + +.. raw:: html + + python 01_basics/01_functions.py --help + +.. program-output:: python ../../examples/01_basics/01_functions.py --help + +------------ + +.. raw:: html + + python 01_basics/01_functions.py --field1 hello + +.. program-output:: python ../../examples/01_basics/01_functions.py --field1 hello + +------------ + +.. raw:: html + + python 01_basics/01_functions.py --field1 hello --field2 10 + +.. program-output:: python ../../examples/01_basics/01_functions.py --field1 hello --field2 10 diff --git a/docs/source/examples/01_basics/02_dataclasses.rst b/docs/source/examples/01_basics/02_dataclasses.rst new file mode 100644 index 000000000..01f7acb78 --- /dev/null +++ b/docs/source/examples/01_basics/02_dataclasses.rst @@ -0,0 +1,57 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Dataclasses +========================================== + +Common pattern: use :func:`tyro.cli()` to instantiate a dataclass. + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + + @dataclasses.dataclass + class Args: + """Description. + This should show up in the helptext!""" + + field1: str + """A string field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 01_basics/02_dataclasses.py --help + +.. program-output:: python ../../examples/01_basics/02_dataclasses.py --help + +------------ + +.. raw:: html + + python 01_basics/02_dataclasses.py --field1 hello + +.. program-output:: python ../../examples/01_basics/02_dataclasses.py --field1 hello + +------------ + +.. raw:: html + + python 01_basics/02_dataclasses.py --field1 hello --field2 5 + +.. program-output:: python ../../examples/01_basics/02_dataclasses.py --field1 hello --field2 5 diff --git a/docs/source/examples/01_basics/03_dataclasses_defaults.rst b/docs/source/examples/01_basics/03_dataclasses_defaults.rst new file mode 100644 index 000000000..ba87d9fac --- /dev/null +++ b/docs/source/examples/01_basics/03_dataclasses_defaults.rst @@ -0,0 +1,80 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Dataclasses + Defaults +========================================== + +The :code:`default=` argument can be used to override default values in dataclass +types. + + +.. warning:: + We advise against mutation of configuration objects from a dataclass's + :code:`__post_init__` method [#f1]_. In the example below, + :code:`__post_init__` would be called twice: once for the :code:`Args()` + object provided as a default value and another time for the :code:`Args()` + objected instantiated by :func:`tyro.cli()`. This can cause confusing + behavior! Instead, we show below one example of how derived fields can be + defined immutably. + + .. [#f1] Official Python docs for ``__post_init__`` can be found `here `_. + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + + @dataclasses.dataclass + class Args: + """Description. + This should show up in the helptext!""" + + field1: str + """A string field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + @property + def derived_field(self) -> str: + return ", ".join([self.field1] * self.field2) + + + if __name__ == "__main__": + args = tyro.cli( + Args, + default=Args( + field1="default string", + field2=tyro.MISSING, + ), + ) + print(args.derived_field) + +------------ + +.. raw:: html + + python 01_basics/03_dataclasses_defaults.py --help + +.. program-output:: python ../../examples/01_basics/03_dataclasses_defaults.py --help + +------------ + +.. raw:: html + + python 01_basics/03_dataclasses_defaults.py --field2 3 + +.. program-output:: python ../../examples/01_basics/03_dataclasses_defaults.py --field2 3 + +------------ + +.. raw:: html + + python 01_basics/03_dataclasses_defaults.py --field1 hello --field2 5 + +.. program-output:: python ../../examples/01_basics/03_dataclasses_defaults.py --field1 hello --field2 5 diff --git a/docs/source/examples/01_basics/04_collections.rst b/docs/source/examples/01_basics/04_collections.rst new file mode 100644 index 000000000..bd9cf55cd --- /dev/null +++ b/docs/source/examples/01_basics/04_collections.rst @@ -0,0 +1,61 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Multi-value Arguments +========================================== + +Arguments of both fixed and variable lengths can be annotated with standard +Python collection types. For Python 3.7 and 3.8, we can use either :code:`from +__future__ import annotations` to support :code:`list[T]` and :code:`tuple[T]`, +or the older API :code:`typing.List[T]` and :code:`typing.Tuple[T1, T2]`. + + +.. code-block:: python + :linenos: + + + import dataclasses + import pathlib + + import tyro + + + @dataclasses.dataclass(frozen=True) + class TrainConfig: + # Example of a variable-length tuple. `list[T]`, `set[T]`, + # `dict[K, V]`, etc are supported as well. + dataset_sources: tuple[pathlib.Path, ...] + """Paths to load training data from. This can be multiple!""" + + # Fixed-length tuples are also okay. + image_dimensions: tuple[int, int] = (32, 32) + """Height and width of some image data.""" + + + if __name__ == "__main__": + config = tyro.cli(TrainConfig) + print(config) + +------------ + +.. raw:: html + + python 01_basics/04_collections.py --help + +.. program-output:: python ../../examples/01_basics/04_collections.py --help + +------------ + +.. raw:: html + + python 01_basics/04_collections.py --dataset-sources ./data --image-dimensions 16 16 + +.. program-output:: python ../../examples/01_basics/04_collections.py --dataset-sources ./data --image-dimensions 16 16 + +------------ + +.. raw:: html + + python 01_basics/04_collections.py --dataset-sources ./data + +.. program-output:: python ../../examples/01_basics/04_collections.py --dataset-sources ./data diff --git a/docs/source/examples/01_basics/05_flags.rst b/docs/source/examples/01_basics/05_flags.rst new file mode 100644 index 000000000..433e6af72 --- /dev/null +++ b/docs/source/examples/01_basics/05_flags.rst @@ -0,0 +1,71 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Booleans and Flags +========================================== + +Booleans can either be expected to be explicitly passed in, or, if given a default +value, automatically converted to flags. + +To turn off conversion, see :class:`tyro.conf.FlagConversionOff`. + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + + @dataclasses.dataclass + class Args: + # Boolean. This expects an explicit "True" or "False". + boolean: bool + + # Optional boolean. Same as above, but can be omitted. + optional_boolean: bool | None = None + + # Pass --flag-a in to set this value to True. + flag_a: bool = False + + # Pass --no-flag-b in to set this value to False. + flag_b: bool = True + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 01_basics/05_flags.py --help + +.. program-output:: python ../../examples/01_basics/05_flags.py --help + +------------ + +.. raw:: html + + python 01_basics/05_flags.py --boolean True + +.. program-output:: python ../../examples/01_basics/05_flags.py --boolean True + +------------ + +.. raw:: html + + python 01_basics/05_flags.py --boolean False --flag-a + +.. program-output:: python ../../examples/01_basics/05_flags.py --boolean False --flag-a + +------------ + +.. raw:: html + + python 01_basics/05_flags.py --boolean False --no-flag-b + +.. program-output:: python ../../examples/01_basics/05_flags.py --boolean False --no-flag-b diff --git a/docs/source/examples/01_basics/06_literals.rst b/docs/source/examples/01_basics/06_literals.rst new file mode 100644 index 000000000..746e7ec61 --- /dev/null +++ b/docs/source/examples/01_basics/06_literals.rst @@ -0,0 +1,40 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Choices +========================================== + +:code:`typing.Literal[]` can be used to restrict inputs to a fixed set of literal choices. + + +.. code-block:: python + :linenos: + + + import dataclasses + from typing import Literal + + import tyro + + + @dataclasses.dataclass(frozen=True) + class Args: + # We can use Literal[] to restrict the set of allowable inputs, for example, over + # a set of strings. + strings: Literal["red", "green"] = "red" + + # Integers also work. (as well as booleans, enums, etc) + numbers: Literal[0, 1, 2] = 0 + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 01_basics/06_literals.py --help + +.. program-output:: python ../../examples/01_basics/06_literals.py --help diff --git a/docs/source/examples/01_basics/07_unions.rst b/docs/source/examples/01_basics/07_unions.rst new file mode 100644 index 000000000..8ce43f511 --- /dev/null +++ b/docs/source/examples/01_basics/07_unions.rst @@ -0,0 +1,57 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Unions +========================================== + +:code:`X | Y` or :code:`typing.Union[X, Y]` can be used to expand inputs to +multiple types. + + +.. code-block:: python + :linenos: + + + import dataclasses + import enum + from typing import Literal, Optional + + import tyro + + + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + + @dataclasses.dataclass(frozen=True) + class Args: + # Unions can be used to specify multiple allowable types. + union_over_types: int | str = 0 + string_or_enum: Literal["red", "green"] | Color = "red" + + # Unions also work over more complex nested types. + union_over_tuples: tuple[int, int] | tuple[str] = ("1",) + + # And can be nested in other types. + tuple_of_string_or_enum: tuple[Literal["red", "green"] | Color, ...] = ( + "red", + Color.RED, + ) + + # Optional[T] is equivalent to `T | None`. + integer: Optional[Literal[0, 1, 2, 3]] = None + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 01_basics/07_unions.py --help + +.. program-output:: python ../../examples/01_basics/07_unions.py --help diff --git a/docs/source/examples/01_basics/08_enums.rst b/docs/source/examples/01_basics/08_enums.rst new file mode 100644 index 000000000..08f142963 --- /dev/null +++ b/docs/source/examples/01_basics/08_enums.rst @@ -0,0 +1,62 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Enums +========================================== + +In addition to literals, enums can also be used to provide a fixed set of +choices. + + +.. code-block:: python + :linenos: + + + import dataclasses + import enum + + import tyro + + + class OptimizerType(enum.Enum): + ADAM = enum.auto() + SGD = enum.auto() + + + @dataclasses.dataclass(frozen=True) + class TrainConfig: + # Enums are handled seamlessly. + optimizer_type: OptimizerType = OptimizerType.ADAM + """Gradient-based optimizer to use.""" + + learning_rate: float = 1e-4 + """Learning rate for optimizer.""" + + + if __name__ == "__main__": + config = tyro.cli(TrainConfig) + print(config) + +------------ + +.. raw:: html + + python 01_basics/08_enums.py --help + +.. program-output:: python ../../examples/01_basics/08_enums.py --help + +------------ + +.. raw:: html + + python 01_basics/08_enums.py --optimizer-type SGD + +.. program-output:: python ../../examples/01_basics/08_enums.py --optimizer-type SGD + +------------ + +.. raw:: html + + python 01_basics/08_enums.py --optimizer-type ADAM --learning-rate 3e-4 + +.. program-output:: python ../../examples/01_basics/08_enums.py --optimizer-type ADAM --learning-rate 3e-4 diff --git a/docs/source/examples/02_nesting/01_nesting.rst b/docs/source/examples/02_nesting/01_nesting.rst new file mode 100644 index 000000000..d56858eeb --- /dev/null +++ b/docs/source/examples/02_nesting/01_nesting.rst @@ -0,0 +1,99 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Hierarchical Configs +========================================== + +Structures (typically dataclasses) can be nested to build hierarchical configuration +objects. This helps with modularity and grouping in larger projects. + + +.. code-block:: python + :linenos: + + + import dataclasses + import enum + import pathlib + + import tyro + + + class OptimizerType(enum.Enum): + ADAM = enum.auto() + SGD = enum.auto() + + + @dataclasses.dataclass + class OptimizerConfig: + # Gradient-based optimizer to use. + algorithm: OptimizerType = OptimizerType.ADAM + + # Learning rate to use. + learning_rate: float = 3e-4 + + # Coefficient for L2 regularization. + weight_decay: float = 1e-2 + + + @dataclasses.dataclass + class ExperimentConfig: + # Various configurable options for our optimizer. + optimizer: OptimizerConfig + + # Batch size. + batch_size: int = 32 + + # Total number of training steps. + train_steps: int = 100_000 + + # Random seed. This is helpful for making sure that our experiments are all + # reproducible! + seed: int = 0 + + + def train( + out_dir: pathlib.Path, + config: ExperimentConfig, + restore_checkpoint: bool = False, + checkpoint_interval: int = 1000, + ) -> None: + """Train a model. + + Args: + out_dir: Where to save logs and checkpoints. + config: Experiment configuration. + restore_checkpoint: Set to restore an existing checkpoint. + checkpoint_interval: Training steps between each checkpoint save. + """ + print(f"{out_dir=}, {restore_checkpoint=}, {checkpoint_interval=}") + print() + print(f"{config=}") + + + if __name__ == "__main__": + tyro.cli(train) + +------------ + +.. raw:: html + + python 02_nesting/01_nesting.py --help + +.. program-output:: python ../../examples/02_nesting/01_nesting.py --help + +------------ + +.. raw:: html + + python 02_nesting/01_nesting.py --out-dir . --config.optimizer.algorithm SGD + +.. program-output:: python ../../examples/02_nesting/01_nesting.py --out-dir . --config.optimizer.algorithm SGD + +------------ + +.. raw:: html + + python 02_nesting/01_nesting.py --out-dir . --restore-checkpoint + +.. program-output:: python ../../examples/02_nesting/01_nesting.py --out-dir . --restore-checkpoint diff --git a/docs/source/examples/02_nesting/02_subcommands.rst b/docs/source/examples/02_nesting/02_subcommands.rst new file mode 100644 index 000000000..1992a3bbc --- /dev/null +++ b/docs/source/examples/02_nesting/02_subcommands.rst @@ -0,0 +1,87 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Subcommands +========================================== + +Unions over nested types (classes or dataclasses) are populated using subcommands. + +For configuring subcommands beyond what can be expressed with type annotations, see +:func:`tyro.conf.subcommand()`. + + +.. code-block:: python + :linenos: + + + from __future__ import annotations + + import dataclasses + + import tyro + + + @dataclasses.dataclass(frozen=True) + class Checkout: + """Checkout a branch.""" + + branch: str + + + @dataclasses.dataclass(frozen=True) + class Commit: + """Commit changes.""" + + message: str + all: bool = False + + + def main(cmd: Checkout | Commit) -> None: + print(cmd) + + + if __name__ == "__main__": + # Note that we can also pass `Checkout | Command` directly into + # `tyro.cli()`; this is understood by tyro and pyright, but unfortunately not by + # mypy. + tyro.cli(main) + +------------ + +.. raw:: html + + python 02_nesting/02_subcommands.py --help + +.. program-output:: python ../../examples/02_nesting/02_subcommands.py --help + +------------ + +.. raw:: html + + python 02_nesting/02_subcommands.py cmd:commit --help + +.. program-output:: python ../../examples/02_nesting/02_subcommands.py cmd:commit --help + +------------ + +.. raw:: html + + python 02_nesting/02_subcommands.py cmd:commit --cmd.message hello --cmd.all + +.. program-output:: python ../../examples/02_nesting/02_subcommands.py cmd:commit --cmd.message hello --cmd.all + +------------ + +.. raw:: html + + python 02_nesting/02_subcommands.py cmd:checkout --help + +.. program-output:: python ../../examples/02_nesting/02_subcommands.py cmd:checkout --help + +------------ + +.. raw:: html + + python 02_nesting/02_subcommands.py cmd:checkout --cmd.branch main + +.. program-output:: python ../../examples/02_nesting/02_subcommands.py cmd:checkout --cmd.branch main diff --git a/docs/source/examples/02_nesting/03_multiple_subcommands.rst b/docs/source/examples/02_nesting/03_multiple_subcommands.rst new file mode 100644 index 000000000..5064916a8 --- /dev/null +++ b/docs/source/examples/02_nesting/03_multiple_subcommands.rst @@ -0,0 +1,103 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Sequenced Subcommands +========================================== + +Multiple unions over nested types are populated using a series of subcommands. + + +.. code-block:: python + :linenos: + + + from __future__ import annotations + + import dataclasses + from typing import Literal + + import tyro + + # Possible dataset configurations. + + + @dataclasses.dataclass + class Mnist: + binary: bool = False + """Set to load binary version of MNIST dataset.""" + + + @dataclasses.dataclass + class ImageNet: + subset: Literal[50, 100, 1000] + """Choose between ImageNet-50, ImageNet-100, ImageNet-1000, etc.""" + + + # Possible optimizer configurations. + + + @dataclasses.dataclass + class Adam: + learning_rate: float = 1e-3 + betas: tuple[float, float] = (0.9, 0.999) + + + @dataclasses.dataclass + class Sgd: + learning_rate: float = 3e-4 + + + # Train script. + + + def train( + dataset: Mnist | ImageNet = Mnist(), + optimizer: Adam | Sgd = Adam(), + ) -> None: + """Example training script. + + Args: + dataset: Dataset to train on. + optimizer: Optimizer to train with. + + Returns: + None: + """ + print(dataset) + print(optimizer) + + + if __name__ == "__main__": + tyro.cli(train, config=(tyro.conf.ConsolidateSubcommandArgs,)) + +------------ + +.. raw:: html + + python 02_nesting/03_multiple_subcommands.py --help + +.. program-output:: python ../../examples/02_nesting/03_multiple_subcommands.py --help + +------------ + +.. raw:: html + + python 02_nesting/03_multiple_subcommands.py dataset:mnist --help + +.. program-output:: python ../../examples/02_nesting/03_multiple_subcommands.py dataset:mnist --help + +------------ + +.. raw:: html + + python 02_nesting/03_multiple_subcommands.py dataset:mnist optimizer:adam --help + +.. program-output:: python ../../examples/02_nesting/03_multiple_subcommands.py dataset:mnist optimizer:adam --help + +------------ + +.. raw:: html + + python 02_nesting/03_multiple_subcommands.py dataset:mnist optimizer:adam --optimizer.learning-rate 3e-4 --dataset.binary + +.. program-output:: python ../../examples/02_nesting/03_multiple_subcommands.py dataset:mnist optimizer:adam --optimizer.learning-rate 3e-4 --dataset.binary diff --git a/docs/source/examples/02_nesting/04_nesting_in_containers.rst b/docs/source/examples/02_nesting/04_nesting_in_containers.rst new file mode 100644 index 000000000..d9d5a812f --- /dev/null +++ b/docs/source/examples/02_nesting/04_nesting_in_containers.rst @@ -0,0 +1,72 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Nesting in Containers +========================================== + +Structures can be nested inside of standard containers. + +Note that lengths must be inferable, either via a fixed-length tuple annotation or by +parsing default values. + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + + class Color: + pass + + + @dataclasses.dataclass + class RGB(Color): + r: int + g: int + b: int + + + @dataclasses.dataclass + class HSL(Color): + h: int + s: int + l: int + + + @dataclasses.dataclass + class Args: + # Example of specifying nested structures via a fixed-length tuple. + color_tuple: tuple[RGB, HSL] + + # Examples of nested structures in variable-length containers. These need a default + # provided for length inference; we don't currently support specifying dynamic + # container lengths directly from the commandline. + color_tuple_alt: tuple[Color, ...] = ( + RGB(255, 0, 0), + HSL(0, 255, 0), + ) + color_map: dict[str, RGB] = dataclasses.field( + # We can't use mutable values as defaults directly. + default_factory={ + "red": RGB(255, 0, 0), + "green": RGB(0, 255, 0), + "blue": RGB(0, 0, 255), + }.copy + ) + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 02_nesting/04_nesting_in_containers.py --help + +.. program-output:: python ../../examples/02_nesting/04_nesting_in_containers.py --help diff --git a/docs/source/examples/02_nesting/05_subcommands_func.rst b/docs/source/examples/02_nesting/05_subcommands_func.rst new file mode 100644 index 000000000..141c9a53d --- /dev/null +++ b/docs/source/examples/02_nesting/05_subcommands_func.rst @@ -0,0 +1,76 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Subcommands from Functions +========================================== + +We provide a shorthand for generating a subcommand CLI from a dictionary. This +is a thin wrapper around :func:`tyro.cli()`'s more verbose, type-based API. If +more generality is needed, the internal working are explained in the docs for +:func:`tyro.extras.subcommand_cli_from_dict()`. + + +.. code-block:: python + :linenos: + + + import tyro + + + def checkout(branch: str) -> None: + """Check out a branch.""" + print(f"{branch=}") + + + def commit(message: str, all: bool = False) -> None: + """Make a commit.""" + print(f"{message=} {all=}") + + + if __name__ == "__main__": + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + } + ) + +------------ + +.. raw:: html + + python 02_nesting/05_subcommands_func.py --help + +.. program-output:: python ../../examples/02_nesting/05_subcommands_func.py --help + +------------ + +.. raw:: html + + python 02_nesting/05_subcommands_func.py commit --help + +.. program-output:: python ../../examples/02_nesting/05_subcommands_func.py commit --help + +------------ + +.. raw:: html + + python 02_nesting/05_subcommands_func.py commit --message hello --all + +.. program-output:: python ../../examples/02_nesting/05_subcommands_func.py commit --message hello --all + +------------ + +.. raw:: html + + python 02_nesting/05_subcommands_func.py checkout --help + +.. program-output:: python ../../examples/02_nesting/05_subcommands_func.py checkout --help + +------------ + +.. raw:: html + + python 02_nesting/05_subcommands_func.py checkout --branch main + +.. program-output:: python ../../examples/02_nesting/05_subcommands_func.py checkout --branch main diff --git a/docs/source/examples/03_config_systems/01_base_configs.rst b/docs/source/examples/03_config_systems/01_base_configs.rst new file mode 100644 index 000000000..0ac47f25e --- /dev/null +++ b/docs/source/examples/03_config_systems/01_base_configs.rst @@ -0,0 +1,132 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Base Configurations +========================================== + +We can integrate `tyro` into common configuration patterns: here, we select +one of multiple possible base configurations, create a subcommand for each one, and then +use the CLI to either override (existing) or fill in (missing) values. + +The helper function used here, :func:`tyro.extras.overridable_config_cli()`, is a +lightweight wrapper over :func:`tyro.cli()`. + + +.. code-block:: python + :linenos: + + + from dataclasses import dataclass + from typing import Callable, Literal + + from torch import nn + + import tyro + + + @dataclass(frozen=True) + class AdamOptimizer: + learning_rate: float = 1e-3 + betas: tuple[float, float] = (0.9, 0.999) + + + @dataclass(frozen=True) + class ExperimentConfig: + # Dataset to run experiment on. + dataset: Literal["mnist", "imagenet-50"] + + # Optimizer parameters. + optimizer: AdamOptimizer + + # Model size. + num_layers: int + units: int + + # Batch size. + batch_size: int + + # Total number of training steps. + train_steps: int + + # Random seed. This is helpful for making sure that our experiments are all + # reproducible! + seed: int + + # Activation to use. Not specifiable via the commandline. + activation: Callable[[], nn.Module] + + + # Note that we could also define this library using separate YAML files (similar to + # `config_path`/`config_name` in Hydra), but staying in Python enables seamless type + # checking + IDE support. + default_configs = { + "small": ( + "Small experiment.", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + ), + "big": ( + "Big experiment.", + ExperimentConfig( + dataset="imagenet-50", + optimizer=AdamOptimizer(), + batch_size=32, + num_layers=8, + units=256, + train_steps=100_000, + seed=0, + activation=nn.GELU, + ), + ), + } + if __name__ == "__main__": + config = tyro.extras.overridable_config_cli(default_configs) + print(config) + +------------ + +.. raw:: html + + python 03_config_systems/01_base_configs.py --help + +.. program-output:: python ../../examples/03_config_systems/01_base_configs.py --help + +------------ + +.. raw:: html + + python 03_config_systems/01_base_configs.py small --help + +.. program-output:: python ../../examples/03_config_systems/01_base_configs.py small --help + +------------ + +.. raw:: html + + python 03_config_systems/01_base_configs.py small --seed 94720 + +.. program-output:: python ../../examples/03_config_systems/01_base_configs.py small --seed 94720 + +------------ + +.. raw:: html + + python 03_config_systems/01_base_configs.py big --help + +.. program-output:: python ../../examples/03_config_systems/01_base_configs.py big --help + +------------ + +.. raw:: html + + python 03_config_systems/01_base_configs.py big --seed 94720 + +.. program-output:: python ../../examples/03_config_systems/01_base_configs.py big --seed 94720 diff --git a/docs/source/examples/03_config_systems/02_overriding_yaml.rst b/docs/source/examples/03_config_systems/02_overriding_yaml.rst new file mode 100644 index 000000000..e926429cc --- /dev/null +++ b/docs/source/examples/03_config_systems/02_overriding_yaml.rst @@ -0,0 +1,61 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Overriding YAML Configs +========================================== + +If you have a library of existing YAML files that you want to use, `tyro` can +be used to override values in them. We generally recommend dataclass configs +for new projects. + + +.. code-block:: python + :linenos: + + + import yaml + + import tyro + + # YAML configuration. Note that this could also be loaded from a file! Environment + # variables are an easy way to select between different YAML files. + default_yaml = r""" + exp_name: test + optimizer: + learning_rate: 0.0001 + type: adam + training: + batch_size: 32 + num_steps: 10000 + checkpoint_steps: + - 500 + - 1000 + - 1500 + """.strip() + + if __name__ == "__main__": + # Convert our YAML config into a nested dictionary. + default_config = yaml.safe_load(default_yaml) + + # Override fields in the dictionary. + overridden_config = tyro.cli(dict, default=default_config) + + # Print the overridden config. + overridden_yaml = yaml.safe_dump(overridden_config) + print(overridden_yaml) + +------------ + +.. raw:: html + + python 03_config_systems/02_overriding_yaml.py --help + +.. program-output:: python ../../examples/03_config_systems/02_overriding_yaml.py --help + +------------ + +.. raw:: html + + python 03_config_systems/02_overriding_yaml.py --training.checkpoint-steps 300 1000 9000 + +.. program-output:: python ../../examples/03_config_systems/02_overriding_yaml.py --training.checkpoint-steps 300 1000 9000 diff --git a/docs/source/examples/04_additional/01_positional_args.rst b/docs/source/examples/04_additional/01_positional_args.rst new file mode 100644 index 000000000..b547f0218 --- /dev/null +++ b/docs/source/examples/04_additional/01_positional_args.rst @@ -0,0 +1,83 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Positional Arguments +========================================== + +Positional-only arguments in functions are converted to positional CLI arguments. + +For more general positional arguments, see :class:`tyro.conf.Positional`. + + +.. code-block:: python + :linenos: + + + from __future__ import annotations + + import dataclasses + import enum + import pathlib + from typing import Tuple + + import tyro + + + def main( + source: pathlib.Path, + dest: pathlib.Path, + /, # Mark the end of positional arguments. + optimizer: OptimizerConfig, + force: bool = False, + verbose: bool = False, + background_rgb: Tuple[float, float, float] = (1.0, 0.0, 0.0), + ) -> None: + """Command-line interface defined using a function signature. Note that this + docstring is parsed to generate helptext. + + Args: + source: Source path. + dest: Destination path. + optimizer: Configuration for our optimizer object. + force: Do not prompt before overwriting. + verbose: Explain what is being done. + background_rgb: Background color. Red by default. + """ + print(f"{source=}\n{dest=}\n{optimizer=}\n{force=}\n{verbose=}\n{background_rgb=}") + + + class OptimizerType(enum.Enum): + ADAM = enum.auto() + SGD = enum.auto() + + + @dataclasses.dataclass(frozen=True) + class OptimizerConfig: + algorithm: OptimizerType = OptimizerType.ADAM + """Gradient-based optimizer to use.""" + + learning_rate: float = 3e-4 + """Learning rate to use.""" + + weight_decay: float = 1e-2 + """Coefficient for L2 regularization.""" + + + if __name__ == "__main__": + tyro.cli(main) + +------------ + +.. raw:: html + + python 04_additional/01_positional_args.py --help + +.. program-output:: python ../../examples/04_additional/01_positional_args.py --help + +------------ + +.. raw:: html + + python 04_additional/01_positional_args.py ./a ./b --optimizer.learning-rate 1e-5 + +.. program-output:: python ../../examples/04_additional/01_positional_args.py ./a ./b --optimizer.learning-rate 1e-5 diff --git a/docs/source/examples/04_additional/02_dictionaries.rst b/docs/source/examples/04_additional/02_dictionaries.rst new file mode 100644 index 000000000..7d28077e6 --- /dev/null +++ b/docs/source/examples/04_additional/02_dictionaries.rst @@ -0,0 +1,82 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Dictionaries and TypedDict +========================================== + +Dictionary inputs can be specified using either a standard `Dict[K, V]` +annotation, or a :code:`TypedDict` subclass. + +For configuring :code:`TypedDict`, we also support :code:`total={True/False}`, +:code:`typing.Required`, and :code:`typing.NotRequired`. See the `Python docs `_ for all :code:`TypedDict` features. + + +.. code-block:: python + :linenos: + + + from typing import TypedDict + + from typing_extensions import NotRequired + + import tyro + + + class DictionarySchemaA( + TypedDict, + # Setting `total=False` specifies that not all keys need to exist. + total=False, + ): + learning_rate: float + betas: tuple[float, float] + + + class DictionarySchemaB(TypedDict): + learning_rate: NotRequired[float] + """NotRequired[] specifies that a particular key doesn't need to exist.""" + betas: tuple[float, float] + + + def main( + typed_dict_a: DictionarySchemaA, + typed_dict_b: DictionarySchemaB, + standard_dict: dict[str, float] = { + "learning_rate": 3e-4, + "beta1": 0.9, + "beta2": 0.999, + }, + ) -> None: + assert isinstance(typed_dict_a, dict) + assert isinstance(typed_dict_b, dict) + assert isinstance(standard_dict, dict) + print("Typed dict A:", typed_dict_a) + print("Typed dict B:", typed_dict_b) + print("Standard dict:", standard_dict) + + + if __name__ == "__main__": + tyro.cli(main) + +------------ + +.. raw:: html + + python 04_additional/02_dictionaries.py --help + +.. program-output:: python ../../examples/04_additional/02_dictionaries.py --help + +------------ + +.. raw:: html + + python 04_additional/02_dictionaries.py --typed-dict-a.learning-rate 3e-4 --typed-dict-b.betas 0.9 0.999 + +.. program-output:: python ../../examples/04_additional/02_dictionaries.py --typed-dict-a.learning-rate 3e-4 --typed-dict-b.betas 0.9 0.999 + +------------ + +.. raw:: html + + python 04_additional/02_dictionaries.py --typed-dict-b.betas 0.9 0.999 + +.. program-output:: python ../../examples/04_additional/02_dictionaries.py --typed-dict-b.betas 0.9 0.999 diff --git a/docs/source/examples/04_additional/03_tuples.rst b/docs/source/examples/04_additional/03_tuples.rst new file mode 100644 index 000000000..cb979eda9 --- /dev/null +++ b/docs/source/examples/04_additional/03_tuples.rst @@ -0,0 +1,65 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Tuples +========================================== + +Example using :func:`tyro.cli()` to instantiate tuple types. :code:`tuple`, +:code:`typing.Tuple`, and :code:`NamedTuple` are all supported. + + +.. code-block:: python + :linenos: + + + from typing import NamedTuple + + import tyro + + + # Named tuples are interpreted as nested structures. + class Color(NamedTuple): + r: int + g: int + b: int + + + class TupleType(NamedTuple): + """Description. + This should show up in the helptext!""" + + # Tuple types can contain raw values. + color: tuple[int, int, int] = (255, 0, 0) + + # Tuple types can contain nested structures. + two_colors: tuple[Color, Color] = (Color(255, 0, 0), Color(0, 255, 0)) + + + if __name__ == "__main__": + x = tyro.cli(TupleType) + assert isinstance(x, tuple) + print(x) + +------------ + +.. raw:: html + + python 04_additional/03_tuples.py --help + +.. program-output:: python ../../examples/04_additional/03_tuples.py --help + +------------ + +.. raw:: html + + python 04_additional/03_tuples.py --color 127 127 127 + +.. program-output:: python ../../examples/04_additional/03_tuples.py --color 127 127 127 + +------------ + +.. raw:: html + + python 04_additional/03_tuples.py --two-colors.1.r 127 --two-colors.1.g 0 --two-colors.1.b 0 + +.. program-output:: python ../../examples/04_additional/03_tuples.py --two-colors.1.r 127 --two-colors.1.g 0 --two-colors.1.b 0 diff --git a/docs/source/examples/04_additional/04_classes.rst b/docs/source/examples/04_additional/04_classes.rst new file mode 100644 index 000000000..bf6f64526 --- /dev/null +++ b/docs/source/examples/04_additional/04_classes.rst @@ -0,0 +1,53 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Instantiating Classes +========================================== + +In addition to functions and dataclasses, we can also generate CLIs from the +constructors of standard Python classes. + + +.. code-block:: python + :linenos: + + + import tyro + + + class Args: + def __init__( + self, + field1: str, + field2: int, + flag: bool = False, + ): + """Arguments. + + Args: + field1: A string field. + field2: A numeric field. + flag: A boolean flag. + """ + self.data = [field1, field2, flag] + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args.data) + +------------ + +.. raw:: html + + python 04_additional/04_classes.py --help + +.. program-output:: python ../../examples/04_additional/04_classes.py --help + +------------ + +.. raw:: html + + python 04_additional/04_classes.py --field1 hello --field2 7 + +.. program-output:: python ../../examples/04_additional/04_classes.py --field1 hello --field2 7 diff --git a/docs/source/examples/04_additional/05_generics.rst b/docs/source/examples/04_additional/05_generics.rst new file mode 100644 index 000000000..6578dcd93 --- /dev/null +++ b/docs/source/examples/04_additional/05_generics.rst @@ -0,0 +1,53 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Generic Types +========================================== + +Example of parsing for generic dataclasses. + + +.. code-block:: python + :linenos: + + + import dataclasses + from typing import Generic, TypeVar + + import tyro + + ScalarType = TypeVar("ScalarType", int, float) + ShapeType = TypeVar("ShapeType") + + + @dataclasses.dataclass(frozen=True) + class Point3(Generic[ScalarType]): + x: ScalarType + y: ScalarType + z: ScalarType + frame_id: str + + + @dataclasses.dataclass(frozen=True) + class Triangle: + a: Point3[float] + b: Point3[float] + c: Point3[float] + + + @dataclasses.dataclass(frozen=True) + class Args(Generic[ShapeType]): + shape: ShapeType + + + if __name__ == "__main__": + args = tyro.cli(Args[Triangle]) + print(args) + +------------ + +.. raw:: html + + python 04_additional/05_generics.py --help + +.. program-output:: python ../../examples/04_additional/05_generics.py --help diff --git a/docs/source/examples/04_additional/06_generics_py312.rst b/docs/source/examples/04_additional/06_generics_py312.rst new file mode 100644 index 000000000..44dbb9d8e --- /dev/null +++ b/docs/source/examples/04_additional/06_generics_py312.rst @@ -0,0 +1,53 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Generic Types (Python 3.12+ syntax) +========================================== + +Example of parsing for generic dataclasses using syntax introduced in Python +3.12 (`PEP 695 `_). + +.. warning:: + If used in conjunction with :code:`from __future__ import annotations`, the updated type parameter syntax requires Python 3.12.4 or newer. For technical details, see `this CPython PR `_. + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + + @dataclasses.dataclass(frozen=True) + class Point3[ScalarType: (int, float)]: + x: ScalarType + y: ScalarType + z: ScalarType + frame_id: str + + + @dataclasses.dataclass(frozen=True) + class Triangle: + a: Point3[float] + b: Point3[float] + c: Point3[float] + + + @dataclasses.dataclass(frozen=True) + class Args[ShapeType]: + shape: ShapeType + + + if __name__ == "__main__": + args = tyro.cli(Args[Triangle]) + print(args) + +------------ + +.. raw:: html + + python 04_additional/06_generics_py312.py --help + +.. program-output:: python ../../examples/04_additional/06_generics_py312.py --help diff --git a/docs/source/examples/04_additional/07_conf.rst b/docs/source/examples/04_additional/07_conf.rst new file mode 100644 index 000000000..2faacdf48 --- /dev/null +++ b/docs/source/examples/04_additional/07_conf.rst @@ -0,0 +1,63 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Configuration via typing.Annotated[] +========================================== + +The :mod:`tyro.conf` module contains utilities that can be used to configure +command-line interfaces beyond what is expressible via static type annotations. + +Features here are supported, but generally unnecessary and should be used sparingly. + + +.. code-block:: python + :linenos: + + + import dataclasses + + from typing_extensions import Annotated + + import tyro + + + @dataclasses.dataclass + class Args: + # A numeric field parsed as a positional argument. + positional: tyro.conf.Positional[int] + + # A boolean field with flag conversion turned off. + boolean: tyro.conf.FlagConversionOff[bool] = False + + # A numeric field that can't be changed via the CLI. + fixed: tyro.conf.Fixed[int] = 5 + + # A field with manually overridden properties. + manual: Annotated[ + str, + tyro.conf.arg( + name="renamed", + metavar="STRING", + help="A field with manually overridden properties!", + ), + ] = "Hello" + + + if __name__ == "__main__": + print(tyro.cli(Args)) + +------------ + +.. raw:: html + + python 04_additional/07_conf.py --help + +.. program-output:: python ../../examples/04_additional/07_conf.py --help + +------------ + +.. raw:: html + + python 04_additional/07_conf.py 5 --boolean True + +.. program-output:: python ../../examples/04_additional/07_conf.py 5 --boolean True diff --git a/docs/source/examples/04_additional/08_pydantic.rst b/docs/source/examples/04_additional/08_pydantic.rst new file mode 100644 index 000000000..9f529d173 --- /dev/null +++ b/docs/source/examples/04_additional/08_pydantic.rst @@ -0,0 +1,54 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Pydantic Integration +========================================== + +In addition to standard dataclasses, :func:`tyro.cli()` also supports +`Pydantic `_ models. + + +.. code-block:: python + :linenos: + + + from pydantic import BaseModel, Field + + import tyro + + + class Args(BaseModel): + """Description. + This should show up in the helptext!""" + + field1: str + field2: int = Field(3, description="An integer field.") + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 04_additional/08_pydantic.py --help + +.. program-output:: python ../../examples/04_additional/08_pydantic.py --help + +------------ + +.. raw:: html + + python 04_additional/08_pydantic.py --field1 hello + +.. program-output:: python ../../examples/04_additional/08_pydantic.py --field1 hello + +------------ + +.. raw:: html + + python 04_additional/08_pydantic.py --field1 hello --field2 5 + +.. program-output:: python ../../examples/04_additional/08_pydantic.py --field1 hello --field2 5 diff --git a/docs/source/examples/04_additional/09_attrs.rst b/docs/source/examples/04_additional/09_attrs.rst new file mode 100644 index 000000000..e52a189f5 --- /dev/null +++ b/docs/source/examples/04_additional/09_attrs.rst @@ -0,0 +1,58 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Attrs Integration +========================================== + +In addition to standard dataclasses, :func:`tyro.cli()` also supports +`attrs `_ classes. + + +.. code-block:: python + :linenos: + + + import attr + + import tyro + + + @attr.s + class Args: + """Description. + This should show up in the helptext!""" + + field1: str = attr.ib() + """A string field.""" + + field2: int = attr.ib(factory=lambda: 5) + """A required integer field.""" + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 04_additional/09_attrs.py --help + +.. program-output:: python ../../examples/04_additional/09_attrs.py --help + +------------ + +.. raw:: html + + python 04_additional/09_attrs.py --field1 hello + +.. program-output:: python ../../examples/04_additional/09_attrs.py --field1 hello + +------------ + +.. raw:: html + + python 04_additional/09_attrs.py --field1 hello --field2 5 + +.. program-output:: python ../../examples/04_additional/09_attrs.py --field1 hello --field2 5 diff --git a/docs/source/examples/04_additional/10_flax.rst b/docs/source/examples/04_additional/10_flax.rst new file mode 100644 index 000000000..aec226d00 --- /dev/null +++ b/docs/source/examples/04_additional/10_flax.rst @@ -0,0 +1,74 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +JAX/Flax Integration +========================================== + +If you use `flax.linen `_, modules can be instantiated +directly from :func:`tyro.cli()`. + + +.. code-block:: python + :linenos: + + + from flax import linen as nn + from jax import numpy as jnp + + import tyro + + + class Classifier(nn.Module): + layers: int + """Layers in our network.""" + units: int = 32 + """Hidden unit count.""" + output_dim: int = 10 + """Number of classes.""" + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # type: ignore + for i in range(self.layers - 1): + x = nn.Dense( + self.units, + kernel_init=nn.initializers.kaiming_normal(), + )(x) + x = nn.relu(x) + + x = nn.Dense( + self.output_dim, + kernel_init=nn.initializers.xavier_normal(), + )(x) + x = nn.sigmoid(x) + return x + + + def train(model: Classifier, num_iterations: int = 1000) -> None: + """Train a model. + + Args: + model: Model to train. + num_iterations: Number of training iterations. + """ + print(f"{model=}") + print(f"{num_iterations=}") + + + if __name__ == "__main__": + tyro.cli(train) + +------------ + +.. raw:: html + + python 04_additional/10_flax.py --help + +.. program-output:: python ../../examples/04_additional/10_flax.py --help + +------------ + +.. raw:: html + + python 04_additional/10_flax.py --model.layers 4 + +.. program-output:: python ../../examples/04_additional/10_flax.py --model.layers 4 diff --git a/docs/source/examples/04_additional/11_custom_constructors.rst b/docs/source/examples/04_additional/11_custom_constructors.rst new file mode 100644 index 000000000..3d91faa9a --- /dev/null +++ b/docs/source/examples/04_additional/11_custom_constructors.rst @@ -0,0 +1,95 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Custom Constructors +========================================== + +For additional flexibility, :func:`tyro.conf.arg()` accepts a +:code:`constructor` argument, which makes it easier to load complex objects. + +.. warning:: + Custom constructors are permitted wherever :func:`tyro.conf.arg()` is. This + annotation corresponds to a single argument, so we expect that it is placed + placed at the root of the field corresponding to that argument. For + example: + + .. code-block:: python + + # Great! + x: Annotated[int, tyro.conf.arg(custom_constructor=...)] + + However, nesting :func:`tyro.conf.arg()` within other types is generally + not supported. For example: + + .. code-block:: python + + # Not supported. + x: tuple[Annotated[int, tyro.conf.arg(custom_constructor=...)], ...] + + This example can be rewritten with a custom constructor that directly + returns a tuple of integers: + + .. code-block:: python + + # Workaround for above. + x: Annotated[tuple[int, ...], tyro.conf.arg(custom_constructor=...)] + + +.. code-block:: python + :linenos: + + + import json as json_ + + from typing_extensions import Annotated + + import tyro + + + def dict_json_constructor(json: str) -> dict: + """Construct a dictionary from a JSON string. Raises a ValueError if the result is + not a dictionary.""" + out = json_.loads(json) + if not isinstance(out, dict): + raise ValueError(f"{json} is not a dictionary!") + return out + + + # A dictionary type, but `tyro` will expect a JSON string from the CLI. + JsonDict = Annotated[dict, tyro.conf.arg(constructor=dict_json_constructor)] + + + def main( + dict1: JsonDict, + dict2: JsonDict = {"default": None}, + ) -> None: + print(f"{dict1=}") + print(f"{dict2=}") + + + if __name__ == "__main__": + tyro.cli(main) + +------------ + +.. raw:: html + + python 04_additional/11_custom_constructors.py --help + +.. program-output:: python ../../examples/04_additional/11_custom_constructors.py --help + +------------ + +.. raw:: html + + python 04_additional/11_custom_constructors.py --dict1.json '{"hello": "world"}' + +.. program-output:: python ../../examples/04_additional/11_custom_constructors.py --dict1.json '{"hello": "world"}' + +------------ + +.. raw:: html + + python 04_additional/11_custom_constructors.py --dict1.json '{"hello": "world"}`' --dict2.json '{"hello": "world"}' + +.. program-output:: python ../../examples/04_additional/11_custom_constructors.py --dict1.json '{"hello": "world"}`' --dict2.json '{"hello": "world"}' diff --git a/docs/source/examples/04_additional/12_aliases.rst b/docs/source/examples/04_additional/12_aliases.rst new file mode 100644 index 000000000..1a8c5aa32 --- /dev/null +++ b/docs/source/examples/04_additional/12_aliases.rst @@ -0,0 +1,96 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Argument Aliases +========================================== + +:func:`tyro.conf.arg()` can be used to attach aliases to arguments. + + +.. code-block:: python + :linenos: + + + from typing_extensions import Annotated + + import tyro + + + def checkout( + branch: Annotated[str, tyro.conf.arg(aliases=["-b"])], + ) -> None: + """Check out a branch.""" + print(f"{branch=}") + + + def commit( + message: Annotated[str, tyro.conf.arg(aliases=["-m"])], + all: Annotated[bool, tyro.conf.arg(aliases=["-a"])] = False, + ) -> None: + """Make a commit.""" + print(f"{message=} {all=}") + + + if __name__ == "__main__": + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + } + ) + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py --help + +.. program-output:: python ../../examples/04_additional/12_aliases.py --help + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py commit --help + +.. program-output:: python ../../examples/04_additional/12_aliases.py commit --help + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py commit --message hello --all + +.. program-output:: python ../../examples/04_additional/12_aliases.py commit --message hello --all + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py commit -m hello -a + +.. program-output:: python ../../examples/04_additional/12_aliases.py commit -m hello -a + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py checkout --help + +.. program-output:: python ../../examples/04_additional/12_aliases.py checkout --help + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py checkout --branch main + +.. program-output:: python ../../examples/04_additional/12_aliases.py checkout --branch main + +------------ + +.. raw:: html + + python 04_additional/12_aliases.py checkout -b main + +.. program-output:: python ../../examples/04_additional/12_aliases.py checkout -b main diff --git a/docs/source/examples/04_additional/12_counters.rst b/docs/source/examples/04_additional/12_counters.rst new file mode 100644 index 000000000..a05decc34 --- /dev/null +++ b/docs/source/examples/04_additional/12_counters.rst @@ -0,0 +1,67 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Counters +========================================== + +Repeatable 'counter' arguments can be specified via :data:`tyro.conf.UseCounterAction`. + + +.. code-block:: python + :linenos: + + + from typing_extensions import Annotated + + import tyro + from tyro.conf import UseCounterAction + + + def main( + verbosity: UseCounterAction[int], + aliased_verbosity: Annotated[UseCounterAction[int], tyro.conf.arg(aliases=["-v"])], + ) -> None: + """Example showing how to use counter actions. + + Args: + verbosity: Verbosity level. + aliased_verbosity: Same as above, but can also be specified with -v, -vv, -vvv, etc. + """ + print("Verbosity level:", verbosity) + print("Verbosity level (aliased):", aliased_verbosity) + + + if __name__ == "__main__": + tyro.cli(main) + +------------ + +.. raw:: html + + python 04_additional/12_counters.py --help + +.. program-output:: python ../../examples/04_additional/12_counters.py --help + +------------ + +.. raw:: html + + python 04_additional/12_counters.py --verbosity + +.. program-output:: python ../../examples/04_additional/12_counters.py --verbosity + +------------ + +.. raw:: html + + python 04_additional/12_counters.py --verbosity --verbosity + +.. program-output:: python ../../examples/04_additional/12_counters.py --verbosity --verbosity + +------------ + +.. raw:: html + + python 04_additional/12_counters.py -vvv + +.. program-output:: python ../../examples/04_additional/12_counters.py -vvv diff --git a/docs/source/examples/04_additional/13_type_statement.rst b/docs/source/examples/04_additional/13_type_statement.rst new file mode 100644 index 000000000..e568a1102 --- /dev/null +++ b/docs/source/examples/04_additional/13_type_statement.rst @@ -0,0 +1,50 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Type Aliases (Python 3.12+) +========================================== + +In Python 3.12, the :code:`type` statement is introduced to create type aliases. + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + # Lazily-evaluated type alias. + type Field1Type = Inner + + + @dataclasses.dataclass + class Inner: + a: int + b: str + + + @dataclasses.dataclass + class Args: + """Description. + This should show up in the helptext!""" + + field1: Field1Type + """A field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + + if __name__ == "__main__": + args = tyro.cli(Args) + print(args) + +------------ + +.. raw:: html + + python 04_additional/13_type_statement.py --help + +.. program-output:: python ../../examples/04_additional/13_type_statement.py --help diff --git a/docs/source/examples/04_additional/14_suppress_console_outputs.rst b/docs/source/examples/04_additional/14_suppress_console_outputs.rst new file mode 100644 index 000000000..e6056ea69 --- /dev/null +++ b/docs/source/examples/04_additional/14_suppress_console_outputs.rst @@ -0,0 +1,57 @@ +.. Comment: this file is automatically generated by `update_example_docs.py`. + It should not be modified manually. + +Cleaner Console Outputs for Scripts with Multiple Workers +========================================== + +The :code:`console_outputs=` argument can be set to :code:`False` to suppress helptext and +error message printing. + +This is useful in PyTorch for distributed training scripts, where you only want +to print the helptext from the main process: + + +.. code-block:: python + + # HuggingFace Accelerate. + args = tyro.cli(Args, console_outputs=accelerator.is_main_process) + + # PyTorch DDP. + args = tyro.cli(Args, console_outputs=(rank == 0)) + + # PyTorch Lightning. + args = tyro.cli(Args, console_outputs=trainer.is_global_zero) + + +.. code-block:: python + :linenos: + + + import dataclasses + + import tyro + + + @dataclasses.dataclass + class Args: + """Description. + This should show up in the helptext!""" + + field1: int + """A field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + + if __name__ == "__main__": + args = tyro.cli(Args, console_outputs=False) + print(args) + +------------ + +.. raw:: html + + python 04_additional/14_suppress_console_outputs.py --help + +.. program-output:: python ../../examples/04_additional/14_suppress_console_outputs.py --help diff --git a/docs/source/goals_and_alternatives.md b/docs/source/goals_and_alternatives.md new file mode 100644 index 000000000..e5f36e0ff --- /dev/null +++ b/docs/source/goals_and_alternatives.md @@ -0,0 +1,100 @@ +# Goals and alternatives + +## Design goals + +The core functionality of `tyro` — generating argument parsers from type +annotations — overlaps significantly with features offered by other libraries. +Usage distinctions are the result of two API goals: + +- **One uninvasive function.** For all core functionality, learning to use + `tyro` should reduce to learning to write type-annotated Python. For example, + types are specified using standard annotations, helptext using docstrings, + choices using the standard `typing.Literal` type, subcommands with + `typing.Union` of nested types, and positional arguments with `/`. + - In contrast, similar libraries have more expansive APIs , and require more + library-specific structures, decorators, or metadata formats for configuring + parsing behavior. +- **Strict typing.** Any type that can be annotated and unambiguously parsed + with an `argparse`-style CLI interface should work out-of-the-box; any public + API that isn't statically analyzable should be avoided. + - In contrast, many similar libraries implement features that depend on + dynamic argparse-style namespaces, or string-based accessors that can't be + statically checked. + +More concretely, we can also compare specific features. A noncomprehensive set: + +| | Dataclasses | Functions | Literals | Docstrings as helptext | Nested structures | Unions over primitives | Unions over nested types | Lists, tuples | Dictionaries | Generics | +| -------------------------------------------- | ----------- | --------- | -------------------- | ---------------------- | ----------------- | ---------------------- | ------------------------- | -------------------- | ------------ | -------- | +| [argparse-dataclass][argparse-dataclass] | ✓ | | | | | | | | | | +| [argparse-dataclasses][argparse-dataclasses] | ✓ | | | | | | | | | | +| [datargs][datargs] | ✓ | | ✓[^datargs_literals] | | | | ✓[^datargs_unions_nested] | ✓ | | | +| [tap][tap] | | | ✓ | ✓ | | ✓ | ~[^tap_unions_nested] | ✓ | | | +| [simple-parsing][simple-parsing] | ✓ | | ✓[^simp_literals] | ✓ | ✓ | ✓ | ✓[^simp_unions_nested] | ✓ | ✓ | | +| [dataclass-cli][dataclass-cli] | ✓ | | | | | | | | | | +| [clout][clout] | ✓ | | | | ✓ | | | | | | +| [hf_argparser][hf_argparser] | ✓ | | | | | | | ✓ | ✓ | | +| [typer][typer] | | ✓ | | | | | ~[^typer_unions_nested] | ~[^typer_containers] | | | +| [pyrallis][pyrallis] | ✓ | | | ✓ | ✓ | | | ✓ | | | +| [yahp][yahp] | ✓ | | | ~[^yahp_docstrings] | ✓ | ✓ | ~[^yahp_unions_nested] | ✓ | | | +| [omegaconf][omegaconf] | ✓ | | | | ✓ | | | ✓ | ✓ | | +| **tyro** | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | + + + +[datargs]: https://github.com/roee30/datargs +[tap]: https://github.com/swansonk14/typed-argument-parser +[simple-parsing]: https://github.com/lebrice/SimpleParsing +[argparse-dataclass]: https://pypi.org/project/argparse-dataclass/ +[argparse-dataclasses]: https://pypi.org/project/argparse-dataclasses/ +[dataclass-cli]: https://github.com/malte-soe/dataclass-cli +[clout]: https://pypi.org/project/clout/ +[hf_argparser]: https://github.com/huggingface/transformers/blob/master/src/transformers/hf_argparser.py +[pyrallis]: https://github.com/eladrich/pyrallis/ +[typer]: https://typer.tiangolo.com/ +[yahp]: https://github.com/mosaicml/yahp +[omegaconf]: https://omegaconf.readthedocs.io/en/2.1_branch/structured_config.html + +[^datargs_unions_nested]: One allowed per class. +[^tap_unions_nested]: Not supported, but API exists for creating subcommands that accomplish a similar goal. +[^simp_unions_nested]: One allowed per class. +[^yahp_unions_nested]: Not supported, but similar functionality available via ["registries"](https://docs.mosaicml.com/projects/yahp/en/stable/examples/registry.html). +[^typer_unions_nested]: Not supported, but API exists for creating subcommands that accomplish a similar goal. +[^simp_literals]: Not supported for mixed (eg `Literal[5, "five"]`) or in container (eg `List[Literal[1, 2]]`) types. +[^datargs_literals]: Not supported for mixed types (eg `Literal[5, "five"]`). +[^typer_containers]: `typer` uses positional arguments for all required fields, which means that only one variable-length argument (such as `List[int]`) without a default is supported per argument parser. +[^yahp_docstrings]: Via the `hp.auto()` function, which can parse docstrings from external classes. Usage is different from the more direct parsing that `tyro`, `tap`, and `simple-parsing`/`pyrallis` support. + + + +Note that other libraries are generally aimed specifically at only one of +dataclasses (`datargs`, `simple-parsing`, `argparse-dataclass`, +`argparse-dataclasses`, `dataclass-cli`, `clout`, `hf_argparser`, `pyrallis`, +`yahp`), custom structures (`tap`), or functions (`typer`) rather than general +types and callables, but offer other features that you might find critical, such +as registration for custom types (`pyrallis`), built-in approaches for +serialization and config files (`tap`, `pyrallis`, `yahp`), and opportunities +for integration with fields generated using standard argparse definitions +(`simple-parsing`). + +## Note on configuration systems + +Our API is designed to be used to configure general applications and +computational experiments written in Python, but intentionally tries to avoid +building a full configuration framework (for example, `hydra`). These frameworks +can typically be broken into three components with varying levels of +integration, which include syntax and logic for **(1)** defining configurable +fields, **(2)** saving and choosing between a set of base configurations, and +**(3)** overriding configurable values at the command-line. + +`tyro` is meant to take advantage of modern Python features for **(1)** and +focus completely on **(3)** in a way that's agnostic to a project's preferred +approach for **(2)**. **(2)** is left as an exercise to the user; it tends to be +the most open-ended and varied in terms of project and personal preferences, but +is typically straightforward to implement given **(1)**. + +In contrast, popular libraries oriented toward configuration often strive to be +more "batteries-included", which is convenient but requires prescribing things +like specific formats, processes, or directory structures for working with saved +configurations. This requires more moving parts, which can be limiting if any +one of them is insufficient for a particular use case. (or if you just don't +like working with YAML formats) diff --git a/docs/source/helptext_generation.md b/docs/source/helptext_generation.md new file mode 100644 index 000000000..8b46e5882 --- /dev/null +++ b/docs/source/helptext_generation.md @@ -0,0 +1,76 @@ +# Helptext generation + +In addition to type annotations, :func:`tyro.cli()` will also parse docstrings +and comments. These are used to automatically generate helptext; see examples +for how these end up being formatted. + +## General callables + +For general callables, field helptext is extracted from the corresponding field +docstring. Our examples use Google-style docstrings, but ReST, Numpydoc-style +and Epydoc docstrings are supported as well. Under the hood, all of these +options use [docstring_parser](https://github.com/rr-/docstring_parser). + +```python +def main( + field1: str, + field2: int = 3, +) -> None: + """Function, whose arguments will be populated from a CLI interface. + + Args: + field1: A string field. + field2: A numeric field, with a default value. + """ + print(field1, field2) +``` + +## Dataclasses, TypedDict, NamedTuple + +For types defined using class attributes, enumerating each argument list in the +class docstring can be cumbersome. + +If they are unavailable, :func:`tyro.cli` will generate helptext from +docstrings and comments on attributes. These are parsed via source code +inspection. + +**(1) Attribute docstrings** + +As per [PEP 257](https://peps.python.org/pep-0257/#what-is-a-docstring). + +```python +@dataclasses.dataclass +class Args: + field1: str + """A string field.""" + field2: int = 3 + """A numeric field, with a default value.""" +``` + +**(2) Inline comments** + +Inline comments can be more succinct than true attribute docstrings. + +```python +@dataclasses.dataclass +class Args: + field1: str # A string field. + field2: int = 3 # A numeric field, with a default value. +``` + +**(3) Preceding comments** + +These can also be handy for semantically grouping fields, such as the two string +fields below. + +```python +@dataclasses.dataclass +class Args: + # String fields. + field1: str + field2: str + + # An integer field. + # Multi-line comments are supported. + field3: int +``` diff --git a/docs/source/index.md b/docs/source/index.md new file mode 100644 index 000000000..05d7e25db --- /dev/null +++ b/docs/source/index.md @@ -0,0 +1,159 @@ +# tyro + +|build| |nbsp| |ruff| |nbsp| |mypy| |nbsp| |pyright| |nbsp| |coverage| |nbsp| |versions| + +:code:`tyro` is a tool for generating command-line interfaces and configuration +objects in Python. + +Our core API, :func:`tyro.cli()`, + +- **Generates CLI interfaces** from a comprehensive set of Python type + constructs. +- **Populates helptext automatically** from defaults, annotations, and + docstrings. +- **Understands nesting** of `dataclasses`, `pydantic`, and `attrs` structures. +- **Prioritizes static analysis** for type checking and autocompletion with + tools like + [Pylance](https://marketplace.visualstudio.com/items?itemName=ms-python.vscode-pylance), + [Pyright](https://github.com/microsoft/pyright), and + [mypy](https://github.com/python/mypy). + +For advanced users, it also supports: + +- **Subcommands**, as well as choosing between and overriding values in + configuration objects. +- **Completion script generation** for `bash`, `zsh`, and `tcsh`. +- **Fine-grained configuration** via [PEP + 593](https://peps.python.org/pep-0593/) annotations (`tyro.conf.*`). + +To get started, we recommend browsing the examples to the left. + +### Why `tyro`? + +1. **Strong typing.** + + Unlike tools dependent on dictionaries, YAML, or dynamic namespaces, + arguments populated by `tyro` benefit from IDE and language server-supported + operations — think tab completion, rename, jump-to-def, docstrings on hover — + as well as static checking tools like `pyright` and `mypy`. + +2. **Minimal overhead.** + + Standard Python type annotations, docstrings, and default values are parsed + to automatically generate command-line interfaces with informative helptext. + + `tyro` works seamlessly with tools you already use: examples are included for + [dataclasses](https://docs.python.org/3/library/dataclasses.html), + [attrs](https://www.attrs.org/), + [pydantic](https://pydantic-docs.helpmanual.io/), + [flax.linen](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html), + and more. + +3. **Modularity.** + + `tyro` supports hierarchical configuration structures, which make it easy to + distribute definitions, defaults, and documentation of configurable fields + across modules or source files. + +4. **Tab completion.** + + By extending [shtab](https://github.com/iterative/shtab), `tyro` + automatically generates tab completion scripts for bash, zsh, and tcsh. + + + +.. toctree:: + :caption: Getting started + :hidden: + :maxdepth: 1 + :titlesonly: + + installation + your_first_cli + +.. toctree:: + :caption: Basics + :hidden: + :maxdepth: 1 + :titlesonly: + :glob: + + examples/01_basics/* + + +.. toctree:: + :caption: Hierarchies + :hidden: + :maxdepth: 1 + :titlesonly: + :glob: + + examples/02_nesting/* + + +.. toctree:: + :caption: Config Management + :hidden: + :maxdepth: 1 + :titlesonly: + :glob: + + examples/03_config_systems/* + + +.. toctree:: + :caption: Additional Features + :hidden: + :maxdepth: 1 + :titlesonly: + :glob: + + examples/04_additional/* + + +.. toctree:: + :caption: Notes + :hidden: + :maxdepth: 5 + :glob: + + goals_and_alternatives + helptext_generation + tab_completion + + +.. toctree:: + :caption: API Reference + :hidden: + :maxdepth: 5 + :titlesonly: + + api/tyro/index + + + +.. |build| image:: https://github.com/brentyi/tyro/actions/workflows/build.yml/badge.svg + :alt: Build status icon + :target: https://github.com/brentyi/tyro +.. |mypy| image:: https://github.com/brentyi/tyro/actions/workflows/mypy.yml/badge.svg + :alt: Mypy status icon + :target: https://github.com/brentyi/tyro +.. |pyright| image:: https://github.com/brentyi/tyro/actions/workflows/pyright.yml/badge.svg + :alt: Mypy status icon + :target: https://github.com/brentyi/tyro +.. |ruff| image:: https://github.com/brentyi/tyro/actions/workflows/ruff.yml/badge.svg + :alt: Lint status icon + :target: https://github.com/brentyi/tyro +.. |coverage| image:: https://codecov.io/gh/brentyi/tyro/branch/main/graph/badge.svg + :alt: Test coverage status icon + :target: https://codecov.io/gh/brentyi/tyro +.. |downloads| image:: https://pepy.tech/badge/tyro + :alt: Download count icon + :target: https://pypi.org/project/tyro/ +.. |versions| image:: https://img.shields.io/pypi/pyversions/tyro + :alt: Version icon + :target: https://pypi.org/project/tyro/ +.. |nbsp| unicode:: 0xA0 + :trim: + + diff --git a/docs/source/installation.md b/docs/source/installation.md new file mode 100644 index 000000000..3ffe61bfd --- /dev/null +++ b/docs/source/installation.md @@ -0,0 +1,24 @@ +# Installation + +## Standard + +Installation is supported on Python >=3.7 via pip. + +```bash +pip install tyro +``` + +## Development + +If you're interested in development, the recommended way to install `tyro` is +via `pip`. + +```bash +# Clone repository and install. +git clone git@github.com:brentyi/tyro.git +cd tyro +python -m pip install -e ".[dev]" + +# Run tests. +pytest +``` diff --git a/docs/source/tab_completion.md b/docs/source/tab_completion.md new file mode 100644 index 000000000..a7ac735f0 --- /dev/null +++ b/docs/source/tab_completion.md @@ -0,0 +1,79 @@ +# Tab completion + +Interfaces built with :func:`tyro.cli()` can be tab completed in interactive +shells without any source code modification. + +Completion scripts can be generated by passing the +`--tyro-write-completion {bash/zsh/tcsh} PATH` flag to a tyro CLI. This +generates a completion script via [shtab](https://docs.iterative.ai/shtab/) and +writes it to a specified file. To set up tab completion, the printed script simply +needs to be written somewhere where your shell will find it. + +## Zsh + +For zsh, one option is to emulate the pattern used for completions in +[poetry](https://python-poetry.org/docs): + +```bash +# Set up zsh autocompletion for 01_functions.py, which is located in +# tyro/examples. + +# (1) Make directory for local completions. +mkdir -p ~/.zfunc + +# (2) Write completion script. The name here (_01_functions_py) doesn't matter, +# as long as it's prefixed with an underscore. +python 01_functions.py --tyro-write-completion zsh ~/.zfunc/_01_functions_py +``` + +And if it's not in your `.zshrc` already: + +```bash +# (3) Add .zfunc to our function search path, then initialize completions. +# Ideally this should go in your .zshrc! +fpath+=~/.zfunc +autoload -Uz compinit && compinit +``` + +## Bash + +Local completion scripts for bash can be written as described in the +[bash-complete documentation](https://github.com/scop/bash-completion/blob/master/README.md#FAQ). + +> **Q.** Where should I install my own local completions? +> +> **A.** Put them in the completions subdir of `$BASH_COMPLETION_USER_DIR` +> (defaults to `$XDG_DATA_HOME/bash-completion` or +> `~/.local/share/bash-completion` if `$XDG_DATA_HOME` is not set) to have them +> loaded automatically on demand when the respective command is being completed +> [...] + +Borrowing from the `bash-completion` source[^bash_completion], we can run: + + + +[^bash_completion]: `completion_dir` query is borrowed from [here](https://github.com/scop/bash-completion/blob/966a4e043822613040546e8c003509798c5fae1a/bash_completion#L2440). + + + +```bash +# Set up bash autocompletion for 01_functions.py, which is located in +# tyro/examples. + +# (1) Find and make completion directory. +completion_dir=${BASH_COMPLETION_USER_DIR:-${XDG_DATA_HOME:-$HOME/.local/share}/bash-completion}/completions/ +mkdir -p $completion_dir + +# (2) Write completion scripts. Note that the name of the completion script must +# match the name of the file. +python 01_functions.py --tyro-write-completion bash ${completion_dir}/01_functions.py +``` + +In contrast to zsh, tab completion in bash requires that scripts are either set +up as an entry point or run as :code:`./01_functions.py `, as opposed to +with the `python` command, :code:`python 01_functions.py `. + +Making a script directly executable typically requires: + +1. A permissions update: `chmod +x ./01_functions.py`. +2. A shebang as the first line of your script: `#!/usr/bin/env python` diff --git a/docs/source/your_first_cli.md b/docs/source/your_first_cli.md new file mode 100644 index 000000000..f73dbfba7 --- /dev/null +++ b/docs/source/your_first_cli.md @@ -0,0 +1,83 @@ +# Your first CLI + +To get started with `tyro`, consider the simple `argparse`-based command-line +interface: + +```python +"""Sum two numbers from argparse.""" +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--a", type=int, required=True) +parser.add_argument("--b", type=int, default=3) +args = parser.parse_args() + +total = args.a + args.b + +print(total) +``` + +This is dramatically cleaner than manually parsing `sys.argv`, but has several +issues: it requires a significant amount of parsing-specific boilerplate, lacks +type checking and IDE support (consider: jumping to definitions, finding +references, docstrings, refactoring and renaming tools), and becomes difficult +to manage for larger projects. + +:func:`tyro.cli()` aims to solve these issues. + +**(1) Command-line interfaces from functions.** + +We can write the same script as above using `tyro.cli()`: + +```python +"""Sum two numbers by calling a function with tyro.""" +import tyro + +def add(a: int, b: int = 3) -> int: + return a + b + +# Populate the inputs of add(), call it, then return the output. +total = tyro.cli(add) + +print(total) +``` + +Or, more succinctly: + +```python +"""Sum two numbers by calling a function with tyro.""" +import tyro + +def add(a: int, b: int = 3) -> None: + print(a + b) + +tyro.cli(add) # Returns `None`. +``` + +**(2) Command-line interfaces from config objects.** + +A class in Python can be treated as a function that returns an instance: + +```python +"""Sum two numbers by instantiating a dataclass with tyro.""" +from dataclasses import dataclass + +import tyro + +@dataclass +class Args: + a: int + b: int = 3 + +args = tyro.cli(Args) +print(args.a + args.b) +``` + +Unlike directly using `argparse`, both the function-based and dataclass-based +approaches are compatible with static analysis; tab completion and type checking +will work out-of-the-box. + +And that's it! By incorporating more standard type annotations, we can specify a +broad range of more advanced behaviors: nested structures, variable-length +inputs, unions over types, subcommands, and more. Our examples walk through a +selection of these features. diff --git a/docs/update_example_docs.py b/docs/update_example_docs.py new file mode 100644 index 000000000..475b9dbe2 --- /dev/null +++ b/docs/update_example_docs.py @@ -0,0 +1,139 @@ +"""Helper script for updating the auto-generated examples pages in the documentation.""" + +from __future__ import annotations + +import dataclasses +import pathlib +import shlex +import shutil +from typing import Iterable + +import tyro + + +@dataclasses.dataclass +class ExampleMetadata: + index: str + index_with_zero: str + source: str + title: str + usages: Iterable[str] + description: str + + @staticmethod + def from_path(path: pathlib.Path) -> ExampleMetadata: + # 01_functions -> 01, _, functions. + index, _, _ = path.stem.partition("_") + + # 01 -> 1. + index_with_zero = index + index = str(int(index)) + + source = path.read_text().strip() + + docstring = source.split('"""')[1].strip() + assert "Usage:" in docstring + + title, _, description = docstring.partition("\n") + description, _, usage_text = description.partition("Usage:") + + example_usages = map( + lambda x: x[1:-1], + filter( + lambda line: line.startswith("`") and line.endswith("`"), + usage_text.split("\n"), + ), + ) + return ExampleMetadata( + index=index, + index_with_zero=index_with_zero, + source=source.partition('"""')[2].partition('"""')[2].strip(), + title=title, + usages=example_usages, + description=description.strip(), + ) + + +def get_example_paths(examples_dir: pathlib.Path) -> Iterable[pathlib.Path]: + return filter( + lambda p: not p.name.startswith("_"), sorted(examples_dir.glob("**/*.py")) + ) + + +REPO_ROOT = pathlib.Path(__file__).absolute().parent.parent + + +def main( + examples_dir: pathlib.Path = REPO_ROOT / "examples", + sphinx_source_dir: pathlib.Path = REPO_ROOT / "docs" / "source", +) -> None: + example_doc_dir = sphinx_source_dir / "examples" + shutil.rmtree(example_doc_dir) + example_doc_dir.mkdir() + + for path in get_example_paths(examples_dir): + ex = ExampleMetadata.from_path(path) + path_for_sphinx = pathlib.Path("..") / ".." / path.relative_to(REPO_ROOT) + + usage_lines = [] + for usage in ex.usages: + args = shlex.split(usage) + python_index = args.index("python") + sphinx_usage = shlex.join( + args[:python_index] + + ["python", path_for_sphinx.as_posix()] + + args[python_index + 2 :] + ) + + # Note that :kbd: in Sphinx does unnecessary stuff we want to avoid, see: + # https://github.com/sphinx-doc/sphinx/issues/7530 + # + # Instead, we just use raw HTML. + assert "../../examples/" in sphinx_usage + command = sphinx_usage.replace("../../examples/", "") + usage_lines += [ + "------------", + "", + ".. raw:: html", + "", + f" {command}", + "", + f".. program-output:: {sphinx_usage}", + "", + ] + + relative_dir = path.parent.relative_to(examples_dir) + target_dir = example_doc_dir / relative_dir + target_dir.mkdir(exist_ok=True, parents=True) + + (target_dir / f"{path.stem}.rst").write_text( + "\n".join( + [ + ( + ".. Comment: this file is automatically generated by" + " `update_example_docs.py`." + ), + " It should not be modified manually.", + "", + f"{ex.title}", + "==========================================", + "", + ex.description, + "", + "", + ".. code-block:: python", + " :linenos:", + "", + "", + "\n".join( + f" {line}".rstrip() for line in ex.source.split("\n") + ), + "", + ] + + usage_lines + ) + ) + + +if __name__ == "__main__": + tyro.cli(main, description=__doc__) diff --git a/examples/01_basics/01_functions.py b/examples/01_basics/01_functions.py new file mode 100644 index 000000000..d9fe17413 --- /dev/null +++ b/examples/01_basics/01_functions.py @@ -0,0 +1,29 @@ +"""Functions + +In the simplest case, :func:`tyro.cli()` can be used to run a function with +arguments populated from the CLI. + +Usage: +`python ./01_functions.py --help` +`python ./01_functions.py --field1 hello` +`python ./01_functions.py --field1 hello --field2 10` +""" + +import tyro + + +def main( + field1: str, + field2: int = 3, +) -> None: + """Function, whose arguments will be populated from a CLI interface. + + Args: + field1: A string field. + field2: A numeric field, with a default value. + """ + print(field1, field2) + + +if __name__ == "__main__": + tyro.cli(main) diff --git a/examples/01_basics/02_dataclasses.py b/examples/01_basics/02_dataclasses.py new file mode 100644 index 000000000..09e94bbfc --- /dev/null +++ b/examples/01_basics/02_dataclasses.py @@ -0,0 +1,30 @@ +"""Dataclasses + +Common pattern: use :func:`tyro.cli()` to instantiate a dataclass. + +Usage: +`python ./02_dataclasses.py --help` +`python ./02_dataclasses.py --field1 hello` +`python ./02_dataclasses.py --field1 hello --field2 5` +""" + +import dataclasses + +import tyro + + +@dataclasses.dataclass +class Args: + """Description. + This should show up in the helptext!""" + + field1: str + """A string field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/01_basics/03_dataclasses_defaults.py b/examples/01_basics/03_dataclasses_defaults.py new file mode 100644 index 000000000..2a5670383 --- /dev/null +++ b/examples/01_basics/03_dataclasses_defaults.py @@ -0,0 +1,54 @@ +"""Dataclasses + Defaults + +The :code:`default=` argument can be used to override default values in dataclass +types. + + +.. warning:: + We advise against mutation of configuration objects from a dataclass's + :code:`__post_init__` method [#f1]_. In the example below, + :code:`__post_init__` would be called twice: once for the :code:`Args()` + object provided as a default value and another time for the :code:`Args()` + objected instantiated by :func:`tyro.cli()`. This can cause confusing + behavior! Instead, we show below one example of how derived fields can be + defined immutably. + + .. [#f1] Official Python docs for ``__post_init__`` can be found `here `_. + + +Usage: +`python ./03_dataclasses_defaults.py --help` +`python ./03_dataclasses_defaults.py --field2 3` +`python ./03_dataclasses_defaults.py --field1 hello --field2 5` +""" + +import dataclasses + +import tyro + + +@dataclasses.dataclass +class Args: + """Description. + This should show up in the helptext!""" + + field1: str + """A string field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + @property + def derived_field(self) -> str: + return ", ".join([self.field1] * self.field2) + + +if __name__ == "__main__": + args = tyro.cli( + Args, + default=Args( + field1="default string", + field2=tyro.MISSING, + ), + ) + print(args.derived_field) diff --git a/examples/01_basics/04_collections.py b/examples/01_basics/04_collections.py new file mode 100644 index 000000000..a3f3b836e --- /dev/null +++ b/examples/01_basics/04_collections.py @@ -0,0 +1,34 @@ +"""Multi-value Arguments + +Arguments of both fixed and variable lengths can be annotated with standard +Python collection types. For Python 3.7 and 3.8, we can use either :code:`from +__future__ import annotations` to support :code:`list[T]` and :code:`tuple[T]`, +or the older API :code:`typing.List[T]` and :code:`typing.Tuple[T1, T2]`. + +Usage: +`python ./03_collections.py --help` +`python ./03_collections.py --dataset-sources ./data --image-dimensions 16 16` +`python ./03_collections.py --dataset-sources ./data` +""" + +import dataclasses +import pathlib + +import tyro + + +@dataclasses.dataclass(frozen=True) +class TrainConfig: + # Example of a variable-length tuple. `list[T]`, `set[T]`, + # `dict[K, V]`, etc are supported as well. + dataset_sources: tuple[pathlib.Path, ...] + """Paths to load training data from. This can be multiple!""" + + # Fixed-length tuples are also okay. + image_dimensions: tuple[int, int] = (32, 32) + """Height and width of some image data.""" + + +if __name__ == "__main__": + config = tyro.cli(TrainConfig) + print(config) diff --git a/examples/01_basics/05_flags.py b/examples/01_basics/05_flags.py new file mode 100644 index 000000000..11c07b8bf --- /dev/null +++ b/examples/01_basics/05_flags.py @@ -0,0 +1,37 @@ +"""Booleans and Flags + +Booleans can either be expected to be explicitly passed in, or, if given a default +value, automatically converted to flags. + +To turn off conversion, see :class:`tyro.conf.FlagConversionOff`. + +Usage: +`python ./05_flags.py --help` +`python ./05_flags.py --boolean True` +`python ./05_flags.py --boolean False --flag-a` +`python ./05_flags.py --boolean False --no-flag-b` +""" + +import dataclasses + +import tyro + + +@dataclasses.dataclass +class Args: + # Boolean. This expects an explicit "True" or "False". + boolean: bool + + # Optional boolean. Same as above, but can be omitted. + optional_boolean: bool | None = None + + # Pass --flag-a in to set this value to True. + flag_a: bool = False + + # Pass --no-flag-b in to set this value to False. + flag_b: bool = True + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/01_basics/06_literals.py b/examples/01_basics/06_literals.py new file mode 100644 index 000000000..5f3017d85 --- /dev/null +++ b/examples/01_basics/06_literals.py @@ -0,0 +1,27 @@ +"""Choices + +:code:`typing.Literal[]` can be used to restrict inputs to a fixed set of literal choices. + +Usage: +`python ./06_literals.py --help` +""" + +import dataclasses +from typing import Literal + +import tyro + + +@dataclasses.dataclass(frozen=True) +class Args: + # We can use Literal[] to restrict the set of allowable inputs, for example, over + # a set of strings. + strings: Literal["red", "green"] = "red" + + # Integers also work. (as well as booleans, enums, etc) + numbers: Literal[0, 1, 2] = 0 + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/01_basics/07_unions.py b/examples/01_basics/07_unions.py new file mode 100644 index 000000000..f99ffb04d --- /dev/null +++ b/examples/01_basics/07_unions.py @@ -0,0 +1,44 @@ +"""Unions + +:code:`X | Y` or :code:`typing.Union[X, Y]` can be used to expand inputs to +multiple types. + +Usage: +`python ./07_unions.py --help` +""" + +import dataclasses +import enum +from typing import Literal, Optional + +import tyro + + +class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + +@dataclasses.dataclass(frozen=True) +class Args: + # Unions can be used to specify multiple allowable types. + union_over_types: int | str = 0 + string_or_enum: Literal["red", "green"] | Color = "red" + + # Unions also work over more complex nested types. + union_over_tuples: tuple[int, int] | tuple[str] = ("1",) + + # And can be nested in other types. + tuple_of_string_or_enum: tuple[Literal["red", "green"] | Color, ...] = ( + "red", + Color.RED, + ) + + # Optional[T] is equivalent to `T | None`. + integer: Optional[Literal[0, 1, 2, 3]] = None + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/01_basics/08_enums.py b/examples/01_basics/08_enums.py new file mode 100644 index 000000000..0fa47eed4 --- /dev/null +++ b/examples/01_basics/08_enums.py @@ -0,0 +1,35 @@ +"""Enums + +In addition to literals, enums can also be used to provide a fixed set of +choices. + +Usage: +`python ./04_enums.py --help` +`python ./04_enums.py --optimizer-type SGD` +`python ./04_enums.py --optimizer-type ADAM --learning-rate 3e-4` +""" + +import dataclasses +import enum + +import tyro + + +class OptimizerType(enum.Enum): + ADAM = enum.auto() + SGD = enum.auto() + + +@dataclasses.dataclass(frozen=True) +class TrainConfig: + # Enums are handled seamlessly. + optimizer_type: OptimizerType = OptimizerType.ADAM + """Gradient-based optimizer to use.""" + + learning_rate: float = 1e-4 + """Learning rate for optimizer.""" + + +if __name__ == "__main__": + config = tyro.cli(TrainConfig) + print(config) diff --git a/examples/02_nesting/01_nesting.py b/examples/02_nesting/01_nesting.py new file mode 100644 index 000000000..b6e673056 --- /dev/null +++ b/examples/02_nesting/01_nesting.py @@ -0,0 +1,72 @@ +"""Hierarchical Configs + +Structures (typically dataclasses) can be nested to build hierarchical configuration +objects. This helps with modularity and grouping in larger projects. + +Usage: +`python ./01_nesting.py --help` +`python ./01_nesting.py --out-dir . --config.optimizer.algorithm SGD` +`python ./01_nesting.py --out-dir . --restore-checkpoint` +""" + +import dataclasses +import enum +import pathlib + +import tyro + + +class OptimizerType(enum.Enum): + ADAM = enum.auto() + SGD = enum.auto() + + +@dataclasses.dataclass +class OptimizerConfig: + # Gradient-based optimizer to use. + algorithm: OptimizerType = OptimizerType.ADAM + + # Learning rate to use. + learning_rate: float = 3e-4 + + # Coefficient for L2 regularization. + weight_decay: float = 1e-2 + + +@dataclasses.dataclass +class ExperimentConfig: + # Various configurable options for our optimizer. + optimizer: OptimizerConfig + + # Batch size. + batch_size: int = 32 + + # Total number of training steps. + train_steps: int = 100_000 + + # Random seed. This is helpful for making sure that our experiments are all + # reproducible! + seed: int = 0 + + +def train( + out_dir: pathlib.Path, + config: ExperimentConfig, + restore_checkpoint: bool = False, + checkpoint_interval: int = 1000, +) -> None: + """Train a model. + + Args: + out_dir: Where to save logs and checkpoints. + config: Experiment configuration. + restore_checkpoint: Set to restore an existing checkpoint. + checkpoint_interval: Training steps between each checkpoint save. + """ + print(f"{out_dir=}, {restore_checkpoint=}, {checkpoint_interval=}") + print() + print(f"{config=}") + + +if __name__ == "__main__": + tyro.cli(train) diff --git a/examples/02_nesting/02_subcommands.py b/examples/02_nesting/02_subcommands.py new file mode 100644 index 000000000..5ccc26610 --- /dev/null +++ b/examples/02_nesting/02_subcommands.py @@ -0,0 +1,46 @@ +"""Subcommands + +Unions over nested types (classes or dataclasses) are populated using subcommands. + +For configuring subcommands beyond what can be expressed with type annotations, see +:func:`tyro.conf.subcommand()`. + +Usage: +`python ./02_subcommands.py --help` +`python ./02_subcommands.py cmd:commit --help` +`python ./02_subcommands.py cmd:commit --cmd.message hello --cmd.all` +`python ./02_subcommands.py cmd:checkout --help` +`python ./02_subcommands.py cmd:checkout --cmd.branch main` +""" + +from __future__ import annotations + +import dataclasses + +import tyro + + +@dataclasses.dataclass(frozen=True) +class Checkout: + """Checkout a branch.""" + + branch: str + + +@dataclasses.dataclass(frozen=True) +class Commit: + """Commit changes.""" + + message: str + all: bool = False + + +def main(cmd: Checkout | Commit) -> None: + print(cmd) + + +if __name__ == "__main__": + # Note that we can also pass `Checkout | Command` directly into + # `tyro.cli()`; this is understood by tyro and pyright, but unfortunately not by + # mypy. + tyro.cli(main) diff --git a/examples/02_nesting/03_multiple_subcommands.py b/examples/02_nesting/03_multiple_subcommands.py new file mode 100644 index 000000000..561232b54 --- /dev/null +++ b/examples/02_nesting/03_multiple_subcommands.py @@ -0,0 +1,69 @@ +"""Sequenced Subcommands + +Multiple unions over nested types are populated using a series of subcommands. + +Usage: +`python ./03_multiple_subcommands.py --help` +`python ./03_multiple_subcommands.py dataset:mnist --help` +`python ./03_multiple_subcommands.py dataset:mnist optimizer:adam --help` +`python ./03_multiple_subcommands.py dataset:mnist optimizer:adam --optimizer.learning-rate 3e-4 --dataset.binary` +""" + +from __future__ import annotations + +import dataclasses +from typing import Literal + +import tyro + +# Possible dataset configurations. + + +@dataclasses.dataclass +class Mnist: + binary: bool = False + """Set to load binary version of MNIST dataset.""" + + +@dataclasses.dataclass +class ImageNet: + subset: Literal[50, 100, 1000] + """Choose between ImageNet-50, ImageNet-100, ImageNet-1000, etc.""" + + +# Possible optimizer configurations. + + +@dataclasses.dataclass +class Adam: + learning_rate: float = 1e-3 + betas: tuple[float, float] = (0.9, 0.999) + + +@dataclasses.dataclass +class Sgd: + learning_rate: float = 3e-4 + + +# Train script. + + +def train( + dataset: Mnist | ImageNet = Mnist(), + optimizer: Adam | Sgd = Adam(), +) -> None: + """Example training script. + + Args: + dataset: Dataset to train on. + optimizer: Optimizer to train with. + + Returns: + None: + """ + print(dataset) + print(optimizer) + + +if __name__ == "__main__": + tyro.cli(train, config=(tyro.conf.ConsolidateSubcommandArgs,)) diff --git a/examples/02_nesting/04_nesting_in_containers.py b/examples/02_nesting/04_nesting_in_containers.py new file mode 100644 index 000000000..18892696d --- /dev/null +++ b/examples/02_nesting/04_nesting_in_containers.py @@ -0,0 +1,60 @@ +"""Nesting in Containers + +Structures can be nested inside of standard containers. + +Note that lengths must be inferable, either via a fixed-length tuple annotation or by +parsing default values. + + +Usage: +`python ./04_nesting_in_containers.py.py --help` +""" + +import dataclasses + +import tyro + + +class Color: + pass + + +@dataclasses.dataclass +class RGB(Color): + r: int + g: int + b: int + + +@dataclasses.dataclass +class HSL(Color): + h: int + s: int + l: int + + +@dataclasses.dataclass +class Args: + # Example of specifying nested structures via a fixed-length tuple. + color_tuple: tuple[RGB, HSL] + + # Examples of nested structures in variable-length containers. These need a default + # provided for length inference; we don't currently support specifying dynamic + # container lengths directly from the commandline. + color_tuple_alt: tuple[Color, ...] = ( + RGB(255, 0, 0), + HSL(0, 255, 0), + ) + color_map: dict[str, RGB] = dataclasses.field( + # We can't use mutable values as defaults directly. + default_factory={ + "red": RGB(255, 0, 0), + "green": RGB(0, 255, 0), + "blue": RGB(0, 0, 255), + }.copy + ) + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/02_nesting/05_subcommands_func.py b/examples/02_nesting/05_subcommands_func.py new file mode 100644 index 000000000..ca2f202f9 --- /dev/null +++ b/examples/02_nesting/05_subcommands_func.py @@ -0,0 +1,36 @@ +"""Subcommands from Functions + +We provide a shorthand for generating a subcommand CLI from a dictionary. This +is a thin wrapper around :func:`tyro.cli()`'s more verbose, type-based API. If +more generality is needed, the internal working are explained in the docs for +:func:`tyro.extras.subcommand_cli_from_dict()`. + + +Usage: +`python ./05_subcommands_func.py --help` +`python ./05_subcommands_func.py commit --help` +`python ./05_subcommands_func.py commit --message hello --all` +`python ./05_subcommands_func.py checkout --help` +`python ./05_subcommands_func.py checkout --branch main` +""" + +import tyro + + +def checkout(branch: str) -> None: + """Check out a branch.""" + print(f"{branch=}") + + +def commit(message: str, all: bool = False) -> None: + """Make a commit.""" + print(f"{message=} {all=}") + + +if __name__ == "__main__": + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + } + ) diff --git a/examples/03_config_systems/01_base_configs.py b/examples/03_config_systems/01_base_configs.py new file mode 100644 index 000000000..50ee71fa4 --- /dev/null +++ b/examples/03_config_systems/01_base_configs.py @@ -0,0 +1,92 @@ +"""Base Configurations + +We can integrate `tyro` into common configuration patterns: here, we select +one of multiple possible base configurations, create a subcommand for each one, and then +use the CLI to either override (existing) or fill in (missing) values. + +The helper function used here, :func:`tyro.extras.overridable_config_cli()`, is a +lightweight wrapper over :func:`tyro.cli()`. + + +Usage: +`python ./01_base_configs.py --help` +`python ./01_base_configs.py small --help` +`python ./01_base_configs.py small --seed 94720` +`python ./01_base_configs.py big --help` +`python ./01_base_configs.py big --seed 94720` +""" + +from dataclasses import dataclass +from typing import Callable, Literal + +from torch import nn + +import tyro + + +@dataclass(frozen=True) +class AdamOptimizer: + learning_rate: float = 1e-3 + betas: tuple[float, float] = (0.9, 0.999) + + +@dataclass(frozen=True) +class ExperimentConfig: + # Dataset to run experiment on. + dataset: Literal["mnist", "imagenet-50"] + + # Optimizer parameters. + optimizer: AdamOptimizer + + # Model size. + num_layers: int + units: int + + # Batch size. + batch_size: int + + # Total number of training steps. + train_steps: int + + # Random seed. This is helpful for making sure that our experiments are all + # reproducible! + seed: int + + # Activation to use. Not specifiable via the commandline. + activation: Callable[[], nn.Module] + + +# Note that we could also define this library using separate YAML files (similar to +# `config_path`/`config_name` in Hydra), but staying in Python enables seamless type +# checking + IDE support. +default_configs = { + "small": ( + "Small experiment.", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + ), + "big": ( + "Big experiment.", + ExperimentConfig( + dataset="imagenet-50", + optimizer=AdamOptimizer(), + batch_size=32, + num_layers=8, + units=256, + train_steps=100_000, + seed=0, + activation=nn.GELU, + ), + ), +} +if __name__ == "__main__": + config = tyro.extras.overridable_config_cli(default_configs) + print(config) diff --git a/examples/03_config_systems/02_overriding_yaml.py b/examples/03_config_systems/02_overriding_yaml.py new file mode 100644 index 000000000..5db0cfd86 --- /dev/null +++ b/examples/03_config_systems/02_overriding_yaml.py @@ -0,0 +1,41 @@ +"""Overriding YAML Configs + +If you have a library of existing YAML files that you want to use, `tyro` can +be used to override values in them. We generally recommend dataclass configs +for new projects. + +Usage: +`python ./02_overriding_yaml.py --help` +`python ./02_overriding_yaml.py --training.checkpoint-steps 300 1000 9000` +""" + +import yaml + +import tyro + +# YAML configuration. Note that this could also be loaded from a file! Environment +# variables are an easy way to select between different YAML files. +default_yaml = r""" +exp_name: test +optimizer: + learning_rate: 0.0001 + type: adam +training: + batch_size: 32 + num_steps: 10000 + checkpoint_steps: + - 500 + - 1000 + - 1500 +""".strip() + +if __name__ == "__main__": + # Convert our YAML config into a nested dictionary. + default_config = yaml.safe_load(default_yaml) + + # Override fields in the dictionary. + overridden_config = tyro.cli(dict, default=default_config) + + # Print the overridden config. + overridden_yaml = yaml.safe_dump(overridden_config) + print(overridden_yaml) diff --git a/examples/04_additional/01_positional_args.py b/examples/04_additional/01_positional_args.py new file mode 100644 index 000000000..66048cd53 --- /dev/null +++ b/examples/04_additional/01_positional_args.py @@ -0,0 +1,63 @@ +"""Positional Arguments + +Positional-only arguments in functions are converted to positional CLI arguments. + +For more general positional arguments, see :class:`tyro.conf.Positional`. + +Usage: +`python ./01_positional_args.py --help` +`python ./01_positional_args.py ./a ./b --optimizer.learning-rate 1e-5` +""" + +from __future__ import annotations + +import dataclasses +import enum +import pathlib +from typing import Tuple + +import tyro + + +def main( + source: pathlib.Path, + dest: pathlib.Path, + /, # Mark the end of positional arguments. + optimizer: OptimizerConfig, + force: bool = False, + verbose: bool = False, + background_rgb: Tuple[float, float, float] = (1.0, 0.0, 0.0), +) -> None: + """Command-line interface defined using a function signature. Note that this + docstring is parsed to generate helptext. + + Args: + source: Source path. + dest: Destination path. + optimizer: Configuration for our optimizer object. + force: Do not prompt before overwriting. + verbose: Explain what is being done. + background_rgb: Background color. Red by default. + """ + print(f"{source=}\n{dest=}\n{optimizer=}\n{force=}\n{verbose=}\n{background_rgb=}") + + +class OptimizerType(enum.Enum): + ADAM = enum.auto() + SGD = enum.auto() + + +@dataclasses.dataclass(frozen=True) +class OptimizerConfig: + algorithm: OptimizerType = OptimizerType.ADAM + """Gradient-based optimizer to use.""" + + learning_rate: float = 3e-4 + """Learning rate to use.""" + + weight_decay: float = 1e-2 + """Coefficient for L2 regularization.""" + + +if __name__ == "__main__": + tyro.cli(main) diff --git a/examples/04_additional/02_dictionaries.py b/examples/04_additional/02_dictionaries.py new file mode 100644 index 000000000..9b0bac6f0 --- /dev/null +++ b/examples/04_additional/02_dictionaries.py @@ -0,0 +1,55 @@ +"""Dictionaries and TypedDict + +Dictionary inputs can be specified using either a standard `Dict[K, V]` +annotation, or a :code:`TypedDict` subclass. + +For configuring :code:`TypedDict`, we also support :code:`total={True/False}`, +:code:`typing.Required`, and :code:`typing.NotRequired`. See the `Python docs `_ for all :code:`TypedDict` features. + +Usage: +`python ./02_dictionaries.py --help` +`python ./02_dictionaries.py --typed-dict-a.learning-rate 3e-4 --typed-dict-b.betas 0.9 0.999` +`python ./02_dictionaries.py --typed-dict-b.betas 0.9 0.999` +""" + +from typing import TypedDict + +from typing_extensions import NotRequired + +import tyro + + +class DictionarySchemaA( + TypedDict, + # Setting `total=False` specifies that not all keys need to exist. + total=False, +): + learning_rate: float + betas: tuple[float, float] + + +class DictionarySchemaB(TypedDict): + learning_rate: NotRequired[float] + """NotRequired[] specifies that a particular key doesn't need to exist.""" + betas: tuple[float, float] + + +def main( + typed_dict_a: DictionarySchemaA, + typed_dict_b: DictionarySchemaB, + standard_dict: dict[str, float] = { + "learning_rate": 3e-4, + "beta1": 0.9, + "beta2": 0.999, + }, +) -> None: + assert isinstance(typed_dict_a, dict) + assert isinstance(typed_dict_b, dict) + assert isinstance(standard_dict, dict) + print("Typed dict A:", typed_dict_a) + print("Typed dict B:", typed_dict_b) + print("Standard dict:", standard_dict) + + +if __name__ == "__main__": + tyro.cli(main) diff --git a/examples/04_additional/03_tuples.py b/examples/04_additional/03_tuples.py new file mode 100644 index 000000000..ef40280c8 --- /dev/null +++ b/examples/04_additional/03_tuples.py @@ -0,0 +1,38 @@ +"""Tuples + +Example using :func:`tyro.cli()` to instantiate tuple types. :code:`tuple`, +:code:`typing.Tuple`, and :code:`NamedTuple` are all supported. + +Usage: +`python ./03_tuples.py --help` +`python ./03_tuples.py --color 127 127 127` +`python ./03_tuples.py --two-colors.1.r 127 --two-colors.1.g 0 --two-colors.1.b 0` +""" + +from typing import NamedTuple + +import tyro + + +# Named tuples are interpreted as nested structures. +class Color(NamedTuple): + r: int + g: int + b: int + + +class TupleType(NamedTuple): + """Description. + This should show up in the helptext!""" + + # Tuple types can contain raw values. + color: tuple[int, int, int] = (255, 0, 0) + + # Tuple types can contain nested structures. + two_colors: tuple[Color, Color] = (Color(255, 0, 0), Color(0, 255, 0)) + + +if __name__ == "__main__": + x = tyro.cli(TupleType) + assert isinstance(x, tuple) + print(x) diff --git a/examples/04_additional/04_classes.py b/examples/04_additional/04_classes.py new file mode 100644 index 000000000..5dbae4639 --- /dev/null +++ b/examples/04_additional/04_classes.py @@ -0,0 +1,33 @@ +"""Instantiating Classes + +In addition to functions and dataclasses, we can also generate CLIs from the +constructors of standard Python classes. + +Usage: +`python ./04_classes.py --help` +`python ./04_classes.py --field1 hello --field2 7` +""" + +import tyro + + +class Args: + def __init__( + self, + field1: str, + field2: int, + flag: bool = False, + ): + """Arguments. + + Args: + field1: A string field. + field2: A numeric field. + flag: A boolean flag. + """ + self.data = [field1, field2, flag] + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args.data) diff --git a/examples/04_additional/05_generics.py b/examples/04_additional/05_generics.py new file mode 100644 index 000000000..89f96b242 --- /dev/null +++ b/examples/04_additional/05_generics.py @@ -0,0 +1,40 @@ +"""Generic Types + +Example of parsing for generic dataclasses. + +Usage: +`python ./05_generics.py --help` +""" + +import dataclasses +from typing import Generic, TypeVar + +import tyro + +ScalarType = TypeVar("ScalarType", int, float) +ShapeType = TypeVar("ShapeType") + + +@dataclasses.dataclass(frozen=True) +class Point3(Generic[ScalarType]): + x: ScalarType + y: ScalarType + z: ScalarType + frame_id: str + + +@dataclasses.dataclass(frozen=True) +class Triangle: + a: Point3[float] + b: Point3[float] + c: Point3[float] + + +@dataclasses.dataclass(frozen=True) +class Args(Generic[ShapeType]): + shape: ShapeType + + +if __name__ == "__main__": + args = tyro.cli(Args[Triangle]) + print(args) diff --git a/examples/04_additional/06_generics_py312.py b/examples/04_additional/06_generics_py312.py new file mode 100644 index 000000000..2fd4b4580 --- /dev/null +++ b/examples/04_additional/06_generics_py312.py @@ -0,0 +1,43 @@ +# mypy: ignore-errors +# +# PEP 695 isn't yet supported in mypy. (April 4, 2024) +"""Generic Types (Python 3.12+ syntax) + +Example of parsing for generic dataclasses using syntax introduced in Python +3.12 (`PEP 695 `_). + +.. warning:: + If used in conjunction with :code:`from __future__ import annotations`, the updated type parameter syntax requires Python 3.12.4 or newer. For technical details, see `this CPython PR `_. + +Usage: +`python ./05_generics.py --help` +""" + +import dataclasses + +import tyro + + +@dataclasses.dataclass(frozen=True) +class Point3[ScalarType: (int, float)]: + x: ScalarType + y: ScalarType + z: ScalarType + frame_id: str + + +@dataclasses.dataclass(frozen=True) +class Triangle: + a: Point3[float] + b: Point3[float] + c: Point3[float] + + +@dataclasses.dataclass(frozen=True) +class Args[ShapeType]: + shape: ShapeType + + +if __name__ == "__main__": + args = tyro.cli(Args[Triangle]) + print(args) diff --git a/examples/04_additional/07_conf.py b/examples/04_additional/07_conf.py new file mode 100644 index 000000000..479a888a9 --- /dev/null +++ b/examples/04_additional/07_conf.py @@ -0,0 +1,43 @@ +"""Configuration via typing.Annotated[] + +The :mod:`tyro.conf` module contains utilities that can be used to configure +command-line interfaces beyond what is expressible via static type annotations. + +Features here are supported, but generally unnecessary and should be used sparingly. + +Usage: +`python ./06_conf.py --help` +`python ./06_conf.py 5 --boolean True` +""" + +import dataclasses + +from typing_extensions import Annotated + +import tyro + + +@dataclasses.dataclass +class Args: + # A numeric field parsed as a positional argument. + positional: tyro.conf.Positional[int] + + # A boolean field with flag conversion turned off. + boolean: tyro.conf.FlagConversionOff[bool] = False + + # A numeric field that can't be changed via the CLI. + fixed: tyro.conf.Fixed[int] = 5 + + # A field with manually overridden properties. + manual: Annotated[ + str, + tyro.conf.arg( + name="renamed", + metavar="STRING", + help="A field with manually overridden properties!", + ), + ] = "Hello" + + +if __name__ == "__main__": + print(tyro.cli(Args)) diff --git a/examples/04_additional/08_pydantic.py b/examples/04_additional/08_pydantic.py new file mode 100644 index 000000000..e9c5ee457 --- /dev/null +++ b/examples/04_additional/08_pydantic.py @@ -0,0 +1,27 @@ +"""Pydantic Integration + +In addition to standard dataclasses, :func:`tyro.cli()` also supports +`Pydantic `_ models. + +Usage: +`python ./08_pydantic.py --help` +`python ./08_pydantic.py --field1 hello` +`python ./08_pydantic.py --field1 hello --field2 5` +""" + +from pydantic import BaseModel, Field + +import tyro + + +class Args(BaseModel): + """Description. + This should show up in the helptext!""" + + field1: str + field2: int = Field(3, description="An integer field.") + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/04_additional/09_attrs.py b/examples/04_additional/09_attrs.py new file mode 100644 index 000000000..effe864e9 --- /dev/null +++ b/examples/04_additional/09_attrs.py @@ -0,0 +1,31 @@ +"""Attrs Integration + +In addition to standard dataclasses, :func:`tyro.cli()` also supports +`attrs `_ classes. + +Usage: +`python ./09_attrs.py --help` +`python ./09_attrs.py --field1 hello` +`python ./09_attrs.py --field1 hello --field2 5` +""" + +import attr + +import tyro + + +@attr.s +class Args: + """Description. + This should show up in the helptext!""" + + field1: str = attr.ib() + """A string field.""" + + field2: int = attr.ib(factory=lambda: 5) + """A required integer field.""" + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/04_additional/10_flax.py b/examples/04_additional/10_flax.py new file mode 100644 index 000000000..ecc5c50fc --- /dev/null +++ b/examples/04_additional/10_flax.py @@ -0,0 +1,54 @@ +"""JAX/Flax Integration + +If you use `flax.linen `_, modules can be instantiated +directly from :func:`tyro.cli()`. + +Usage: +`python ./07_flax.py --help` +`python ./07_flax.py --model.layers 4` +""" + +from flax import linen as nn +from jax import numpy as jnp + +import tyro + + +class Classifier(nn.Module): + layers: int + """Layers in our network.""" + units: int = 32 + """Hidden unit count.""" + output_dim: int = 10 + """Number of classes.""" + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # type: ignore + for i in range(self.layers - 1): + x = nn.Dense( + self.units, + kernel_init=nn.initializers.kaiming_normal(), + )(x) + x = nn.relu(x) + + x = nn.Dense( + self.output_dim, + kernel_init=nn.initializers.xavier_normal(), + )(x) + x = nn.sigmoid(x) + return x + + +def train(model: Classifier, num_iterations: int = 1000) -> None: + """Train a model. + + Args: + model: Model to train. + num_iterations: Number of training iterations. + """ + print(f"{model=}") + print(f"{num_iterations=}") + + +if __name__ == "__main__": + tyro.cli(train) diff --git a/examples/04_additional/11_custom_constructors.py b/examples/04_additional/11_custom_constructors.py new file mode 100644 index 000000000..0bbcebb2d --- /dev/null +++ b/examples/04_additional/11_custom_constructors.py @@ -0,0 +1,69 @@ +"""Custom Constructors + +For additional flexibility, :func:`tyro.conf.arg()` accepts a +:code:`constructor` argument, which makes it easier to load complex objects. + +.. warning:: + Custom constructors are permitted wherever :func:`tyro.conf.arg()` is. This + annotation corresponds to a single argument, so we expect that it is placed + placed at the root of the field corresponding to that argument. For + example: + + .. code-block:: python + + # Great! + x: Annotated[int, tyro.conf.arg(custom_constructor=...)] + + However, nesting :func:`tyro.conf.arg()` within other types is generally + not supported. For example: + + .. code-block:: python + + # Not supported. + x: tuple[Annotated[int, tyro.conf.arg(custom_constructor=...)], ...] + + This example can be rewritten with a custom constructor that directly + returns a tuple of integers: + + .. code-block:: python + + # Workaround for above. + x: Annotated[tuple[int, ...], tyro.conf.arg(custom_constructor=...)] + + +Usage: +`python ./10_custom_constructors.py --help` +`python ./10_custom_constructors.py --dict1.json "{\"hello\": \"world\"}"` +`python ./10_custom_constructors.py --dict1.json "{\"hello\": \"world\"}"` --dict2.json "{\"hello\": \"world\"}"` +""" + +import json as json_ + +from typing_extensions import Annotated + +import tyro + + +def dict_json_constructor(json: str) -> dict: + """Construct a dictionary from a JSON string. Raises a ValueError if the result is + not a dictionary.""" + out = json_.loads(json) + if not isinstance(out, dict): + raise ValueError(f"{json} is not a dictionary!") + return out + + +# A dictionary type, but `tyro` will expect a JSON string from the CLI. +JsonDict = Annotated[dict, tyro.conf.arg(constructor=dict_json_constructor)] + + +def main( + dict1: JsonDict, + dict2: JsonDict = {"default": None}, +) -> None: + print(f"{dict1=}") + print(f"{dict2=}") + + +if __name__ == "__main__": + tyro.cli(main) diff --git a/examples/04_additional/12_aliases.py b/examples/04_additional/12_aliases.py new file mode 100644 index 000000000..5ee33de9c --- /dev/null +++ b/examples/04_additional/12_aliases.py @@ -0,0 +1,41 @@ +"""Argument Aliases + +:func:`tyro.conf.arg()` can be used to attach aliases to arguments. + +Usage: +`python ./11_aliases.py --help` +`python ./11_aliases.py commit --help` +`python ./11_aliases.py commit --message hello --all` +`python ./11_aliases.py commit -m hello -a` +`python ./11_aliases.py checkout --help` +`python ./11_aliases.py checkout --branch main` +`python ./11_aliases.py checkout -b main` +""" + +from typing_extensions import Annotated + +import tyro + + +def checkout( + branch: Annotated[str, tyro.conf.arg(aliases=["-b"])], +) -> None: + """Check out a branch.""" + print(f"{branch=}") + + +def commit( + message: Annotated[str, tyro.conf.arg(aliases=["-m"])], + all: Annotated[bool, tyro.conf.arg(aliases=["-a"])] = False, +) -> None: + """Make a commit.""" + print(f"{message=} {all=}") + + +if __name__ == "__main__": + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + } + ) diff --git a/examples/04_additional/12_counters.py b/examples/04_additional/12_counters.py new file mode 100644 index 000000000..2a6dd2bbf --- /dev/null +++ b/examples/04_additional/12_counters.py @@ -0,0 +1,33 @@ +"""Counters + +Repeatable 'counter' arguments can be specified via :data:`tyro.conf.UseCounterAction`. + +Usage: +`python ./12_counters.py --help` +`python ./12_counters.py --verbosity` +`python ./12_counters.py --verbosity --verbosity` +`python ./12_counters.py -vvv` +""" + +from typing_extensions import Annotated + +import tyro +from tyro.conf import UseCounterAction + + +def main( + verbosity: UseCounterAction[int], + aliased_verbosity: Annotated[UseCounterAction[int], tyro.conf.arg(aliases=["-v"])], +) -> None: + """Example showing how to use counter actions. + + Args: + verbosity: Verbosity level. + aliased_verbosity: Same as above, but can also be specified with -v, -vv, -vvv, etc. + """ + print("Verbosity level:", verbosity) + print("Verbosity level (aliased):", aliased_verbosity) + + +if __name__ == "__main__": + tyro.cli(main) diff --git a/examples/04_additional/13_type_statement.py b/examples/04_additional/13_type_statement.py new file mode 100644 index 000000000..023d5c6f6 --- /dev/null +++ b/examples/04_additional/13_type_statement.py @@ -0,0 +1,40 @@ +# mypy: ignore-errors +# +# PEP 695 isn't yet supported in mypy. (April 4, 2024) +"""Type Aliases (Python 3.12+) + +In Python 3.12, the :code:`type` statement is introduced to create type aliases. + +Usage: +`python ./13_type_statement.py --help` +""" + +import dataclasses + +import tyro + +# Lazily-evaluated type alias. +type Field1Type = Inner + + +@dataclasses.dataclass +class Inner: + a: int + b: str + + +@dataclasses.dataclass +class Args: + """Description. + This should show up in the helptext!""" + + field1: Field1Type + """A field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + +if __name__ == "__main__": + args = tyro.cli(Args) + print(args) diff --git a/examples/04_additional/14_suppress_console_outputs.py b/examples/04_additional/14_suppress_console_outputs.py new file mode 100644 index 000000000..c8b48840a --- /dev/null +++ b/examples/04_additional/14_suppress_console_outputs.py @@ -0,0 +1,45 @@ +"""Cleaner Console Outputs for Scripts with Multiple Workers + +The :code:`console_outputs=` argument can be set to :code:`False` to suppress helptext and +error message printing. + +This is useful in PyTorch for distributed training scripts, where you only want +to print the helptext from the main process: + + +.. code-block:: python + + # HuggingFace Accelerate. + args = tyro.cli(Args, console_outputs=accelerator.is_main_process) + + # PyTorch DDP. + args = tyro.cli(Args, console_outputs=(rank == 0)) + + # PyTorch Lightning. + args = tyro.cli(Args, console_outputs=trainer.is_global_zero) + + +Usage: +`python ./14_suppress_console_outputs.py --help` +""" + +import dataclasses + +import tyro + + +@dataclasses.dataclass +class Args: + """Description. + This should show up in the helptext!""" + + field1: int + """A field.""" + + field2: int = 3 + """A numeric field, with a default value.""" + + +if __name__ == "__main__": + args = tyro.cli(Args, console_outputs=False) + print(args) diff --git a/examples/_rename_example.py b/examples/_rename_example.py new file mode 100755 index 000000000..503831c8d --- /dev/null +++ b/examples/_rename_example.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python + +import pathlib + +import tyro + + +def main(source: pathlib.Path, dest: pathlib.Path, /) -> None: + """Rename an example, while also replacing any occurrences of its old name within + its contents.""" + assert not dest.exists() + source.write_text(source.read_text().replace(str(source), str(dest))) + source.rename(dest) + + +if __name__ == "__main__": + print("starting") + tyro.cli(main) diff --git a/examples/example.py b/examples/example.py deleted file mode 100644 index c9c909bfe..000000000 --- a/examples/example.py +++ /dev/null @@ -1,43 +0,0 @@ -"""An argument parsing example. - -Note that there are multiple possible ways to document dataclass attributes, all -of which are supported by the automatic helptext generator. -""" - -import dataclasses -import enum - -import dcargs - - -class OptimizerType(enum.Enum): - ADAM = enum.auto() - SGD = enum.auto() - - -@dataclasses.dataclass -class OptimizerConfig: - # Variant of SGD to use. - type: OptimizerType - - # Learning rate to use. - learning_rate: float = 3e-4 - - # Coefficient for L2 regularization. - weight_decay: float = 1e-2 - - -@dataclasses.dataclass -class ExperimentConfig: - experiment_name: str # Experiment name to use. - - optimizer: OptimizerConfig - - seed: int = 0 - """Random seed. This is helpful for making sure that our experiments are - all reproducible!""" - - -if __name__ == "__main__": - config = dcargs.parse(ExperimentConfig, description=__doc__) - print(config) diff --git a/examples/generics.py b/examples/generics.py deleted file mode 100644 index 9fb1f675d..000000000 --- a/examples/generics.py +++ /dev/null @@ -1,37 +0,0 @@ -import dataclasses -from typing import Generic, Optional, TypeVar - -import dcargs - -ScalarType = TypeVar("ScalarType") - - -@dataclasses.dataclass -class Point3(Generic[ScalarType]): - x: ScalarType - y: ScalarType - z: ScalarType - frame_id: str - - -@dataclasses.dataclass -class Triangle(Generic[ScalarType]): - a: Point3[ScalarType] - b: Point3[ScalarType] - c: Point3[ScalarType] - - -@dataclasses.dataclass -class Args: - point_continuous: Point3[float] - point_discrete: Point3[int] - - triangle_continuous: Triangle[float] - triangle_discrete: Triangle[int] - - triangle_optional_coords: Triangle[Optional[float]] - - -if __name__ == "__main__": - args = dcargs.parse(Args) - print(args) diff --git a/examples/simple.py b/examples/simple.py deleted file mode 100644 index ae067b638..000000000 --- a/examples/simple.py +++ /dev/null @@ -1,14 +0,0 @@ -import dataclasses - -import dcargs - - -@dataclasses.dataclass -class Args: - field1: str # A string field. - field2: int # A numeric field. - - -if __name__ == "__main__": - args = dcargs.parse(Args) - print(args) diff --git a/examples/subparsers.py b/examples/subparsers.py deleted file mode 100644 index 1ba651671..000000000 --- a/examples/subparsers.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -import dataclasses -from typing import Union - -import dcargs - - -@dataclasses.dataclass -class Args: - command: Union[Checkout, Commit] - - -@dataclasses.dataclass -class Checkout: - branch: str - - -@dataclasses.dataclass -class Commit: - message: str - all: bool = False - - -if __name__ == "__main": - args = dcargs.parse(Args) - print(args) diff --git a/mypy.ini b/mypy.ini deleted file mode 100644 index 976ba0294..000000000 --- a/mypy.ini +++ /dev/null @@ -1,2 +0,0 @@ -[mypy] -ignore_missing_imports = True diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..fb6709c56 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,137 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "tyro" +authors = [ + {name = "brentyi", email = "brentyi@berkeley.edu"}, +] +version = "0.8.10" # TODO: currently needs to be synchronized manually with __init__.py. +description = "Strongly typed, zero-effort CLI interfaces" +readme = "README.md" +license = { text="MIT" } +requires-python = ">=3.7" +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent" +] +dependencies = [ + "docstring-parser>=0.16", + "typing-extensions>=4.7.0", + "backports.cached-property>=1.0.2; python_version<'3.8'", + "colorama>=0.4.0; platform_system=='Windows'", + "rich>=11.1.0", + "shtab>=1.5.6", + "eval_type_backport>=0.1.3; python_version<'3.10'", +] + +[project.optional-dependencies] +dev = [ + "PyYAML>=6.0", + "frozendict>=2.3.4", + "pytest>=7.1.2", + "pytest-cov>=3.0.0", + "omegaconf>=2.2.2", + "attrs>=21.4.0", + "torch>=1.10.0", + "pyright>=1.1.349,!=1.1.379", + "ruff>=0.1.13", + "mypy>=1.4.1", + "numpy>=1.20.0", + # As of 7/27/2023, flax install fails for Python 3.7 without pinning to an + # old version. But doing so breaks other Python versions. + "flax>=0.6.9;python_version>='3.8'", + "pydantic>=2.5.2", + "coverage[toml]>=6.5.0", + "eval_type_backport>=0.1.3", +] + +[project.urls] +"GitHub" = "https://github.com/brentyi/tyro" + +[tool.setuptools.package-data] +tyro = ["py.typed"] + +[tool.isort] +profile = "black" + +[tool.mypy] +python_version = "3.12" +ignore_missing_imports = true +warn_unused_configs = true +exclude = ["^tests/test_py311_generated/.*", "_argparse\\.py"] + +[tool.coverage.run] +omit = ["**/_argparse.py"] + +[tool.coverage.report] +exclude_lines = [ + # Have to re-enable the standard pragma + "pragma: no cover", + + # Don't compute coverage for abstract methods, properties + "@abstract", + "@abc\\.abstract", + + # or warnings + "warnings", + + # or empty function bodies + "pass", + "\\.\\.\\.", + + # or typing imports + "TYPE_CHECKING", + + # or assert statements & errors + "assert", + "raise AssertionError", + + # or anything that's not implemented + "NotImplementedError()", + + # or fallback imports + "except ImportError:", + + # or anything that's deprecated + "deprecated" +] + +[tool.ruff] +src = ["src"] # Needed to recognize first-party import location in GitHub action. +lint.select = [ + "E", # pycodestyle errors. + "F", # Pyflakes rules. + "PLC", # Pylint convention warnings. + "PLE", # Pylint errors. + "PLR", # Pylint refactor recommendations. + "PLW", # Pylint warnings. + "I" # Import sorting. +] +lint.ignore = [ + "E741", # Ambiguous variable name. (l, O, or I) + "E501", # Line too long. + "PLR2004", # Magic value used in comparison. + "PLR0915", # Too many statements. + "PLR0913", # Too many arguments. + "PLC0414", # Import alias does not rename variable. (this is used for exporting names) + "PLC1901", # Use falsey strings. + "PLR5501", # Use `elif` instead of `else if`. + "PLR0911", # Too many return statements. + "PLR0912", # Too many branches. + "PLW0603", # Global statement updates are discouraged. + "PLW2901" # For loop variable overwritten. +] +extend-exclude = ["**/_argparse.py"] + +[tool.pyright] +pythonVersion = "3.12" +ignore = ["**/_argparse.py"] diff --git a/setup.py b/setup.py deleted file mode 100644 index 113e7bd29..000000000 --- a/setup.py +++ /dev/null @@ -1,37 +0,0 @@ -from setuptools import find_packages, setup - -with open("README.md", "r") as fh: - long_description = fh.read() - -setup( - name="dcargs", - version="0.0.4", - description="Portable, reusable, strongly typed CLIs from dataclass definitions", - long_description=long_description, - long_description_content_type="text/markdown", - url="http://github.com/brentyi/dcargs", - author="brentyi", - author_email="brentyi@berkeley.edu", - license="MIT", - packages=find_packages(), - package_data={"dcargs": ["py.typed"]}, - python_requires=">=3.7", - install_requires=["typing_extensions"], - extras_require={ - "testing": [ - "pytest", - "pytest-cov", - ], - "type-checking": [ - "mypy", - ], - }, - classifiers=[ - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "License :: OSI Approved :: MIT License", - "Operating System :: OS Independent", - ], -) diff --git a/src/tyro/__init__.py b/src/tyro/__init__.py new file mode 100644 index 000000000..594d61c7f --- /dev/null +++ b/src/tyro/__init__.py @@ -0,0 +1,17 @@ +from typing import TYPE_CHECKING + +from . import conf as conf +from . import extras as extras +from ._cli import cli as cli +from ._fields import MISSING as MISSING +from ._instantiators import ( + UnsupportedTypeAnnotationError as UnsupportedTypeAnnotationError, +) + +# Deprecated interface. +if not TYPE_CHECKING: + from ._deprecated import * # noqa + + +# TODO: this should be synchronized automatically with the pyproject.toml. +__version__ = "0.8.10" diff --git a/src/tyro/_argparse.py b/src/tyro/_argparse.py new file mode 100644 index 000000000..f5dcdf0f4 --- /dev/null +++ b/src/tyro/_argparse.py @@ -0,0 +1,2672 @@ +# Mirror of Python's argparse module. +# Should be an exact copy from: +# https://github.com/python/cpython/blob/3.12/Lib/argparse.py + +# Author: Steven J. Bethard . +# New maintainer as of 29 August 2019: Raymond Hettinger + +"""Command-line parsing library + +This module is an optparse-inspired command-line parsing library that: + + - handles both optional and positional arguments + - produces highly informative usage messages + - supports parsers that dispatch to sub-parsers + +The following is a simple usage example that sums integers from the +command-line and writes the result to a file:: + + parser = argparse.ArgumentParser( + description='sum the integers at the command line') + parser.add_argument( + 'integers', metavar='int', nargs='+', type=int, + help='an integer to be summed') + parser.add_argument( + '--log', default=sys.stdout, type=argparse.FileType('w'), + help='the file where the sum should be written') + args = parser.parse_args() + args.log.write('%s' % sum(args.integers)) + args.log.close() + +The module contains the following public classes: + + - ArgumentParser -- The main entry point for command-line parsing. As the + example above shows, the add_argument() method is used to populate + the parser with actions for optional and positional arguments. Then + the parse_args() method is invoked to convert the args at the + command-line into an object with attributes. + + - ArgumentError -- The exception raised by ArgumentParser objects when + there are errors with the parser's actions. Errors raised while + parsing the command-line are caught by ArgumentParser and emitted + as command-line messages. + + - FileType -- A factory for defining types of files to be created. As the + example above shows, instances of FileType are typically passed as + the type= argument of add_argument() calls. + + - Action -- The base class for parser actions. Typically actions are + selected by passing strings like 'store_true' or 'append_const' to + the action= argument of add_argument(). However, for greater + customization of ArgumentParser actions, subclasses of Action may + be defined and passed as the action= argument. + + - HelpFormatter, RawDescriptionHelpFormatter, RawTextHelpFormatter, + ArgumentDefaultsHelpFormatter -- Formatter classes which + may be passed as the formatter_class= argument to the + ArgumentParser constructor. HelpFormatter is the default, + RawDescriptionHelpFormatter and RawTextHelpFormatter tell the parser + not to change the formatting for help text, and + ArgumentDefaultsHelpFormatter adds information about argument defaults + to the help. + +All other classes in this module are considered implementation details. +(Also note that HelpFormatter and RawDescriptionHelpFormatter are only +considered public as object names -- the API of the formatter objects is +still considered an implementation detail.) +""" + +__version__ = '1.1' +__all__ = [ + 'ArgumentParser', + 'ArgumentError', + 'ArgumentTypeError', + 'BooleanOptionalAction', + 'FileType', + 'HelpFormatter', + 'ArgumentDefaultsHelpFormatter', + 'RawDescriptionHelpFormatter', + 'RawTextHelpFormatter', + 'MetavarTypeHelpFormatter', + 'Namespace', + 'Action', + 'ONE_OR_MORE', + 'OPTIONAL', + 'PARSER', + 'REMAINDER', + 'SUPPRESS', + 'ZERO_OR_MORE', +] + + +import os as _os +import re as _re +import sys as _sys + +import warnings + +from gettext import gettext as _, ngettext + +SUPPRESS = '==SUPPRESS==' + +OPTIONAL = '?' +ZERO_OR_MORE = '*' +ONE_OR_MORE = '+' +PARSER = 'A...' +REMAINDER = '...' +_UNRECOGNIZED_ARGS_ATTR = '_unrecognized_args' + +# ============================= +# Utility functions and classes +# ============================= + +class _AttributeHolder(object): + """Abstract base class that provides __repr__. + + The __repr__ method returns a string in the format:: + ClassName(attr=name, attr=name, ...) + The attributes are determined either by a class-level attribute, + '_kwarg_names', or by inspecting the instance __dict__. + """ + + def __repr__(self): + type_name = type(self).__name__ + arg_strings = [] + star_args = {} + for arg in self._get_args(): + arg_strings.append(repr(arg)) + for name, value in self._get_kwargs(): + if name.isidentifier(): + arg_strings.append('%s=%r' % (name, value)) + else: + star_args[name] = value + if star_args: + arg_strings.append('**%s' % repr(star_args)) + return '%s(%s)' % (type_name, ', '.join(arg_strings)) + + def _get_kwargs(self): + return list(self.__dict__.items()) + + def _get_args(self): + return [] + + +def _copy_items(items): + if items is None: + return [] + # The copy module is used only in the 'append' and 'append_const' + # actions, and it is needed only when the default value isn't a list. + # Delay its import for speeding up the common case. + if type(items) is list: + return items[:] + import copy + return copy.copy(items) + + +# =============== +# Formatting Help +# =============== + + +class HelpFormatter(object): + """Formatter for generating usage messages and argument help strings. + + Only the name of this class is considered a public API. All the methods + provided by the class are considered an implementation detail. + """ + + def __init__(self, + prog, + indent_increment=2, + max_help_position=24, + width=None): + + # default setting for width + if width is None: + import shutil + width = shutil.get_terminal_size().columns + width -= 2 + + self._prog = prog + self._indent_increment = indent_increment + self._max_help_position = min(max_help_position, + max(width - 20, indent_increment * 2)) + self._width = width + + self._current_indent = 0 + self._level = 0 + self._action_max_length = 0 + + self._root_section = self._Section(self, None) + self._current_section = self._root_section + + self._whitespace_matcher = _re.compile(r'\s+', _re.ASCII) + self._long_break_matcher = _re.compile(r'\n\n\n+') + + # =============================== + # Section and indentation methods + # =============================== + def _indent(self): + self._current_indent += self._indent_increment + self._level += 1 + + def _dedent(self): + self._current_indent -= self._indent_increment + assert self._current_indent >= 0, 'Indent decreased below 0.' + self._level -= 1 + + class _Section(object): + + def __init__(self, formatter, parent, heading=None): + self.formatter = formatter + self.parent = parent + self.heading = heading + self.items = [] + + def format_help(self): + # format the indented section + if self.parent is not None: + self.formatter._indent() + join = self.formatter._join_parts + item_help = join([func(*args) for func, args in self.items]) + if self.parent is not None: + self.formatter._dedent() + + # return nothing if the section was empty + if not item_help: + return '' + + # add the heading if the section was non-empty + if self.heading is not SUPPRESS and self.heading is not None: + current_indent = self.formatter._current_indent + heading_text = _('%(heading)s:') % dict(heading=self.heading) + heading = '%*s%s\n' % (current_indent, '', heading_text) + else: + heading = '' + + # join the section-initial newline, the heading and the help + return join(['\n', heading, item_help, '\n']) + + def _add_item(self, func, args): + self._current_section.items.append((func, args)) + + # ======================== + # Message building methods + # ======================== + def start_section(self, heading): + self._indent() + section = self._Section(self, self._current_section, heading) + self._add_item(section.format_help, []) + self._current_section = section + + def end_section(self): + self._current_section = self._current_section.parent + self._dedent() + + def add_text(self, text): + if text is not SUPPRESS and text is not None: + self._add_item(self._format_text, [text]) + + def add_usage(self, usage, actions, groups, prefix=None): + if usage is not SUPPRESS: + args = usage, actions, groups, prefix + self._add_item(self._format_usage, args) + + def add_argument(self, action): + if action.help is not SUPPRESS: + + # find all invocations + get_invocation = self._format_action_invocation + invocations = [get_invocation(action)] + for subaction in self._iter_indented_subactions(action): + invocations.append(get_invocation(subaction)) + + # update the maximum item length + invocation_length = max(map(len, invocations)) + action_length = invocation_length + self._current_indent + self._action_max_length = max(self._action_max_length, + action_length) + + # add the item to the list + self._add_item(self._format_action, [action]) + + def add_arguments(self, actions): + for action in actions: + self.add_argument(action) + + # ======================= + # Help-formatting methods + # ======================= + def format_help(self): + help = self._root_section.format_help() + if help: + help = self._long_break_matcher.sub('\n\n', help) + help = help.strip('\n') + '\n' + return help + + def _join_parts(self, part_strings): + return ''.join([part + for part in part_strings + if part and part is not SUPPRESS]) + + def _format_usage(self, usage, actions, groups, prefix): + if prefix is None: + prefix = _('usage: ') + + # if usage is specified, use that + if usage is not None: + usage = usage % dict(prog=self._prog) + + # if no optionals or positionals are available, usage is just prog + elif usage is None and not actions: + usage = '%(prog)s' % dict(prog=self._prog) + + # if optionals and positionals are available, calculate usage + elif usage is None: + prog = '%(prog)s' % dict(prog=self._prog) + + # split optionals from positionals + optionals = [] + positionals = [] + for action in actions: + if action.option_strings: + optionals.append(action) + else: + positionals.append(action) + + # build full usage string + format = self._format_actions_usage + action_usage = format(optionals + positionals, groups) + usage = ' '.join([s for s in [prog, action_usage] if s]) + + # wrap the usage parts if it's too long + text_width = self._width - self._current_indent + if len(prefix) + len(usage) > text_width: + + # break usage into wrappable parts + part_regexp = ( + r'\(.*?\)+(?=\s|$)|' + r'\[.*?\]+(?=\s|$)|' + r'\S+' + ) + opt_usage = format(optionals, groups) + pos_usage = format(positionals, groups) + opt_parts = _re.findall(part_regexp, opt_usage) + pos_parts = _re.findall(part_regexp, pos_usage) + assert ' '.join(opt_parts) == opt_usage + assert ' '.join(pos_parts) == pos_usage + + # helper for wrapping lines + def get_lines(parts, indent, prefix=None): + lines = [] + line = [] + indent_length = len(indent) + if prefix is not None: + line_len = len(prefix) - 1 + else: + line_len = indent_length - 1 + for part in parts: + if line_len + 1 + len(part) > text_width and line: + lines.append(indent + ' '.join(line)) + line = [] + line_len = indent_length - 1 + line.append(part) + line_len += len(part) + 1 + if line: + lines.append(indent + ' '.join(line)) + if prefix is not None: + lines[0] = lines[0][indent_length:] + return lines + + # if prog is short, follow it with optionals or positionals + if len(prefix) + len(prog) <= 0.75 * text_width: + indent = ' ' * (len(prefix) + len(prog) + 1) + if opt_parts: + lines = get_lines([prog] + opt_parts, indent, prefix) + lines.extend(get_lines(pos_parts, indent)) + elif pos_parts: + lines = get_lines([prog] + pos_parts, indent, prefix) + else: + lines = [prog] + + # if prog is long, put it on its own line + else: + indent = ' ' * len(prefix) + parts = opt_parts + pos_parts + lines = get_lines(parts, indent) + if len(lines) > 1: + lines = [] + lines.extend(get_lines(opt_parts, indent)) + lines.extend(get_lines(pos_parts, indent)) + lines = [prog] + lines + + # join lines into usage + usage = '\n'.join(lines) + + # prefix with 'usage:' + return '%s%s\n\n' % (prefix, usage) + + def _format_actions_usage(self, actions, groups): + # find group indices and identify actions in groups + group_actions = set() + inserts = {} + for group in groups: + if not group._group_actions: + raise ValueError(f'empty group {group}') + + try: + start = actions.index(group._group_actions[0]) + except ValueError: + continue + else: + group_action_count = len(group._group_actions) + end = start + group_action_count + if actions[start:end] == group._group_actions: + + suppressed_actions_count = 0 + for action in group._group_actions: + group_actions.add(action) + if action.help is SUPPRESS: + suppressed_actions_count += 1 + + exposed_actions_count = group_action_count - suppressed_actions_count + if not exposed_actions_count: + continue + + if not group.required: + if start in inserts: + inserts[start] += ' [' + else: + inserts[start] = '[' + if end in inserts: + inserts[end] += ']' + else: + inserts[end] = ']' + elif exposed_actions_count > 1: + if start in inserts: + inserts[start] += ' (' + else: + inserts[start] = '(' + if end in inserts: + inserts[end] += ')' + else: + inserts[end] = ')' + for i in range(start + 1, end): + inserts[i] = '|' + + # collect all actions format strings + parts = [] + for i, action in enumerate(actions): + + # suppressed arguments are marked with None + # remove | separators for suppressed arguments + if action.help is SUPPRESS: + parts.append(None) + if inserts.get(i) == '|': + inserts.pop(i) + elif inserts.get(i + 1) == '|': + inserts.pop(i + 1) + + # produce all arg strings + elif not action.option_strings: + default = self._get_default_metavar_for_positional(action) + part = self._format_args(action, default) + + # if it's in a group, strip the outer [] + if action in group_actions: + if part[0] == '[' and part[-1] == ']': + part = part[1:-1] + + # add the action string to the list + parts.append(part) + + # produce the first way to invoke the option in brackets + else: + option_string = action.option_strings[0] + + # if the Optional doesn't take a value, format is: + # -s or --long + if action.nargs == 0: + part = action.format_usage() + + # if the Optional takes a value, format is: + # -s ARGS or --long ARGS + else: + default = self._get_default_metavar_for_optional(action) + args_string = self._format_args(action, default) + part = '%s %s' % (option_string, args_string) + + # make it look optional if it's not required or in a group + if not action.required and action not in group_actions: + part = '[%s]' % part + + # add the action string to the list + parts.append(part) + + # insert things at the necessary indices + for i in sorted(inserts, reverse=True): + parts[i:i] = [inserts[i]] + + # join all the action items with spaces + text = ' '.join([item for item in parts if item is not None]) + + # clean up separators for mutually exclusive groups + open = r'[\[(]' + close = r'[\])]' + text = _re.sub(r'(%s) ' % open, r'\1', text) + text = _re.sub(r' (%s)' % close, r'\1', text) + text = _re.sub(r'%s *%s' % (open, close), r'', text) + text = text.strip() + + # return the text + return text + + def _format_text(self, text): + if '%(prog)' in text: + text = text % dict(prog=self._prog) + text_width = max(self._width - self._current_indent, 11) + indent = ' ' * self._current_indent + return self._fill_text(text, text_width, indent) + '\n\n' + + def _format_action(self, action): + # determine the required width and the entry label + help_position = min(self._action_max_length + 2, + self._max_help_position) + help_width = max(self._width - help_position, 11) + action_width = help_position - self._current_indent - 2 + action_header = self._format_action_invocation(action) + + # no help; start on same line and add a final newline + if not action.help: + tup = self._current_indent, '', action_header + action_header = '%*s%s\n' % tup + + # short action name; start on the same line and pad two spaces + elif len(action_header) <= action_width: + tup = self._current_indent, '', action_width, action_header + action_header = '%*s%-*s ' % tup + indent_first = 0 + + # long action name; start on the next line + else: + tup = self._current_indent, '', action_header + action_header = '%*s%s\n' % tup + indent_first = help_position + + # collect the pieces of the action help + parts = [action_header] + + # if there was help for the action, add lines of help text + if action.help and action.help.strip(): + help_text = self._expand_help(action) + if help_text: + help_lines = self._split_lines(help_text, help_width) + parts.append('%*s%s\n' % (indent_first, '', help_lines[0])) + for line in help_lines[1:]: + parts.append('%*s%s\n' % (help_position, '', line)) + + # or add a newline if the description doesn't end with one + elif not action_header.endswith('\n'): + parts.append('\n') + + # if there are any sub-actions, add their help as well + for subaction in self._iter_indented_subactions(action): + parts.append(self._format_action(subaction)) + + # return a single string + return self._join_parts(parts) + + def _format_action_invocation(self, action): + if not action.option_strings: + default = self._get_default_metavar_for_positional(action) + metavar, = self._metavar_formatter(action, default)(1) + return metavar + + else: + parts = [] + + # if the Optional doesn't take a value, format is: + # -s, --long + if action.nargs == 0: + parts.extend(action.option_strings) + + # if the Optional takes a value, format is: + # -s ARGS, --long ARGS + else: + default = self._get_default_metavar_for_optional(action) + args_string = self._format_args(action, default) + for option_string in action.option_strings: + parts.append('%s %s' % (option_string, args_string)) + + return ', '.join(parts) + + def _metavar_formatter(self, action, default_metavar): + if action.metavar is not None: + result = action.metavar + elif action.choices is not None: + choice_strs = [str(choice) for choice in action.choices] + result = '{%s}' % ','.join(choice_strs) + else: + result = default_metavar + + def format(tuple_size): + if isinstance(result, tuple): + return result + else: + return (result, ) * tuple_size + return format + + def _format_args(self, action, default_metavar): + get_metavar = self._metavar_formatter(action, default_metavar) + if action.nargs is None: + result = '%s' % get_metavar(1) + elif action.nargs == OPTIONAL: + result = '[%s]' % get_metavar(1) + elif action.nargs == ZERO_OR_MORE: + metavar = get_metavar(1) + if len(metavar) == 2: + result = '[%s [%s ...]]' % metavar + else: + result = '[%s ...]' % metavar + elif action.nargs == ONE_OR_MORE: + result = '%s [%s ...]' % get_metavar(2) + elif action.nargs == REMAINDER: + result = '...' + elif action.nargs == PARSER: + result = '%s ...' % get_metavar(1) + elif action.nargs == SUPPRESS: + result = '' + else: + try: + formats = ['%s' for _ in range(action.nargs)] + except TypeError: + raise ValueError("invalid nargs value") from None + result = ' '.join(formats) % get_metavar(action.nargs) + return result + + def _expand_help(self, action): + params = dict(vars(action), prog=self._prog) + for name in list(params): + if params[name] is SUPPRESS: + del params[name] + for name in list(params): + if hasattr(params[name], '__name__'): + params[name] = params[name].__name__ + if params.get('choices') is not None: + choices_str = ', '.join([str(c) for c in params['choices']]) + params['choices'] = choices_str + return self._get_help_string(action) % params + + def _iter_indented_subactions(self, action): + try: + get_subactions = action._get_subactions + except AttributeError: + pass + else: + self._indent() + yield from get_subactions() + self._dedent() + + def _split_lines(self, text, width): + text = self._whitespace_matcher.sub(' ', text).strip() + # The textwrap module is used only for formatting help. + # Delay its import for speeding up the common usage of argparse. + import textwrap + return textwrap.wrap(text, width) + + def _fill_text(self, text, width, indent): + text = self._whitespace_matcher.sub(' ', text).strip() + import textwrap + return textwrap.fill(text, width, + initial_indent=indent, + subsequent_indent=indent) + + def _get_help_string(self, action): + return action.help + + def _get_default_metavar_for_optional(self, action): + return action.dest.upper() + + def _get_default_metavar_for_positional(self, action): + return action.dest + + +class RawDescriptionHelpFormatter(HelpFormatter): + """Help message formatter which retains any formatting in descriptions. + + Only the name of this class is considered a public API. All the methods + provided by the class are considered an implementation detail. + """ + + def _fill_text(self, text, width, indent): + return ''.join(indent + line for line in text.splitlines(keepends=True)) + + +class RawTextHelpFormatter(RawDescriptionHelpFormatter): + """Help message formatter which retains formatting of all help text. + + Only the name of this class is considered a public API. All the methods + provided by the class are considered an implementation detail. + """ + + def _split_lines(self, text, width): + return text.splitlines() + + +class ArgumentDefaultsHelpFormatter(HelpFormatter): + """Help message formatter which adds default values to argument help. + + Only the name of this class is considered a public API. All the methods + provided by the class are considered an implementation detail. + """ + + def _get_help_string(self, action): + """ + Add the default value to the option help message. + + ArgumentDefaultsHelpFormatter and BooleanOptionalAction when it isn't + already present. This code will do that, detecting cornercases to + prevent duplicates or cases where it wouldn't make sense to the end + user. + """ + help = action.help + if help is None: + help = '' + + if '%(default)' not in help: + if action.default is not SUPPRESS: + defaulting_nargs = [OPTIONAL, ZERO_OR_MORE] + if action.option_strings or action.nargs in defaulting_nargs: + help += _(' (default: %(default)s)') + return help + + + +class MetavarTypeHelpFormatter(HelpFormatter): + """Help message formatter which uses the argument 'type' as the default + metavar value (instead of the argument 'dest') + + Only the name of this class is considered a public API. All the methods + provided by the class are considered an implementation detail. + """ + + def _get_default_metavar_for_optional(self, action): + return action.type.__name__ + + def _get_default_metavar_for_positional(self, action): + return action.type.__name__ + + +# ===================== +# Options and Arguments +# ===================== + +def _get_action_name(argument): + if argument is None: + return None + elif argument.option_strings: + return '/'.join(argument.option_strings) + elif argument.metavar not in (None, SUPPRESS): + return argument.metavar + elif argument.dest not in (None, SUPPRESS): + return argument.dest + elif argument.choices: + return '{' + ','.join(argument.choices) + '}' + else: + return None + + +# Use exceptions from system-level argparse module. +from argparse import ArgumentError, ArgumentTypeError + +# class ArgumentError(Exception): +# """An error from creating or using an argument (optional or positional). +# +# The string value of this exception is the message, augmented with +# information about the argument that caused it. +# """ +# +# def __init__(self, argument, message): +# self.argument_name = _get_action_name(argument) +# self.message = message +# +# def __str__(self): +# if self.argument_name is None: +# format = '%(message)s' +# else: +# format = _('argument %(argument_name)s: %(message)s') +# return format % dict(message=self.message, +# argument_name=self.argument_name) +# +# +# class ArgumentTypeError(Exception): +# """An error from trying to convert a command line string to a type.""" +# pass + + +# ============== +# Action classes +# ============== + +class Action(_AttributeHolder): + """Information about how to convert command line strings to Python objects. + + Action objects are used by an ArgumentParser to represent the information + needed to parse a single argument from one or more strings from the + command line. The keyword arguments to the Action constructor are also + all attributes of Action instances. + + Keyword Arguments: + + - option_strings -- A list of command-line option strings which + should be associated with this action. + + - dest -- The name of the attribute to hold the created object(s) + + - nargs -- The number of command-line arguments that should be + consumed. By default, one argument will be consumed and a single + value will be produced. Other values include: + - N (an integer) consumes N arguments (and produces a list) + - '?' consumes zero or one arguments + - '*' consumes zero or more arguments (and produces a list) + - '+' consumes one or more arguments (and produces a list) + Note that the difference between the default and nargs=1 is that + with the default, a single value will be produced, while with + nargs=1, a list containing a single value will be produced. + + - const -- The value to be produced if the option is specified and the + option uses an action that takes no values. + + - default -- The value to be produced if the option is not specified. + + - type -- A callable that accepts a single string argument, and + returns the converted value. The standard Python types str, int, + float, and complex are useful examples of such callables. If None, + str is used. + + - choices -- A container of values that should be allowed. If not None, + after a command-line argument has been converted to the appropriate + type, an exception will be raised if it is not a member of this + collection. + + - required -- True if the action must always be specified at the + command line. This is only meaningful for optional command-line + arguments. + + - help -- The help string describing the argument. + + - metavar -- The name to be used for the option's argument with the + help string. If None, the 'dest' value will be used as the name. + """ + + def __init__(self, + option_strings, + dest, + nargs=None, + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None): + self.option_strings = option_strings + self.dest = dest + self.nargs = nargs + self.const = const + self.default = default + self.type = type + self.choices = choices + self.required = required + self.help = help + self.metavar = metavar + + def _get_kwargs(self): + names = [ + 'option_strings', + 'dest', + 'nargs', + 'const', + 'default', + 'type', + 'choices', + 'required', + 'help', + 'metavar', + ] + return [(name, getattr(self, name)) for name in names] + + def format_usage(self): + return self.option_strings[0] + + def __call__(self, parser, namespace, values, option_string=None): + raise NotImplementedError(_('.__call__() not defined')) + + +# FIXME: remove together with `BooleanOptionalAction` deprecated arguments. +_deprecated_default = object() + +class BooleanOptionalAction(Action): + def __init__(self, + option_strings, + dest, + default=None, + type=_deprecated_default, + choices=_deprecated_default, + required=False, + help=None, + metavar=_deprecated_default): + + _option_strings = [] + for option_string in option_strings: + _option_strings.append(option_string) + + if option_string.startswith('--'): + option_string = '--no-' + option_string[2:] + _option_strings.append(option_string) + + # We need `_deprecated` special value to ban explicit arguments that + # match default value. Like: + # parser.add_argument('-f', action=BooleanOptionalAction, type=int) + for field_name in ('type', 'choices', 'metavar'): + if locals()[field_name] is not _deprecated_default: + warnings._deprecated( + field_name, + "{name!r} is deprecated as of Python 3.12 and will be " + "removed in Python {remove}.", + remove=(3, 14)) + + if type is _deprecated_default: + type = None + if choices is _deprecated_default: + choices = None + if metavar is _deprecated_default: + metavar = None + + super().__init__( + option_strings=_option_strings, + dest=dest, + nargs=0, + default=default, + type=type, + choices=choices, + required=required, + help=help, + metavar=metavar) + + + def __call__(self, parser, namespace, values, option_string=None): + if option_string in self.option_strings: + setattr(namespace, self.dest, not option_string.startswith('--no-')) + + def format_usage(self): + return ' | '.join(self.option_strings) + + +class _StoreAction(Action): + + def __init__(self, + option_strings, + dest, + nargs=None, + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None): + if nargs == 0: + raise ValueError('nargs for store actions must be != 0; if you ' + 'have nothing to store, actions such as store ' + 'true or store const may be more appropriate') + if const is not None and nargs != OPTIONAL: + raise ValueError('nargs must be %r to supply const' % OPTIONAL) + super(_StoreAction, self).__init__( + option_strings=option_strings, + dest=dest, + nargs=nargs, + const=const, + default=default, + type=type, + choices=choices, + required=required, + help=help, + metavar=metavar) + + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, values) + + +class _StoreConstAction(Action): + + def __init__(self, + option_strings, + dest, + const=None, + default=None, + required=False, + help=None, + metavar=None): + super(_StoreConstAction, self).__init__( + option_strings=option_strings, + dest=dest, + nargs=0, + const=const, + default=default, + required=required, + help=help) + + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, self.const) + + +class _StoreTrueAction(_StoreConstAction): + + def __init__(self, + option_strings, + dest, + default=False, + required=False, + help=None): + super(_StoreTrueAction, self).__init__( + option_strings=option_strings, + dest=dest, + const=True, + default=default, + required=required, + help=help) + + +class _StoreFalseAction(_StoreConstAction): + + def __init__(self, + option_strings, + dest, + default=True, + required=False, + help=None): + super(_StoreFalseAction, self).__init__( + option_strings=option_strings, + dest=dest, + const=False, + default=default, + required=required, + help=help) + + +class _AppendAction(Action): + + def __init__(self, + option_strings, + dest, + nargs=None, + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None): + if nargs == 0: + raise ValueError('nargs for append actions must be != 0; if arg ' + 'strings are not supplying the value to append, ' + 'the append const action may be more appropriate') + if const is not None and nargs != OPTIONAL: + raise ValueError('nargs must be %r to supply const' % OPTIONAL) + super(_AppendAction, self).__init__( + option_strings=option_strings, + dest=dest, + nargs=nargs, + const=const, + default=default, + type=type, + choices=choices, + required=required, + help=help, + metavar=metavar) + + def __call__(self, parser, namespace, values, option_string=None): + items = getattr(namespace, self.dest, None) + items = _copy_items(items) + items.append(values) + setattr(namespace, self.dest, items) + + +class _AppendConstAction(Action): + + def __init__(self, + option_strings, + dest, + const=None, + default=None, + required=False, + help=None, + metavar=None): + super(_AppendConstAction, self).__init__( + option_strings=option_strings, + dest=dest, + nargs=0, + const=const, + default=default, + required=required, + help=help, + metavar=metavar) + + def __call__(self, parser, namespace, values, option_string=None): + items = getattr(namespace, self.dest, None) + items = _copy_items(items) + items.append(self.const) + setattr(namespace, self.dest, items) + + +class _CountAction(Action): + + def __init__(self, + option_strings, + dest, + default=None, + required=False, + help=None): + super(_CountAction, self).__init__( + option_strings=option_strings, + dest=dest, + nargs=0, + default=default, + required=required, + help=help) + + def __call__(self, parser, namespace, values, option_string=None): + count = getattr(namespace, self.dest, None) + if count is None: + count = 0 + setattr(namespace, self.dest, count + 1) + + +class _HelpAction(Action): + + def __init__(self, + option_strings, + dest=SUPPRESS, + default=SUPPRESS, + help=None): + super(_HelpAction, self).__init__( + option_strings=option_strings, + dest=dest, + default=default, + nargs=0, + help=help) + + def __call__(self, parser, namespace, values, option_string=None): + parser.print_help() + parser.exit() + + +class _VersionAction(Action): + + def __init__(self, + option_strings, + version=None, + dest=SUPPRESS, + default=SUPPRESS, + help=None): + if help is None: + help = _("show program's version number and exit") + super(_VersionAction, self).__init__( + option_strings=option_strings, + dest=dest, + default=default, + nargs=0, + help=help) + self.version = version + + def __call__(self, parser, namespace, values, option_string=None): + version = self.version + if version is None: + version = parser.version + formatter = parser._get_formatter() + formatter.add_text(version) + parser._print_message(formatter.format_help(), _sys.stdout) + parser.exit() + + +class _SubParsersAction(Action): + + class _ChoicesPseudoAction(Action): + + def __init__(self, name, aliases, help): + metavar = dest = name + if aliases: + metavar += ' (%s)' % ', '.join(aliases) + sup = super(_SubParsersAction._ChoicesPseudoAction, self) + sup.__init__(option_strings=[], dest=dest, help=help, + metavar=metavar) + + def __init__(self, + option_strings, + prog, + parser_class, + dest=SUPPRESS, + required=False, + help=None, + metavar=None): + + self._prog_prefix = prog + self._parser_class = parser_class + self._name_parser_map = {} + self._choices_actions = [] + + super(_SubParsersAction, self).__init__( + option_strings=option_strings, + dest=dest, + nargs=PARSER, + choices=self._name_parser_map, + required=required, + help=help, + metavar=metavar) + + def add_parser(self, name, **kwargs): + # set prog from the existing prefix + if kwargs.get('prog') is None: + kwargs['prog'] = '%s %s' % (self._prog_prefix, name) + + aliases = kwargs.pop('aliases', ()) + + if name in self._name_parser_map: + raise ArgumentError(self, _('conflicting subparser: %s') % name) + for alias in aliases: + if alias in self._name_parser_map: + raise ArgumentError( + self, _('conflicting subparser alias: %s') % alias) + + # create a pseudo-action to hold the choice help + if 'help' in kwargs: + help = kwargs.pop('help') + choice_action = self._ChoicesPseudoAction(name, aliases, help) + self._choices_actions.append(choice_action) + + # create the parser and add it to the map + parser = self._parser_class(**kwargs) + self._name_parser_map[name] = parser + + # make parser available under aliases also + for alias in aliases: + self._name_parser_map[alias] = parser + + return parser + + def _get_subactions(self): + return self._choices_actions + + def __call__(self, parser, namespace, values, option_string=None): + parser_name = values[0] + arg_strings = values[1:] + + # set the parser name if requested + if self.dest is not SUPPRESS: + setattr(namespace, self.dest, parser_name) + + # select the parser + try: + parser = self._name_parser_map[parser_name] + except KeyError: + args = {'parser_name': parser_name, + 'choices': ', '.join(self._name_parser_map)} + msg = _('unknown parser %(parser_name)r (choices: %(choices)s)') % args + raise ArgumentError(self, msg) + + # parse all the remaining options into the namespace + # store any unrecognized options on the object, so that the top + # level parser can decide what to do with them + + # In case this subparser defines new defaults, we parse them + # in a new namespace object and then update the original + # namespace for the relevant parts. + subnamespace, arg_strings = parser.parse_known_args(arg_strings, None) + for key, value in vars(subnamespace).items(): + setattr(namespace, key, value) + + if arg_strings: + vars(namespace).setdefault(_UNRECOGNIZED_ARGS_ATTR, []) + getattr(namespace, _UNRECOGNIZED_ARGS_ATTR).extend(arg_strings) + +class _ExtendAction(_AppendAction): + def __call__(self, parser, namespace, values, option_string=None): + items = getattr(namespace, self.dest, None) + items = _copy_items(items) + items.extend(values) + setattr(namespace, self.dest, items) + +# ============== +# Type classes +# ============== + +class FileType(object): + """Factory for creating file object types + + Instances of FileType are typically passed as type= arguments to the + ArgumentParser add_argument() method. + + Keyword Arguments: + - mode -- A string indicating how the file is to be opened. Accepts the + same values as the builtin open() function. + - bufsize -- The file's desired buffer size. Accepts the same values as + the builtin open() function. + - encoding -- The file's encoding. Accepts the same values as the + builtin open() function. + - errors -- A string indicating how encoding and decoding errors are to + be handled. Accepts the same value as the builtin open() function. + """ + + def __init__(self, mode='r', bufsize=-1, encoding=None, errors=None): + self._mode = mode + self._bufsize = bufsize + self._encoding = encoding + self._errors = errors + + def __call__(self, string): + # the special argument "-" means sys.std{in,out} + if string == '-': + if 'r' in self._mode: + return _sys.stdin.buffer if 'b' in self._mode else _sys.stdin + elif any(c in self._mode for c in 'wax'): + return _sys.stdout.buffer if 'b' in self._mode else _sys.stdout + else: + msg = _('argument "-" with mode %r') % self._mode + raise ValueError(msg) + + # all other arguments are used as file names + try: + return open(string, self._mode, self._bufsize, self._encoding, + self._errors) + except OSError as e: + args = {'filename': string, 'error': e} + message = _("can't open '%(filename)s': %(error)s") + raise ArgumentTypeError(message % args) + + def __repr__(self): + args = self._mode, self._bufsize + kwargs = [('encoding', self._encoding), ('errors', self._errors)] + args_str = ', '.join([repr(arg) for arg in args if arg != -1] + + ['%s=%r' % (kw, arg) for kw, arg in kwargs + if arg is not None]) + return '%s(%s)' % (type(self).__name__, args_str) + +# =========================== +# Optional and Positional Parsing +# =========================== + +class Namespace(_AttributeHolder): + """Simple object for storing attributes. + + Implements equality by attribute names and values, and provides a simple + string representation. + """ + + def __init__(self, **kwargs): + for name in kwargs: + setattr(self, name, kwargs[name]) + + def __eq__(self, other): + if not isinstance(other, Namespace): + return NotImplemented + return vars(self) == vars(other) + + def __contains__(self, key): + return key in self.__dict__ + + +class _ActionsContainer(object): + + def __init__(self, + description, + prefix_chars, + argument_default, + conflict_handler): + super(_ActionsContainer, self).__init__() + + self.description = description + self.argument_default = argument_default + self.prefix_chars = prefix_chars + self.conflict_handler = conflict_handler + + # set up registries + self._registries = {} + + # register actions + self.register('action', None, _StoreAction) + self.register('action', 'store', _StoreAction) + self.register('action', 'store_const', _StoreConstAction) + self.register('action', 'store_true', _StoreTrueAction) + self.register('action', 'store_false', _StoreFalseAction) + self.register('action', 'append', _AppendAction) + self.register('action', 'append_const', _AppendConstAction) + self.register('action', 'count', _CountAction) + self.register('action', 'help', _HelpAction) + self.register('action', 'version', _VersionAction) + self.register('action', 'parsers', _SubParsersAction) + self.register('action', 'extend', _ExtendAction) + + # raise an exception if the conflict handler is invalid + self._get_handler() + + # action storage + self._actions = [] + self._option_string_actions = {} + + # groups + self._action_groups = [] + self._mutually_exclusive_groups = [] + + # defaults storage + self._defaults = {} + + # determines whether an "option" looks like a negative number + self._negative_number_matcher = _re.compile(r'^-\d+$|^-\d*\.\d+$') + + # whether or not there are any optionals that look like negative + # numbers -- uses a list so it can be shared and edited + self._has_negative_number_optionals = [] + + # ==================== + # Registration methods + # ==================== + def register(self, registry_name, value, object): + registry = self._registries.setdefault(registry_name, {}) + registry[value] = object + + def _registry_get(self, registry_name, value, default=None): + return self._registries[registry_name].get(value, default) + + # ================================== + # Namespace default accessor methods + # ================================== + def set_defaults(self, **kwargs): + self._defaults.update(kwargs) + + # if these defaults match any existing arguments, replace + # the previous default on the object with the new one + for action in self._actions: + if action.dest in kwargs: + action.default = kwargs[action.dest] + + def get_default(self, dest): + for action in self._actions: + if action.dest == dest and action.default is not None: + return action.default + return self._defaults.get(dest, None) + + + # ======================= + # Adding argument actions + # ======================= + def add_argument(self, *args, **kwargs): + """ + add_argument(dest, ..., name=value, ...) + add_argument(option_string, option_string, ..., name=value, ...) + """ + + # if no positional args are supplied or only one is supplied and + # it doesn't look like an option string, parse a positional + # argument + chars = self.prefix_chars + if not args or len(args) == 1 and args[0][0] not in chars: + if args and 'dest' in kwargs: + raise ValueError('dest supplied twice for positional argument') + kwargs = self._get_positional_kwargs(*args, **kwargs) + + # otherwise, we're adding an optional argument + else: + kwargs = self._get_optional_kwargs(*args, **kwargs) + + # if no default was supplied, use the parser-level default + if 'default' not in kwargs: + dest = kwargs['dest'] + if dest in self._defaults: + kwargs['default'] = self._defaults[dest] + elif self.argument_default is not None: + kwargs['default'] = self.argument_default + + # create the action object, and add it to the parser + action_class = self._pop_action_class(kwargs) + if not callable(action_class): + raise ValueError('unknown action "%s"' % (action_class,)) + action = action_class(**kwargs) + + # raise an error if the action type is not callable + type_func = self._registry_get('type', action.type, action.type) + if not callable(type_func): + raise ValueError('%r is not callable' % (type_func,)) + + if type_func is FileType: + raise ValueError('%r is a FileType class object, instance of it' + ' must be passed' % (type_func,)) + + # raise an error if the metavar does not match the type + if hasattr(self, "_get_formatter"): + try: + self._get_formatter()._format_args(action, None) + except TypeError: + raise ValueError("length of metavar tuple does not match nargs") + + return self._add_action(action) + + def add_argument_group(self, *args, **kwargs): + group = _ArgumentGroup(self, *args, **kwargs) + self._action_groups.append(group) + return group + + def add_mutually_exclusive_group(self, **kwargs): + group = _MutuallyExclusiveGroup(self, **kwargs) + self._mutually_exclusive_groups.append(group) + return group + + def _add_action(self, action): + # resolve any conflicts + self._check_conflict(action) + + # add to actions list + self._actions.append(action) + action.container = self + + # index the action by any option strings it has + for option_string in action.option_strings: + self._option_string_actions[option_string] = action + + # set the flag if any option strings look like negative numbers + for option_string in action.option_strings: + if self._negative_number_matcher.match(option_string): + if not self._has_negative_number_optionals: + self._has_negative_number_optionals.append(True) + + # return the created action + return action + + def _remove_action(self, action): + self._actions.remove(action) + + def _add_container_actions(self, container): + # collect groups by titles + title_group_map = {} + for group in self._action_groups: + if group.title in title_group_map: + msg = _('cannot merge actions - two groups are named %r') + raise ValueError(msg % (group.title)) + title_group_map[group.title] = group + + # map each action to its group + group_map = {} + for group in container._action_groups: + + # if a group with the title exists, use that, otherwise + # create a new group matching the container's group + if group.title not in title_group_map: + title_group_map[group.title] = self.add_argument_group( + title=group.title, + description=group.description, + conflict_handler=group.conflict_handler) + + # map the actions to their new group + for action in group._group_actions: + group_map[action] = title_group_map[group.title] + + # add container's mutually exclusive groups + # NOTE: if add_mutually_exclusive_group ever gains title= and + # description= then this code will need to be expanded as above + for group in container._mutually_exclusive_groups: + mutex_group = self.add_mutually_exclusive_group( + required=group.required) + + # map the actions to their new mutex group + for action in group._group_actions: + group_map[action] = mutex_group + + # add all actions to this container or their group + for action in container._actions: + group_map.get(action, self)._add_action(action) + + def _get_positional_kwargs(self, dest, **kwargs): + # make sure required is not specified + if 'required' in kwargs: + msg = _("'required' is an invalid argument for positionals") + raise TypeError(msg) + + # mark positional arguments as required if at least one is + # always required + if kwargs.get('nargs') not in [OPTIONAL, ZERO_OR_MORE]: + kwargs['required'] = True + if kwargs.get('nargs') == ZERO_OR_MORE and 'default' not in kwargs: + kwargs['required'] = True + + # return the keyword arguments with no option strings + return dict(kwargs, dest=dest, option_strings=[]) + + def _get_optional_kwargs(self, *args, **kwargs): + # determine short and long option strings + option_strings = [] + long_option_strings = [] + for option_string in args: + # error on strings that don't start with an appropriate prefix + if not option_string[0] in self.prefix_chars: + args = {'option': option_string, + 'prefix_chars': self.prefix_chars} + msg = _('invalid option string %(option)r: ' + 'must start with a character %(prefix_chars)r') + raise ValueError(msg % args) + + # strings starting with two prefix characters are long options + option_strings.append(option_string) + if len(option_string) > 1 and option_string[1] in self.prefix_chars: + long_option_strings.append(option_string) + + # infer destination, '--foo-bar' -> 'foo_bar' and '-x' -> 'x' + dest = kwargs.pop('dest', None) + if dest is None: + if long_option_strings: + dest_option_string = long_option_strings[0] + else: + dest_option_string = option_strings[0] + dest = dest_option_string.lstrip(self.prefix_chars) + if not dest: + msg = _('dest= is required for options like %r') + raise ValueError(msg % option_string) + dest = dest.replace('-', '_') + + # return the updated keyword arguments + return dict(kwargs, dest=dest, option_strings=option_strings) + + def _pop_action_class(self, kwargs, default=None): + action = kwargs.pop('action', default) + return self._registry_get('action', action, action) + + def _get_handler(self): + # determine function from conflict handler string + handler_func_name = '_handle_conflict_%s' % self.conflict_handler + try: + return getattr(self, handler_func_name) + except AttributeError: + msg = _('invalid conflict_resolution value: %r') + raise ValueError(msg % self.conflict_handler) + + def _check_conflict(self, action): + + # find all options that conflict with this option + confl_optionals = [] + for option_string in action.option_strings: + if option_string in self._option_string_actions: + confl_optional = self._option_string_actions[option_string] + confl_optionals.append((option_string, confl_optional)) + + # resolve any conflicts + if confl_optionals: + conflict_handler = self._get_handler() + conflict_handler(action, confl_optionals) + + def _handle_conflict_error(self, action, conflicting_actions): + message = ngettext('conflicting option string: %s', + 'conflicting option strings: %s', + len(conflicting_actions)) + conflict_string = ', '.join([option_string + for option_string, action + in conflicting_actions]) + raise ArgumentError(action, message % conflict_string) + + def _handle_conflict_resolve(self, action, conflicting_actions): + + # remove all conflicting options + for option_string, action in conflicting_actions: + + # remove the conflicting option + action.option_strings.remove(option_string) + self._option_string_actions.pop(option_string, None) + + # if the option now has no option string, remove it from the + # container holding it + if not action.option_strings: + action.container._remove_action(action) + + +class _ArgumentGroup(_ActionsContainer): + + def __init__(self, container, title=None, description=None, **kwargs): + # add any missing keyword arguments by checking the container + update = kwargs.setdefault + update('conflict_handler', container.conflict_handler) + update('prefix_chars', container.prefix_chars) + update('argument_default', container.argument_default) + super_init = super(_ArgumentGroup, self).__init__ + super_init(description=description, **kwargs) + + # group attributes + self.title = title + self._group_actions = [] + + # share most attributes with the container + self._registries = container._registries + self._actions = container._actions + self._option_string_actions = container._option_string_actions + self._defaults = container._defaults + self._has_negative_number_optionals = \ + container._has_negative_number_optionals + self._mutually_exclusive_groups = container._mutually_exclusive_groups + + def _add_action(self, action): + action = super(_ArgumentGroup, self)._add_action(action) + self._group_actions.append(action) + return action + + def _remove_action(self, action): + super(_ArgumentGroup, self)._remove_action(action) + self._group_actions.remove(action) + + def add_argument_group(self, *args, **kwargs): + warnings.warn( + "Nesting argument groups is deprecated.", + category=DeprecationWarning, + stacklevel=2 + ) + return super().add_argument_group(*args, **kwargs) + + +class _MutuallyExclusiveGroup(_ArgumentGroup): + + def __init__(self, container, required=False): + super(_MutuallyExclusiveGroup, self).__init__(container) + self.required = required + self._container = container + + def _add_action(self, action): + if action.required: + msg = _('mutually exclusive arguments must be optional') + raise ValueError(msg) + action = self._container._add_action(action) + self._group_actions.append(action) + return action + + def _remove_action(self, action): + self._container._remove_action(action) + self._group_actions.remove(action) + + def add_mutually_exclusive_group(self, *args, **kwargs): + warnings.warn( + "Nesting mutually exclusive groups is deprecated.", + category=DeprecationWarning, + stacklevel=2 + ) + return super().add_mutually_exclusive_group(*args, **kwargs) + + +class ArgumentParser(_AttributeHolder, _ActionsContainer): + """Object for parsing command line strings into Python objects. + + Keyword Arguments: + - prog -- The name of the program (default: + ``os.path.basename(sys.argv[0])``) + - usage -- A usage message (default: auto-generated from arguments) + - description -- A description of what the program does + - epilog -- Text following the argument descriptions + - parents -- Parsers whose arguments should be copied into this one + - formatter_class -- HelpFormatter class for printing help messages + - prefix_chars -- Characters that prefix optional arguments + - fromfile_prefix_chars -- Characters that prefix files containing + additional arguments + - argument_default -- The default value for all arguments + - conflict_handler -- String indicating how to handle conflicts + - add_help -- Add a -h/-help option + - allow_abbrev -- Allow long options to be abbreviated unambiguously + - exit_on_error -- Determines whether or not ArgumentParser exits with + error info when an error occurs + """ + + def __init__(self, + prog=None, + usage=None, + description=None, + epilog=None, + parents=[], + formatter_class=HelpFormatter, + prefix_chars='-', + fromfile_prefix_chars=None, + argument_default=None, + conflict_handler='error', + add_help=True, + allow_abbrev=True, + exit_on_error=True): + + superinit = super(ArgumentParser, self).__init__ + superinit(description=description, + prefix_chars=prefix_chars, + argument_default=argument_default, + conflict_handler=conflict_handler) + + # default setting for prog + if prog is None: + prog = _os.path.basename(_sys.argv[0]) + + self.prog = prog + self.usage = usage + self.epilog = epilog + self.formatter_class = formatter_class + self.fromfile_prefix_chars = fromfile_prefix_chars + self.add_help = add_help + self.allow_abbrev = allow_abbrev + self.exit_on_error = exit_on_error + + add_group = self.add_argument_group + self._positionals = add_group(_('positional arguments')) + self._optionals = add_group(_('options')) + self._subparsers = None + + # register types + def identity(string): + return string + self.register('type', None, identity) + + # add help argument if necessary + # (using explicit default to override global argument_default) + default_prefix = '-' if '-' in prefix_chars else prefix_chars[0] + if self.add_help: + self.add_argument( + default_prefix+'h', default_prefix*2+'help', + action='help', default=SUPPRESS, + help=_('show this help message and exit')) + + # add parent arguments and defaults + for parent in parents: + self._add_container_actions(parent) + try: + defaults = parent._defaults + except AttributeError: + pass + else: + self._defaults.update(defaults) + + # ======================= + # Pretty __repr__ methods + # ======================= + def _get_kwargs(self): + names = [ + 'prog', + 'usage', + 'description', + 'formatter_class', + 'conflict_handler', + 'add_help', + ] + return [(name, getattr(self, name)) for name in names] + + # ================================== + # Optional/Positional adding methods + # ================================== + def add_subparsers(self, **kwargs): + if self._subparsers is not None: + self.error(_('cannot have multiple subparser arguments')) + + # add the parser class to the arguments if it's not present + kwargs.setdefault('parser_class', type(self)) + + if 'title' in kwargs or 'description' in kwargs: + title = _(kwargs.pop('title', 'subcommands')) + description = _(kwargs.pop('description', None)) + self._subparsers = self.add_argument_group(title, description) + else: + self._subparsers = self._positionals + + # prog defaults to the usage message of this parser, skipping + # optional arguments and with no "usage:" prefix + if kwargs.get('prog') is None: + formatter = self._get_formatter() + positionals = self._get_positional_actions() + groups = self._mutually_exclusive_groups + formatter.add_usage(self.usage, positionals, groups, '') + kwargs['prog'] = formatter.format_help().strip() + + # create the parsers action and add it to the positionals list + parsers_class = self._pop_action_class(kwargs, 'parsers') + action = parsers_class(option_strings=[], **kwargs) + self._subparsers._add_action(action) + + # return the created parsers action + return action + + def _add_action(self, action): + if action.option_strings: + self._optionals._add_action(action) + else: + self._positionals._add_action(action) + return action + + def _get_optional_actions(self): + return [action + for action in self._actions + if action.option_strings] + + def _get_positional_actions(self): + return [action + for action in self._actions + if not action.option_strings] + + # ===================================== + # Command line argument parsing methods + # ===================================== + def parse_args(self, args=None, namespace=None): + args, argv = self.parse_known_args(args, namespace) + if argv: + msg = _('unrecognized arguments: %s') + self.error(msg % ' '.join(argv)) + return args + + def parse_known_args(self, args=None, namespace=None): + if args is None: + # args default to the system args + args = _sys.argv[1:] + else: + # make sure that args are mutable + args = list(args) + + # default Namespace built from parser defaults + if namespace is None: + namespace = Namespace() + + # add any action defaults that aren't present + for action in self._actions: + if action.dest is not SUPPRESS: + if not hasattr(namespace, action.dest): + if action.default is not SUPPRESS: + setattr(namespace, action.dest, action.default) + + # add any parser defaults that aren't present + for dest in self._defaults: + if not hasattr(namespace, dest): + setattr(namespace, dest, self._defaults[dest]) + + # parse the arguments and exit if there are any errors + if self.exit_on_error: + try: + namespace, args = self._parse_known_args(args, namespace) + except ArgumentError as err: + self.error(str(err)) + else: + namespace, args = self._parse_known_args(args, namespace) + + if hasattr(namespace, _UNRECOGNIZED_ARGS_ATTR): + args.extend(getattr(namespace, _UNRECOGNIZED_ARGS_ATTR)) + delattr(namespace, _UNRECOGNIZED_ARGS_ATTR) + return namespace, args + + def _parse_known_args(self, arg_strings, namespace): + # replace arg strings that are file references + if self.fromfile_prefix_chars is not None: + arg_strings = self._read_args_from_files(arg_strings) + + # map all mutually exclusive arguments to the other arguments + # they can't occur with + action_conflicts = {} + for mutex_group in self._mutually_exclusive_groups: + group_actions = mutex_group._group_actions + for i, mutex_action in enumerate(mutex_group._group_actions): + conflicts = action_conflicts.setdefault(mutex_action, []) + conflicts.extend(group_actions[:i]) + conflicts.extend(group_actions[i + 1:]) + + # find all option indices, and determine the arg_string_pattern + # which has an 'O' if there is an option at an index, + # an 'A' if there is an argument, or a '-' if there is a '--' + option_string_indices = {} + arg_string_pattern_parts = [] + arg_strings_iter = iter(arg_strings) + for i, arg_string in enumerate(arg_strings_iter): + + # all args after -- are non-options + if arg_string == '--': + arg_string_pattern_parts.append('-') + for arg_string in arg_strings_iter: + arg_string_pattern_parts.append('A') + + # otherwise, add the arg to the arg strings + # and note the index if it was an option + else: + option_tuple = self._parse_optional(arg_string) + if option_tuple is None: + pattern = 'A' + else: + option_string_indices[i] = option_tuple + pattern = 'O' + arg_string_pattern_parts.append(pattern) + + # join the pieces together to form the pattern + arg_strings_pattern = ''.join(arg_string_pattern_parts) + + # converts arg strings to the appropriate and then takes the action + seen_actions = set() + seen_non_default_actions = set() + + def take_action(action, argument_strings, option_string=None): + seen_actions.add(action) + argument_values = self._get_values(action, argument_strings) + + # error if this argument is not allowed with other previously + # seen arguments, assuming that actions that use the default + # value don't really count as "present" + if argument_values is not action.default: + seen_non_default_actions.add(action) + for conflict_action in action_conflicts.get(action, []): + if conflict_action in seen_non_default_actions: + msg = _('not allowed with argument %s') + action_name = _get_action_name(conflict_action) + raise ArgumentError(action, msg % action_name) + + # take the action if we didn't receive a SUPPRESS value + # (e.g. from a default) + if argument_values is not SUPPRESS: + action(self, namespace, argument_values, option_string) + + # function to convert arg_strings into an optional action + def consume_optional(start_index): + + # get the optional identified at this index + option_tuple = option_string_indices[start_index] + action, option_string, sep, explicit_arg = option_tuple + + # identify additional optionals in the same arg string + # (e.g. -xyz is the same as -x -y -z if no args are required) + match_argument = self._match_argument + action_tuples = [] + while True: + + # if we found no optional action, skip it + if action is None: + extras.append(arg_strings[start_index]) + return start_index + 1 + + # if there is an explicit argument, try to match the + # optional's string arguments to only this + if explicit_arg is not None: + arg_count = match_argument(action, 'A') + + # if the action is a single-dash option and takes no + # arguments, try to parse more single-dash options out + # of the tail of the option string + chars = self.prefix_chars + if ( + arg_count == 0 + and option_string[1] not in chars + and explicit_arg != '' + ): + if sep or explicit_arg[0] in chars: + msg = _('ignored explicit argument %r') + raise ArgumentError(action, msg % explicit_arg) + action_tuples.append((action, [], option_string)) + char = option_string[0] + option_string = char + explicit_arg[0] + optionals_map = self._option_string_actions + if option_string in optionals_map: + action = optionals_map[option_string] + explicit_arg = explicit_arg[1:] + if not explicit_arg: + sep = explicit_arg = None + elif explicit_arg[0] == '=': + sep = '=' + explicit_arg = explicit_arg[1:] + else: + sep = '' + else: + extras.append(char + explicit_arg) + stop = start_index + 1 + break + # if the action expect exactly one argument, we've + # successfully matched the option; exit the loop + elif arg_count == 1: + stop = start_index + 1 + args = [explicit_arg] + action_tuples.append((action, args, option_string)) + break + + # error if a double-dash option did not use the + # explicit argument + else: + msg = _('ignored explicit argument %r') + raise ArgumentError(action, msg % explicit_arg) + + # if there is no explicit argument, try to match the + # optional's string arguments with the following strings + # if successful, exit the loop + else: + start = start_index + 1 + selected_patterns = arg_strings_pattern[start:] + arg_count = match_argument(action, selected_patterns) + stop = start + arg_count + args = arg_strings[start:stop] + action_tuples.append((action, args, option_string)) + break + + # add the Optional to the list and return the index at which + # the Optional's string args stopped + assert action_tuples + for action, args, option_string in action_tuples: + take_action(action, args, option_string) + return stop + + # the list of Positionals left to be parsed; this is modified + # by consume_positionals() + positionals = self._get_positional_actions() + + # function to convert arg_strings into positional actions + def consume_positionals(start_index): + # match as many Positionals as possible + match_partial = self._match_arguments_partial + selected_pattern = arg_strings_pattern[start_index:] + arg_counts = match_partial(positionals, selected_pattern) + + # slice off the appropriate arg strings for each Positional + # and add the Positional and its args to the list + for action, arg_count in zip(positionals, arg_counts): + args = arg_strings[start_index: start_index + arg_count] + start_index += arg_count + take_action(action, args) + + # slice off the Positionals that we just parsed and return the + # index at which the Positionals' string args stopped + positionals[:] = positionals[len(arg_counts):] + return start_index + + # consume Positionals and Optionals alternately, until we have + # passed the last option string + extras = [] + start_index = 0 + if option_string_indices: + max_option_string_index = max(option_string_indices) + else: + max_option_string_index = -1 + while start_index <= max_option_string_index: + + # consume any Positionals preceding the next option + next_option_string_index = min([ + index + for index in option_string_indices + if index >= start_index]) + if start_index != next_option_string_index: + positionals_end_index = consume_positionals(start_index) + + # only try to parse the next optional if we didn't consume + # the option string during the positionals parsing + if positionals_end_index > start_index: + start_index = positionals_end_index + continue + else: + start_index = positionals_end_index + + # if we consumed all the positionals we could and we're not + # at the index of an option string, there were extra arguments + if start_index not in option_string_indices: + strings = arg_strings[start_index:next_option_string_index] + extras.extend(strings) + start_index = next_option_string_index + + # consume the next optional and any arguments for it + start_index = consume_optional(start_index) + + # consume any positionals following the last Optional + stop_index = consume_positionals(start_index) + + # if we didn't consume all the argument strings, there were extras + extras.extend(arg_strings[stop_index:]) + + # make sure all required actions were present and also convert + # action defaults which were not given as arguments + required_actions = [] + for action in self._actions: + if action not in seen_actions: + if action.required: + required_actions.append(_get_action_name(action)) + else: + # Convert action default now instead of doing it before + # parsing arguments to avoid calling convert functions + # twice (which may fail) if the argument was given, but + # only if it was defined already in the namespace + if (action.default is not None and + isinstance(action.default, str) and + hasattr(namespace, action.dest) and + action.default is getattr(namespace, action.dest)): + setattr(namespace, action.dest, + self._get_value(action, action.default)) + + if required_actions: + self.error(_('the following arguments are required: %s') % + ', '.join(required_actions)) + + # make sure all required groups had one option present + for group in self._mutually_exclusive_groups: + if group.required: + for action in group._group_actions: + if action in seen_non_default_actions: + break + + # if no actions were used, report the error + else: + names = [_get_action_name(action) + for action in group._group_actions + if action.help is not SUPPRESS] + msg = _('one of the arguments %s is required') + self.error(msg % ' '.join(names)) + + # return the updated namespace and the extra arguments + return namespace, extras + + def _read_args_from_files(self, arg_strings): + # expand arguments referencing files + new_arg_strings = [] + for arg_string in arg_strings: + + # for regular arguments, just add them back into the list + if not arg_string or arg_string[0] not in self.fromfile_prefix_chars: + new_arg_strings.append(arg_string) + + # replace arguments referencing files with the file content + else: + try: + with open(arg_string[1:], + encoding=_sys.getfilesystemencoding(), + errors=_sys.getfilesystemencodeerrors()) as args_file: + arg_strings = [] + for arg_line in args_file.read().splitlines(): + for arg in self.convert_arg_line_to_args(arg_line): + arg_strings.append(arg) + arg_strings = self._read_args_from_files(arg_strings) + new_arg_strings.extend(arg_strings) + except OSError as err: + self.error(str(err)) + + # return the modified argument list + return new_arg_strings + + def convert_arg_line_to_args(self, arg_line): + return [arg_line] + + def _match_argument(self, action, arg_strings_pattern): + # match the pattern for this action to the arg strings + nargs_pattern = self._get_nargs_pattern(action) + match = _re.match(nargs_pattern, arg_strings_pattern) + + # raise an exception if we weren't able to find a match + if match is None: + nargs_errors = { + None: _('expected one argument'), + OPTIONAL: _('expected at most one argument'), + ONE_OR_MORE: _('expected at least one argument'), + } + msg = nargs_errors.get(action.nargs) + if msg is None: + msg = ngettext('expected %s argument', + 'expected %s arguments', + action.nargs) % action.nargs + raise ArgumentError(action, msg) + + # return the number of arguments matched + return len(match.group(1)) + + def _match_arguments_partial(self, actions, arg_strings_pattern): + # progressively shorten the actions list by slicing off the + # final actions until we find a match + result = [] + for i in range(len(actions), 0, -1): + actions_slice = actions[:i] + pattern = ''.join([self._get_nargs_pattern(action) + for action in actions_slice]) + match = _re.match(pattern, arg_strings_pattern) + if match is not None: + result.extend([len(string) for string in match.groups()]) + break + + # return the list of arg string counts + return result + + def _parse_optional(self, arg_string): + # if it's an empty string, it was meant to be a positional + if not arg_string: + return None + + # if it doesn't start with a prefix, it was meant to be positional + if not arg_string[0] in self.prefix_chars: + return None + + # if the option string is present in the parser, return the action + if arg_string in self._option_string_actions: + action = self._option_string_actions[arg_string] + return action, arg_string, None, None + + # if it's just a single character, it was meant to be positional + if len(arg_string) == 1: + return None + + # if the option string before the "=" is present, return the action + option_string, sep, explicit_arg = arg_string.partition('=') + if sep and option_string in self._option_string_actions: + action = self._option_string_actions[option_string] + return action, option_string, sep, explicit_arg + + # search through all possible prefixes of the option string + # and all actions in the parser for possible interpretations + option_tuples = self._get_option_tuples(arg_string) + + # if multiple actions match, the option string was ambiguous + if len(option_tuples) > 1: + options = ', '.join([option_string + for action, option_string, sep, explicit_arg in option_tuples]) + args = {'option': arg_string, 'matches': options} + msg = _('ambiguous option: %(option)s could match %(matches)s') + self.error(msg % args) + + # if exactly one action matched, this segmentation is good, + # so return the parsed action + elif len(option_tuples) == 1: + option_tuple, = option_tuples + return option_tuple + + # if it was not found as an option, but it looks like a negative + # number, it was meant to be positional + # unless there are negative-number-like options + if self._negative_number_matcher.match(arg_string): + if not self._has_negative_number_optionals: + return None + + # if it contains a space, it was meant to be a positional + if ' ' in arg_string: + return None + + # it was meant to be an optional but there is no such option + # in this parser (though it might be a valid option in a subparser) + return None, arg_string, None, None + + def _get_option_tuples(self, option_string): + result = [] + + # option strings starting with two prefix characters are only + # split at the '=' + chars = self.prefix_chars + if option_string[0] in chars and option_string[1] in chars: + if self.allow_abbrev: + option_prefix, sep, explicit_arg = option_string.partition('=') + if not sep: + sep = explicit_arg = None + for option_string in self._option_string_actions: + if option_string.startswith(option_prefix): + action = self._option_string_actions[option_string] + tup = action, option_string, sep, explicit_arg + result.append(tup) + + # single character options can be concatenated with their arguments + # but multiple character options always have to have their argument + # separate + elif option_string[0] in chars and option_string[1] not in chars: + option_prefix = option_string + short_option_prefix = option_string[:2] + short_explicit_arg = option_string[2:] + + for option_string in self._option_string_actions: + if option_string == short_option_prefix: + action = self._option_string_actions[option_string] + tup = action, option_string, '', short_explicit_arg + result.append(tup) + elif option_string.startswith(option_prefix): + action = self._option_string_actions[option_string] + tup = action, option_string, None, None + result.append(tup) + + # shouldn't ever get here + else: + self.error(_('unexpected option string: %s') % option_string) + + # return the collected option tuples + return result + + def _get_nargs_pattern(self, action): + # in all examples below, we have to allow for '--' args + # which are represented as '-' in the pattern + nargs = action.nargs + + # the default (None) is assumed to be a single argument + if nargs is None: + nargs_pattern = '(-*A-*)' + + # allow zero or one arguments + elif nargs == OPTIONAL: + nargs_pattern = '(-*A?-*)' + + # allow zero or more arguments + elif nargs == ZERO_OR_MORE: + nargs_pattern = '(-*[A-]*)' + + # allow one or more arguments + elif nargs == ONE_OR_MORE: + nargs_pattern = '(-*A[A-]*)' + + # allow any number of options or arguments + elif nargs == REMAINDER: + nargs_pattern = '([-AO]*)' + + # allow one argument followed by any number of options or arguments + elif nargs == PARSER: + nargs_pattern = '(-*A[-AO]*)' + + # suppress action, like nargs=0 + elif nargs == SUPPRESS: + nargs_pattern = '(-*-*)' + + # all others should be integers + else: + nargs_pattern = '(-*%s-*)' % '-*'.join('A' * nargs) + + # if this is an optional action, -- is not allowed + if action.option_strings: + nargs_pattern = nargs_pattern.replace('-*', '') + nargs_pattern = nargs_pattern.replace('-', '') + + # return the pattern + return nargs_pattern + + # ======================== + # Alt command line argument parsing, allowing free intermix + # ======================== + + def parse_intermixed_args(self, args=None, namespace=None): + args, argv = self.parse_known_intermixed_args(args, namespace) + if argv: + msg = _('unrecognized arguments: %s') + self.error(msg % ' '.join(argv)) + return args + + def parse_known_intermixed_args(self, args=None, namespace=None): + # returns a namespace and list of extras + # + # positional can be freely intermixed with optionals. optionals are + # first parsed with all positional arguments deactivated. The 'extras' + # are then parsed. If the parser definition is incompatible with the + # intermixed assumptions (e.g. use of REMAINDER, subparsers) a + # TypeError is raised. + # + # positionals are 'deactivated' by setting nargs and default to + # SUPPRESS. This blocks the addition of that positional to the + # namespace + + positionals = self._get_positional_actions() + a = [action for action in positionals + if action.nargs in [PARSER, REMAINDER]] + if a: + raise TypeError('parse_intermixed_args: positional arg' + ' with nargs=%s'%a[0].nargs) + + if [action.dest for group in self._mutually_exclusive_groups + for action in group._group_actions if action in positionals]: + raise TypeError('parse_intermixed_args: positional in' + ' mutuallyExclusiveGroup') + + try: + save_usage = self.usage + try: + if self.usage is None: + # capture the full usage for use in error messages + self.usage = self.format_usage()[7:] + for action in positionals: + # deactivate positionals + action.save_nargs = action.nargs + # action.nargs = 0 + action.nargs = SUPPRESS + action.save_default = action.default + action.default = SUPPRESS + namespace, remaining_args = self.parse_known_args(args, + namespace) + for action in positionals: + # remove the empty positional values from namespace + if (hasattr(namespace, action.dest) + and getattr(namespace, action.dest)==[]): + from warnings import warn + warn('Do not expect %s in %s' % (action.dest, namespace)) + delattr(namespace, action.dest) + finally: + # restore nargs and usage before exiting + for action in positionals: + action.nargs = action.save_nargs + action.default = action.save_default + optionals = self._get_optional_actions() + try: + # parse positionals. optionals aren't normally required, but + # they could be, so make sure they aren't. + for action in optionals: + action.save_required = action.required + action.required = False + for group in self._mutually_exclusive_groups: + group.save_required = group.required + group.required = False + namespace, extras = self.parse_known_args(remaining_args, + namespace) + finally: + # restore parser values before exiting + for action in optionals: + action.required = action.save_required + for group in self._mutually_exclusive_groups: + group.required = group.save_required + finally: + self.usage = save_usage + return namespace, extras + + # ======================== + # Value conversion methods + # ======================== + def _get_values(self, action, arg_strings): + # for everything but PARSER, REMAINDER args, strip out first '--' + if not action.option_strings and action.nargs not in [PARSER, REMAINDER]: + try: + arg_strings.remove('--') + except ValueError: + pass + + # optional argument produces a default when not present + if not arg_strings and action.nargs == OPTIONAL: + if action.option_strings: + value = action.const + else: + value = action.default + if isinstance(value, str): + value = self._get_value(action, value) + self._check_value(action, value) + + # when nargs='*' on a positional, if there were no command-line + # args, use the default if it is anything other than None + elif (not arg_strings and action.nargs == ZERO_OR_MORE and + not action.option_strings): + if action.default is not None: + value = action.default + self._check_value(action, value) + else: + # since arg_strings is always [] at this point + # there is no need to use self._check_value(action, value) + value = arg_strings + + # single argument or optional argument produces a single value + elif len(arg_strings) == 1 and action.nargs in [None, OPTIONAL]: + arg_string, = arg_strings + value = self._get_value(action, arg_string) + self._check_value(action, value) + + # REMAINDER arguments convert all values, checking none + elif action.nargs == REMAINDER: + value = [self._get_value(action, v) for v in arg_strings] + + # PARSER arguments convert all values, but check only the first + elif action.nargs == PARSER: + value = [self._get_value(action, v) for v in arg_strings] + self._check_value(action, value[0]) + + # SUPPRESS argument does not put anything in the namespace + elif action.nargs == SUPPRESS: + value = SUPPRESS + + # all other types of nargs produce a list + else: + value = [self._get_value(action, v) for v in arg_strings] + for v in value: + self._check_value(action, v) + + # return the converted value + return value + + def _get_value(self, action, arg_string): + type_func = self._registry_get('type', action.type, action.type) + if not callable(type_func): + msg = _('%r is not callable') + raise ArgumentError(action, msg % type_func) + + # convert the value to the appropriate type + try: + result = type_func(arg_string) + + # ArgumentTypeErrors indicate errors + except ArgumentTypeError as err: + msg = str(err) + raise ArgumentError(action, msg) + + # TypeErrors or ValueErrors also indicate errors + except (TypeError, ValueError): + name = getattr(action.type, '__name__', repr(action.type)) + args = {'type': name, 'value': arg_string} + msg = _('invalid %(type)s value: %(value)r') + raise ArgumentError(action, msg % args) + + # return the converted value + return result + + def _check_value(self, action, value): + # converted value must be one of the choices (if specified) + if action.choices is not None and value not in action.choices: + args = {'value': value, + 'choices': ', '.join(map(repr, action.choices))} + msg = _('invalid choice: %(value)r (choose from %(choices)s)') + raise ArgumentError(action, msg % args) + + # ======================= + # Help-formatting methods + # ======================= + def format_usage(self): + formatter = self._get_formatter() + formatter.add_usage(self.usage, self._actions, + self._mutually_exclusive_groups) + return formatter.format_help() + + def format_help(self): + formatter = self._get_formatter() + + # usage + formatter.add_usage(self.usage, self._actions, + self._mutually_exclusive_groups) + + # description + formatter.add_text(self.description) + + # positionals, optionals and user-defined groups + for action_group in self._action_groups: + formatter.start_section(action_group.title) + formatter.add_text(action_group.description) + formatter.add_arguments(action_group._group_actions) + formatter.end_section() + + # epilog + formatter.add_text(self.epilog) + + # determine help from format above + return formatter.format_help() + + def _get_formatter(self): + return self.formatter_class(prog=self.prog) + + # ===================== + # Help-printing methods + # ===================== + def print_usage(self, file=None): + if file is None: + file = _sys.stdout + self._print_message(self.format_usage(), file) + + def print_help(self, file=None): + if file is None: + file = _sys.stdout + self._print_message(self.format_help(), file) + + def _print_message(self, message, file=None): + if message: + file = file or _sys.stderr + try: + file.write(message) + except (AttributeError, OSError): + pass + + # =============== + # Exiting methods + # =============== + def exit(self, status=0, message=None): + if message: + self._print_message(message, _sys.stderr) + _sys.exit(status) + + def error(self, message): + """error(message: string) + + Prints a usage message incorporating the message to stderr and + exits. + + If you override this in a subclass, it should not return -- it + should either exit or raise an exception. + """ + self.print_usage(_sys.stderr) + args = {'prog': self.prog, 'message': message} + self.exit(2, _('%(prog)s: error: %(message)s\n') % args) diff --git a/src/tyro/_argparse_formatter.py b/src/tyro/_argparse_formatter.py new file mode 100644 index 000000000..69753e4f0 --- /dev/null +++ b/src/tyro/_argparse_formatter.py @@ -0,0 +1,1364 @@ +"""Utilities and functions for helptext formatting. We replace argparse's simple help +messages with ones that: + - Are more nicely formatted! + - Support multiple columns when many fields are defined. + - Use `rich` for formatting. + - Can be themed with an accent color. + +This is largely built by fussing around in argparse implementation details. It's +chaotic as a result; for stability we mirror argparse at _argparse.py. +""" + +from __future__ import annotations + +import argparse as argparse_sys +import contextlib +import dataclasses +import difflib +import itertools +import re as _re +import shlex +import shutil +import sys +from gettext import gettext as _ +from typing import Any, Dict, Generator, Iterable, List, NoReturn, Optional, Set, Tuple + +from rich.columns import Columns +from rich.console import Console, Group, RenderableType +from rich.padding import Padding +from rich.panel import Panel +from rich.rule import Rule +from rich.style import Style +from rich.table import Table +from rich.text import Text +from rich.theme import Theme +from typing_extensions import override + +from . import _argparse as argparse +from . import _arguments, _strings, conf +from ._parsers import ParserSpecification + + +@dataclasses.dataclass +class TyroTheme: + border: Style = Style() + description: Style = Style() + invocation: Style = Style() + metavar: Style = Style() + metavar_fixed: Style = Style() + helptext: Style = Style() + helptext_required: Style = Style() + helptext_default: Style = Style() + + def as_rich_theme(self) -> Theme: + return Theme(vars(self)) + + +def set_accent_color(accent_color: Optional[str]) -> None: + """Set an accent color to use in help messages. Takes any color supported by `rich`, + see `python -m rich.color`. Experimental.""" + THEME.border = Style(color=accent_color, dim=True) + THEME.description = Style(color=accent_color, bold=True) + THEME.invocation = Style() + THEME.metavar = Style(color=accent_color, bold=True) + THEME.metavar_fixed = Style(color="red", bold=True) + THEME.helptext = Style(dim=True) + THEME.helptext_required = Style(color="bright_red", bold=True) + THEME.helptext_default = Style( + color="cyan" if accent_color != "cyan" else "magenta" + # Another option: make default color match accent color. This is maybe more + # visually consistent, but harder to read. + # color=accent_color if accent_color is not None else "cyan", + # dim=accent_color is not None, + ) + + +def recursive_arg_search( + args: List[str], + parser_spec: ParserSpecification, + prog: str, + unrecognized_arguments: Set[str], +) -> Tuple[List[_ArgumentInfo], bool, bool]: + """Recursively search for arguments in a ParserSpecification. Used for error message + printing. + + Returns a list of arguments, whether the parser has subcommands or not, and -- if + `unrecognized_arguments` is passed in --- whether an unrecognized argument exists + under a different subparser. + + Args: + args: Arguments being parsed. Used for heuristics on subcommands. + parser_spec: Argument parser specification. + subcommands: Prog corresponding to parser_spec. + unrecognized_arguments: Used for same_exists return value. + """ + # Argument name => subcommands it came from. + arguments: List[_ArgumentInfo] = [] + has_subcommands = False + same_exists = False + + def _recursive_arg_search( + parser_spec: ParserSpecification, + prog: str, + subcommand_match_score: float, + ) -> None: + """Find all possible arguments that could have been passed in.""" + + # When tyro.conf.ConsolidateSubcommandArgs is turned on, arguments will + # only appear in the help message for "leaf" subparsers. + help_flag = ( + " (other subcommands) --help" + if parser_spec.consolidate_subcommand_args + and parser_spec.subparsers is not None + else " --help" + ) + for arg in parser_spec.args: + if arg.field.is_positional() or arg.lowered.is_fixed(): + # Skip positional arguments. + continue + + # Skip suppressed arguments. + if conf.Suppress in arg.field.markers or ( + conf.SuppressFixed in arg.field.markers + and conf.Fixed in arg.field.markers + ): + continue + + option_strings = (arg.lowered.name_or_flag,) + + # Handle actions, eg BooleanOptionalAction will map ("--flag",) to + # ("--flag", "--no-flag"). + if ( + arg.lowered.action is not None + # Actions are sometimes strings in Python 3.7, eg "append". + # We'll ignore these, but this kind of thing is a good reason + # for just forking argparse. + and callable(arg.lowered.action) + ): + option_strings = arg.lowered.action( + option_strings, + dest="", # dest should not matter. + ).option_strings # type: ignore + + arguments.append( + _ArgumentInfo( + # Currently doesn't handle actions well, eg boolean optional + # arguments. + option_strings, + metavar=arg.lowered.metavar, + usage_hint=prog + help_flag, + help=arg.lowered.help, + subcommand_match_score=subcommand_match_score, + ) + ) + + # An unrecognized argument. + nonlocal same_exists + if not same_exists and arg.lowered.name_or_flag in unrecognized_arguments: + same_exists = True + + if parser_spec.subparsers is not None: + nonlocal has_subcommands + has_subcommands = True + for ( + subparser_name, + subparser, + ) in parser_spec.subparsers.parser_from_name.items(): + _recursive_arg_search( + subparser, + prog + " " + subparser_name, + # Leaky (!!) heuristic for if this subcommand is matched or not. + subcommand_match_score=subcommand_match_score + + (1 if subparser_name in args else -0.001), + ) + + for child in parser_spec.child_from_prefix.values(): + _recursive_arg_search(child, prog, subcommand_match_score) + + _recursive_arg_search(parser_spec, prog, 0) + + return arguments, has_subcommands, same_exists + + +# TODO: this is a prototype; for a v1.0.0 release we should revisit whether the global +# state here is acceptable or not. +THEME = TyroTheme() +set_accent_color(None) + + +def monkeypatch_len(obj: Any) -> int: + if isinstance(obj, str): + return len(_strings.strip_ansi_sequences(obj)) + else: + return len(obj) + + +@contextlib.contextmanager +def ansi_context() -> Generator[None, None, None]: + """Context for working with ANSI codes + argparse: + - Applies a temporary monkey patch for making argparse ignore ANSI codes when + wrapping usage text. + - Enables support for Windows via colorama. + """ + + if not hasattr(argparse, "len"): + # Sketchy, but seems to work. + argparse.len = monkeypatch_len # type: ignore + try: # pragma: no cover + # Use Colorama to support coloring in Windows shells. + import colorama # type: ignore + + # Notes: + # + # (1) This context manager looks very nice and local, but under-the-hood + # does some global operations which look likely to cause unexpected + # behavior if another library relies on `colorama.init()` and + # `colorama.deinit()`. + # + # (2) SSHed into a non-Windows machine from a WinAPI terminal => this + # won't work. + # + # Fixing these issues doesn't seem worth it: it doesn't seem like there + # are low-effort solutions for either problem, and more modern terminals + # in Windows (PowerShell, MSYS2, ...) do support ANSI codes anyways. + with colorama.colorama_text(): + yield + + except ImportError: + yield + + del argparse.len # type: ignore + else: + # No-op when the context manager is nested. + yield + + +def str_from_rich( + renderable: RenderableType, width: Optional[int] = None, soft_wrap: bool = False +) -> str: + dummy_console = Console(width=width, theme=THEME.as_rich_theme()) + with dummy_console.capture() as out: + dummy_console.print(renderable, soft_wrap=soft_wrap) + return out.get().rstrip("\n") + + +@dataclasses.dataclass(frozen=True) +class _ArgumentInfo: + option_strings: Tuple[str, ...] + metavar: Optional[str] + usage_hint: str + help: Optional[str] + subcommand_match_score: float + """Priority value used when an argument is in the current subcommand tree.""" + + +# By default, unrecognized arguments won't raise an error in the case of: +# +# # We've misspelled `binary`! +# python 03_multiple_subcommands.py dataset:mnist --dataset.binayr True +# +# When there's a subcommand that follows dataset:mnist. Instead, +# --dataset.binayr is consumed and we get an error that `True` is not a valid +# subcommand. This can be really confusing when we have a lot of keyword +# arguments. +# +# Our current solution is to manually track unrecognized arguments in _parse_known_args, +# and in error() override other errors when unrecognized arguments are present. +global_unrecognized_arg_and_prog: List[Tuple[str, str]] = [] + + +# We inherit from both our local mirror of argparse and the upstream one. +# Including the latter is purely for `isinstance()`-style checks. +class TyroArgumentParser(argparse.ArgumentParser, argparse_sys.ArgumentParser): # type: ignore + _parser_specification: ParserSpecification + _parsing_known_args: bool + _console_outputs: bool + _args: List[str] + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + @override + def _check_value(self, action, value): + """We override _check_value to ignore sentinel values defined by tyro. + + This solves a choices error raised by argparse in a very specific edge case: + literals in containers as positional arguments. + """ + from ._fields import MISSING_SINGLETONS + + if value in MISSING_SINGLETONS: + return + return super()._check_value(action, value) + + @override + def _print_message(self, message, file=None): + if message and self._console_outputs: + file = file or sys.stderr + try: + file.write(message) + except (AttributeError, OSError): # pragma: no cover + pass + + @override + def _parse_known_args(self, arg_strings, namespace): # pragma: no cover + """We override _parse_known_args() to improve error messages in the presence of + subcommands. Difference is marked with ... below.""" + + # + # Reset the unused argument list in the root parser. + # Subparsers will have spaces in self.prog. + if " " not in self.prog: + global global_unrecognized_arg_and_prog + global_unrecognized_arg_and_prog = [] + # + + # replace arg strings that are file references + if self.fromfile_prefix_chars is not None: + arg_strings = self._read_args_from_files(arg_strings) + + # map all mutually exclusive arguments to the other arguments + # they can't occur with + action_conflicts = {} + for mutex_group in self._mutually_exclusive_groups: + group_actions = mutex_group._group_actions + for i, mutex_action in enumerate(mutex_group._group_actions): + conflicts = action_conflicts.setdefault(mutex_action, []) + conflicts.extend(group_actions[:i]) + conflicts.extend(group_actions[i + 1 :]) + + # find all option indices, and determine the arg_string_pattern + # which has an 'O' if there is an option at an index, + # an 'A' if there is an argument, or a '-' if there is a '--' + option_string_indices = {} + arg_string_pattern_parts = [] + arg_strings_iter = iter(arg_strings) + for i, arg_string in enumerate(arg_strings_iter): + # all args after -- are non-options + if arg_string == "--": + arg_string_pattern_parts.append("-") + for arg_string in arg_strings_iter: + arg_string_pattern_parts.append("A") + + # otherwise, add the arg to the arg strings + # and note the index if it was an option + else: + option_tuple = self._parse_optional(arg_string) + if option_tuple is None: + pattern = "A" + else: + option_string_indices[i] = option_tuple + pattern = "O" + arg_string_pattern_parts.append(pattern) + + # join the pieces together to form the pattern + arg_strings_pattern = "".join(arg_string_pattern_parts) + + # converts arg strings to the appropriate and then takes the action + seen_actions = set() + seen_non_default_actions = set() + + def take_action(action, argument_strings, option_string=None): + seen_actions.add(action) + argument_values = self._get_values(action, argument_strings) + + # error if this argument is not allowed with other previously + # seen arguments, assuming that actions that use the default + # value don't really count as "present" + if argument_values is not action.default: + seen_non_default_actions.add(action) + for conflict_action in action_conflicts.get(action, []): + if conflict_action in seen_non_default_actions: + msg = _("not allowed with argument %s") + action_name = argparse._get_action_name(conflict_action) + raise argparse.ArgumentError(action, msg % action_name) + + # take the action if we didn't receive a SUPPRESS value + # (e.g. from a default) + if argument_values is not argparse.SUPPRESS: + action(self, namespace, argument_values, option_string) + + # function to convert arg_strings into an optional action + def consume_optional(start_index): + # get the optional identified at this index + option_tuple = option_string_indices[start_index] + action, option_string, sep, explicit_arg = option_tuple + + # identify additional optionals in the same arg string + # (e.g. -xyz is the same as -x -y -z if no args are required) + match_argument = self._match_argument + action_tuples = [] + while True: + # if we found no optional action, skip it + if action is None: + # + # Manually track unused arguments to assist with error messages + # later. + if not self._parsing_known_args: + global_unrecognized_arg_and_prog.append( + (option_string, self.prog) + ) + # + extras.append(arg_strings[start_index]) + return start_index + 1 + + # if there is an explicit argument, try to match the + # optional's string arguments to only this + if explicit_arg is not None: + arg_count = match_argument(action, "A") + + # if the action is a single-dash option and takes no + # arguments, try to parse more single-dash options out + # of the tail of the option string + chars = self.prefix_chars + if ( + arg_count == 0 + and option_string[1] not in chars + and explicit_arg != "" + ): + if sep or explicit_arg[0] in chars: + msg = _("ignored explicit argument %r") + raise argparse.ArgumentError(action, msg % explicit_arg) + action_tuples.append((action, [], option_string)) + char = option_string[0] + option_string = char + explicit_arg[0] + optionals_map = self._option_string_actions + if option_string in optionals_map: + action = optionals_map[option_string] + explicit_arg = explicit_arg[1:] + if not explicit_arg: + sep = explicit_arg = None + elif explicit_arg[0] == "=": + sep = "=" + explicit_arg = explicit_arg[1:] + else: + sep = "" + else: + extras.append(char + explicit_arg) + stop = start_index + 1 + break + # if the action expect exactly one argument, we've + # successfully matched the option; exit the loop + elif arg_count == 1: + stop = start_index + 1 + args = [explicit_arg] + action_tuples.append((action, args, option_string)) + break + + # error if a double-dash option did not use the + # explicit argument + else: + msg = _("ignored explicit argument %r") + raise argparse.ArgumentError(action, msg % explicit_arg) + + # if there is no explicit argument, try to match the + # optional's string arguments with the following strings + # if successful, exit the loop + else: + start = start_index + 1 + selected_patterns = arg_strings_pattern[start:] + arg_count = match_argument(action, selected_patterns) + stop = start + arg_count + args = arg_strings[start:stop] + action_tuples.append((action, args, option_string)) + break + + # add the Optional to the list and return the index at which + # the Optional's string args stopped + assert action_tuples + for action, args, option_string in action_tuples: + take_action(action, args, option_string) + return stop + + # the list of Positionals left to be parsed; this is modified + # by consume_positionals() + positionals = self._get_positional_actions() + + # function to convert arg_strings into positional actions + def consume_positionals(start_index): + # match as many Positionals as possible + match_partial = self._match_arguments_partial + selected_pattern = arg_strings_pattern[start_index:] + arg_counts = match_partial(positionals, selected_pattern) + + # slice off the appropriate arg strings for each Positional + # and add the Positional and its args to the list + for action, arg_count in zip(positionals, arg_counts): + args = arg_strings[start_index : start_index + arg_count] + start_index += arg_count + take_action(action, args) + + # slice off the Positionals that we just parsed and return the + # index at which the Positionals' string args stopped + positionals[:] = positionals[len(arg_counts) :] + return start_index + + # consume Positionals and Optionals alternately, until we have + # passed the last option string + extras = [] + start_index = 0 + if option_string_indices: + max_option_string_index = max(option_string_indices) + else: + max_option_string_index = -1 + while start_index <= max_option_string_index: + # consume any Positionals preceding the next option + next_option_string_index = min( + [index for index in option_string_indices if index >= start_index] + ) + if start_index != next_option_string_index: + positionals_end_index = consume_positionals(start_index) + + # only try to parse the next optional if we didn't consume + # the option string during the positionals parsing + if positionals_end_index > start_index: + start_index = positionals_end_index + continue + else: + start_index = positionals_end_index + + # if we consumed all the positionals we could and we're not + # at the index of an option string, there were extra arguments + if start_index not in option_string_indices: + strings = arg_strings[start_index:next_option_string_index] + extras.extend(strings) + start_index = next_option_string_index + + # consume the next optional and any arguments for it + start_index = consume_optional(start_index) + + # consume any positionals following the last Optional + stop_index = consume_positionals(start_index) + + # if we didn't consume all the argument strings, there were extras + extras.extend(arg_strings[stop_index:]) + + # make sure all required actions were present and also convert + # action defaults which were not given as arguments + required_actions = [] + for action in self._actions: + if action not in seen_actions: + if action.required: + required_actions.append(argparse._get_action_name(action)) + else: + # Convert action default now instead of doing it before + # parsing arguments to avoid calling convert functions + # twice (which may fail) if the argument was given, but + # only if it was defined already in the namespace + if ( + action.default is not None + and isinstance(action.default, str) + and hasattr(namespace, action.dest) + and action.default is getattr(namespace, action.dest) + ): + setattr( + namespace, + action.dest, + self._get_value(action, action.default), + ) + + if required_actions: + self.error( + _("the following arguments are required: %s") + % ", ".join(required_actions) + ) + + # make sure all required groups had one option present + for group in self._mutually_exclusive_groups: + if group.required: + for action in group._group_actions: + if action in seen_non_default_actions: + break + + # if no actions were used, report the error + else: + names = [ + argparse._get_action_name(action) + for action in group._group_actions + if action.help is not argparse.SUPPRESS + ] + msg = _("one of the arguments %s is required") + self.error(msg % " ".join(names)) # type: ignore + + # return the updated namespace and the extra arguments + return namespace, extras + + @override + def error(self, message: str) -> NoReturn: + """Improve error messages from argparse. + + error(message: string) + + Prints a usage message incorporating the message to stderr and + exits. + + If you override this in a subclass, it should not return -- it + should either exit or raise an exception. + """ + + extra_info: List[RenderableType] = [] + message_title = "Parsing error" + + if len(global_unrecognized_arg_and_prog) > 0: + message_title = "Unrecognized options" + message = f"Unrecognized options: {' '.join([arg for arg, _ in global_unrecognized_arg_and_prog])}" + unrecognized_arguments = set( + arg + for arg, _ in global_unrecognized_arg_and_prog + # If we pass in `--spell-chekc on`, we only want `spell-chekc` and not + # `on`. + if arg.startswith("--") + ) + arguments, has_subcommands, same_exists = recursive_arg_search( + args=self._args, + parser_spec=self._parser_specification, + prog=self.prog.partition(" ")[0], + unrecognized_arguments=unrecognized_arguments, + ) + + if has_subcommands and same_exists: + message = "Unrecognized or misplaced options:\n\n" + for arg, prog in global_unrecognized_arg_and_prog: + message += f" {arg} (applied to [green]{prog}[/green])\n" + message += "\nArguments are applied to the directly preceding subcommand, so ordering matters." + + # Show similar arguments for keyword options. + for unrecognized_argument in unrecognized_arguments: + # Sort arguments by similarity. + scored_arguments: List[Tuple[_ArgumentInfo, float]] = [] + for arg_info in arguments: + # Compute a score for each argument. + assert unrecognized_argument.startswith("--") + + def get_score(option_string: str) -> float: + if option_string.endswith( + unrecognized_argument[2:] + ) or option_string.startswith(unrecognized_argument[2:]): + return 0.9 + elif len(unrecognized_argument) >= 4 and all( + map( + lambda part: part in option_string, + unrecognized_argument[2:].split("."), + ) + ): + return 0.9 + else: + return difflib.SequenceMatcher( + a=unrecognized_argument, b=option_string + ).ratio() + + scored_arguments.append( + (arg_info, max(map(get_score, arg_info.option_strings))) + ) + + # Add information about similar arguments. + prev_arg_option_strings: Optional[Tuple[str, ...]] = None + show_arguments: List[_ArgumentInfo] = [] + unique_counter = 0 + for arg_info, score in ( + # Sort scores greatest to least. + sorted( + scored_arguments, + key=lambda arg_score: ( + # Highest scores first. + -arg_score[1], + # Prefer arguments available in the currently specified + # subcommands. + -arg_score[0].subcommand_match_score, + # Cluster by flag name, usage hint, help message. + arg_score[0].option_strings[0], + arg_score[0].usage_hint, + arg_score[0].help, + ), + ) + ): + if score < 0.8: + break + if ( + score < 0.9 + and unique_counter >= 3 + and prev_arg_option_strings != arg_info.option_strings + ): + break + unique_counter += prev_arg_option_strings != arg_info.option_strings + + show_arguments.append(arg_info) + prev_arg_option_strings = arg_info.option_strings + + prev_arg_option_strings = None + prev_argument_help: Optional[str] = None + same_counter = 0 + dots_printed = False + if len(show_arguments) > 0: + # Add a header before the first similar argument. + extra_info.append(Rule(style=Style(color="red"))) + extra_info.append( + "Perhaps you meant:" + if len(unrecognized_arguments) == 1 + else f"Arguments similar to {unrecognized_argument}:" + ) + + unique_counter = 0 + for arg_info in show_arguments: + same_counter += 1 + if arg_info.option_strings != prev_arg_option_strings: + same_counter = 0 + if unique_counter >= 10: + break + unique_counter += 1 + + # For arguments with the same name, only show a limited number of + # subcommands / help messages. + if ( + len(show_arguments) >= 8 + and same_counter >= 4 + and arg_info.option_strings == prev_arg_option_strings + ): + if not dots_printed: + extra_info.append( + Padding( + "[...]", + (0, 0, 0, 12), + ) + ) + dots_printed = True + continue + + if not ( + has_subcommands + and arg_info.option_strings == prev_arg_option_strings + ): + extra_info.append( + Padding( + "[bold]" + + ( + ", ".join(arg_info.option_strings) + if arg_info.metavar is None + else ", ".join(arg_info.option_strings) + + " " + + arg_info.metavar + ) + + "[/bold]", + (0, 0, 0, 4), + ) + ) + + # Uncomment to show similarity metric. + # extra_info.append( + # Padding( + # f"[green]Similarity: {score:.02f}[/green]", (0, 0, 0, 8) + # ) + # ) + + if arg_info.help is not None and ( + # Only print help messages if it's not the same as the previous + # one. + arg_info.help != prev_argument_help + or arg_info.option_strings != prev_arg_option_strings + ): + extra_info.append(Padding(arg_info.help, (0, 0, 0, 8))) + + # Show the subcommand that this argument is available in. + if has_subcommands: + extra_info.append( + Padding( + f"in [green]{arg_info.usage_hint}[/green]", + (0, 0, 0, 12), + ) + ) + + prev_arg_option_strings = arg_info.option_strings + prev_argument_help = arg_info.help + + elif message.startswith("the following arguments are required:"): + message_title = "Required options" + + info_from_required_arg: Dict[str, Optional[_ArgumentInfo]] = {} + for arg in message.partition(":")[2].strip().split(", "): + info_from_required_arg[arg] = None + + arguments, has_subcommands, same_exists = recursive_arg_search( + args=self._args, + parser_spec=self._parser_specification, + prog=self.prog.partition(" ")[0], + unrecognized_arguments=set(), + ) + del same_exists + + for arg_info in arguments: + # Iterate over each option string separately. This can help us support + # aliases in the future. + for option_string in arg_info.option_strings: + # If the option string was found... + if option_string in info_from_required_arg and ( + # And it's the first time it was found... + info_from_required_arg[option_string] is None + # Or we found a better one... + or arg_info.subcommand_match_score + > info_from_required_arg[option_string].subcommand_match_score # type: ignore + ): + # Record the argument info. + info_from_required_arg[option_string] = arg_info + + # Try to print help text for required arguments. + first = True + for maybe_arg in info_from_required_arg.values(): + if maybe_arg is None: + # No argument info found. This will currently happen for + # subcommands. + continue + + if first: + extra_info.extend( + [ + Rule(style=Style(color="red")), + "Argument helptext:", + ] + ) + first = False + + extra_info.append( + Padding( + "[bold]" + + ( + ", ".join(maybe_arg.option_strings) + if maybe_arg.metavar is None + else ", ".join(maybe_arg.option_strings) + + " " + + maybe_arg.metavar + ) + + "[/bold]", + (0, 0, 0, 4), + ) + ) + if maybe_arg.help is not None: + extra_info.append(Padding(maybe_arg.help, (0, 0, 0, 8))) + if has_subcommands: + # We are explicit about where the argument helptext is being + # extracted from because the `subcommand_match_score` heuristic + # above is flawed. + # + # The stars really need to be aligned for it to fail, but this makes + # sure that if it does fail that it's obvious to the user. + extra_info.append( + Padding( + f"in [green]{maybe_arg.usage_hint}[/green]", + (0, 0, 0, 12), + ) + ) + + if self._console_outputs: + console = Console(theme=THEME.as_rich_theme(), stderr=True) + console.print( + Panel( + Group( + f"{message[0].upper() + message[1:]}" + if len(message) > 0 + else "", + *extra_info, + Rule(style=Style(color="red")), + f"For full helptext, run [bold]{self.prog} --help[/bold]", + ), + title=f"[bold]{message_title}[/bold]", + title_align="left", + border_style=Style(color="bright_red"), + expand=False, + ) + ) + sys.exit(2) + + +class TyroArgparseHelpFormatter(argparse.RawDescriptionHelpFormatter): + def __init__(self, prog: str): + indent_increment = 4 + width = shutil.get_terminal_size().columns - 2 + max_help_position = 24 + self._fixed_help_position = False + + # TODO: hacky. Refactor this. + self._strip_ansi_sequences = not _arguments.USE_RICH + + super().__init__(prog, indent_increment, max_help_position, width) + + @override + def _format_args(self, action, default_metavar): + """Override _format_args() to ignore nargs and always expect single string + metavars.""" + get_metavar = self._metavar_formatter(action, default_metavar) + + out = get_metavar(1)[0] + if isinstance(out, str): + # Can result in an failed argparse assertion if we turn off soft wrapping. + return ( + out + if self._strip_ansi_sequences + else str_from_rich( + Text.from_ansi( + out, + style=( + THEME.metavar_fixed if out == "{fixed}" else THEME.metavar + ), + ), + soft_wrap=True, + ) + ) + return out + + @override + def add_argument(self, action): # pragma: no cover + # Patch to avoid super long arguments from shifting the helptext of all of the + # fields. + prev_max_length = self._action_max_length + super().add_argument(action) + if self._action_max_length > self._max_help_position + 2: + self._action_max_length = prev_max_length + + def _split_lines(self, text, width): # pragma: no cover + text = self._whitespace_matcher.sub(" ", text).strip() + # The textwrap module is used only for formatting help. + # Delay its import for speeding up the common usage of argparse. + import textwrap as textwrap + + # Sketchy, but seems to work. + textwrap.len = monkeypatch_len # type: ignore + out = textwrap.wrap(text, width) + del textwrap.len # type: ignore + return out + + @override + def _fill_text(self, text, width, indent): + return "".join(indent + line for line in text.splitlines(keepends=True)) + + @override + def format_help(self): + # Try with and without a fixed help position, then return the shorter help + # message. + # For dense multi-column layouts, the fixed help position is often shorter. + # For wider layouts, using the default help position settings can be more + # efficient. + self._tyro_rule = None + self._fixed_help_position = False + help1 = super().format_help() + + self._tyro_rule = None + self._fixed_help_position = True + help2 = super().format_help() + + out = help1 if help1.count("\n") < help2.count("\n") else help2 + + if self._strip_ansi_sequences: + return _strings.strip_ansi_sequences(out) + else: + return out + + @override + class _Section(object): # type: ignore + def __init__(self, formatter, parent, heading=None): + self.formatter = formatter + self.parent = parent + self.heading = heading + self.items = [] + self.formatter._tyro_rule = None + + def format_help(self): + if self.parent is None: + return self._tyro_format_root() + else: + return self._tyro_format_nonroot() + + def _tyro_format_root(self): + dummy_console = Console( + width=self.formatter._width, theme=THEME.as_rich_theme() + ) + with dummy_console.capture() as capture: + # Get rich renderables from items. + top_parts = [] + column_parts = [] + column_parts_lines = [] + for func, args in self.items: + item_content = func(*args) + if item_content is None: + pass + + # Add strings. (usage, description, etc) + elif isinstance(item_content, str): + if item_content.strip() == "": + continue + top_parts.append(Text.from_ansi(item_content)) + + # Add panels. (argument groups, subcommands, etc) + else: + assert isinstance(item_content, Panel) + column_parts.append(item_content) + # Estimate line count. This won't correctly account for + # wrapping, as we don't know the column layout yet. + column_parts_lines.append( + str_from_rich(item_content, width=65).strip().count("\n") + + 1 + ) + + # Split into columns. + min_column_width = 65 + height_breakpoint = 50 + column_count = max( + 1, + min( + sum(column_parts_lines) // height_breakpoint + 1, + self.formatter._width // min_column_width, + len(column_parts), + ), + ) + if column_count > 1: # pragma: no cover + column_width = self.formatter._width // column_count - 1 + # Correct the line count for each panel using the known column + # width. This will account for word wrap. + column_parts_lines = map( + lambda p: str_from_rich(p, width=column_width) + .strip() + .count("\n") + + 1, + column_parts, + ) + else: + column_width = None + + column_lines = [0 for i in range(column_count)] + column_parts_grouped = [[] for i in range(column_count)] + for p, l in zip(column_parts, column_parts_lines): + chosen_column = column_lines.index(min(column_lines)) + column_parts_grouped[chosen_column].append(p) + column_lines[chosen_column] += l + columns = Columns( + [Group(*g) for g in column_parts_grouped], + column_first=True, + width=column_width, + ) + + dummy_console.print(Group(*top_parts)) + dummy_console.print(columns) + return capture.get() + + def _format_action(self, action: argparse.Action): + invocation = self.formatter._format_action_invocation(action) + indent = self.formatter._current_indent + help_position = min( + self.formatter._action_max_length + 4, + self.formatter._max_help_position, + ) + if self.formatter._fixed_help_position: + help_position = 4 + + item_parts: List[RenderableType] = [] + + # Put invocation and help side-by-side. + if action.option_strings == ["-h", "--help"]: + # Darken helptext for --help flag. This makes it visually consistent + # with the helptext strings defined via docstrings and set by + # _arguments.py. + assert action.help is not None + action.help = str_from_rich( + Text.from_markup("[helptext]" + action.help + "[/helptext]") + ) + + # Unescape % signs, which need special handling in argparse. + if action.help is not None: + assert isinstance(action.help, str) + helptext = ( + Text.from_ansi(action.help.replace("%%", "%")) + if _strings.strip_ansi_sequences(action.help) != action.help + else Text.from_markup(action.help.replace("%%", "%")) + ) + else: + helptext = Text("") + + if ( + action.help + and len(_strings.strip_ansi_sequences(invocation)) + indent + < help_position - 1 + and not self.formatter._fixed_help_position + ): + table = Table(show_header=False, box=None, padding=0) + table.add_column(width=help_position - indent) + table.add_column() + table.add_row( + Text.from_ansi( + invocation, + style=THEME.invocation, + ), + helptext, + ) + item_parts.append(table) + + # Put invocation and help on separate lines. + else: + item_parts.append( + Text.from_ansi( + invocation + "\n", + style=THEME.invocation, + ) + ) + if action.help: + item_parts.append( + Padding( + # Unescape % signs, which need special handling in argparse. + helptext, + pad=(0, 0, 0, help_position - indent), + ) + ) + + # Add subactions, indented. + try: + subaction: argparse.Action + for subaction in action._get_subactions(): # type: ignore + self.formatter._indent() + item_parts.append( + Padding( + Group(*self._format_action(subaction)), + pad=(0, 0, 0, self.formatter._indent_increment), + ) + ) + self.formatter._dedent() + except AttributeError: + pass + + return item_parts + + def _tyro_format_nonroot(self): + # Add each child item as a rich renderable. + description_part = None + item_parts = [] + for func, args in self.items: + if ( + getattr(func, "__func__", None) + is TyroArgparseHelpFormatter._format_action + ): + (action,) = args + assert isinstance(action, argparse.Action) + item_parts.extend(self._format_action(action)) + + else: + item_content = func(*args) + assert isinstance(item_content, str) + if item_content.strip() != "": + assert ( + description_part is None + ) # Should only have one description part. + description_part = Text.from_ansi( + item_content.strip() + "\n", + style=THEME.description, + ) + + if len(item_parts) == 0: + return None + + # Get heading. + if self.heading is not argparse.SUPPRESS and self.heading is not None: + current_indent = self.formatter._current_indent + heading = "%*s%s:\n" % (current_indent, "", self.heading) + # Remove colon from heading. + heading = heading.strip()[:-1] + else: + heading = "" + + # Determine width for divider below description text. This is shared across + # all sections in a particular formatter. + lines = list( + itertools.chain( + *map( + lambda p: _strings.strip_ansi_sequences( + str_from_rich( + p, width=self.formatter._width, soft_wrap=True + ) + ) + .rstrip() + .split("\n"), + ( + item_parts + [description_part] + if description_part is not None + else item_parts + ), + ) + ) + ) + max_width = max(map(len, lines)) + + if self.formatter._tyro_rule is None: + # We don't use rich.rule.Rule() because this will make all of the panels + # expand to fill the full width of the console. This only impacts + # single-column layouts. + self.formatter._tyro_rule = Text.from_ansi( + "─" * max_width, style=THEME.border, overflow="crop" + ) + elif len(self.formatter._tyro_rule._text[0]) < max_width: + self.formatter._tyro_rule._text = ["─" * max_width] + + # Add description text if needed. + if description_part is not None: + item_parts = [ + description_part, + self.formatter._tyro_rule, + ] + item_parts + + return Panel( + Group(*item_parts), + title=heading, + title_align="left", + border_style=THEME.border, + # padding=(1, 1, 0, 1), + ) + + def _format_actions_usage(self, actions, groups): # pragma: no cover + """Backporting from Python 3.10, primarily to call format_usage() on actions.""" + + # find group indices and identify actions in groups + group_actions = set() + inserts = {} + for group in groups: + if not group._group_actions: + raise ValueError(f"empty group {group}") + + try: + start = actions.index(group._group_actions[0]) # type: ignore + except ValueError: + continue + else: + group_action_count = len(group._group_actions) + end = start + group_action_count + if actions[start:end] == group._group_actions: # type: ignore + suppressed_actions_count = 0 + for action in group._group_actions: + group_actions.add(action) + if action.help is argparse.SUPPRESS: + suppressed_actions_count += 1 + + exposed_actions_count = ( + group_action_count - suppressed_actions_count + ) + + if not group.required: # type: ignore + if start in inserts: + inserts[start] += " [" + else: + inserts[start] = "[" + if end in inserts: + inserts[end] += "]" + else: + inserts[end] = "]" + elif exposed_actions_count > 1: + if start in inserts: + inserts[start] += " (" + else: + inserts[start] = "(" + if end in inserts: + inserts[end] += ")" + else: + inserts[end] = ")" + for i in range(start + 1, end): + inserts[i] = "|" + + # collect all actions format strings + parts = [] + for i, action in enumerate(actions): + # suppressed arguments are marked with None + # remove | separators for suppressed arguments + if action.help is argparse.SUPPRESS: + parts.append(None) + if inserts.get(i) == "|": + inserts.pop(i) + elif inserts.get(i + 1) == "|": + inserts.pop(i + 1) + + # produce all arg strings + elif not action.option_strings: + default = self._get_default_metavar_for_positional(action) + part = self._format_args(action, default) + + # if it's in a group, strip the outer [] + if action in group_actions: + if part[0] == "[" and part[-1] == "]": + part = part[1:-1] + + # add the action string to the list + parts.append(part) + + # produce the first way to invoke the option in brackets + else: + option_string = action.option_strings[0] + + # if the Optional doesn't take a value, format is: + # -s or --long + if action.nargs == 0: + part = ( + action.format_usage() + if hasattr(action, "format_usage") + else "%s" % option_string + ) + + # if the Optional takes a value, format is: + # -s ARGS or --long ARGS + else: + default = self._get_default_metavar_for_optional(action) + args_string = self._format_args(action, default) + part = "%s %s" % (option_string, args_string) + + # make it look optional if it's not required or in a group + if not action.required and action not in group_actions: + part = "[%s]" % part + + # add the action string to the list + parts.append(part) + + # insert things at the necessary indices + for i in sorted(inserts, reverse=True): + parts[i:i] = [inserts[i]] + + # join all the action items with spaces + text = " ".join([item for item in parts if item is not None]) + + # clean up separators for mutually exclusive groups + open = r"[\[(]" + close = r"[\])]" + text = _re.sub(r"(%s) " % open, r"\1", text) + text = _re.sub(r" (%s)" % close, r"\1", text) + text = _re.sub(r"%s *%s" % (open, close), r"", text) + text = text.strip() + + # return the text + return text + + @override + def _format_usage( + self, usage, actions: Iterable[argparse.Action], groups, prefix + ) -> str: + assert isinstance(actions, list) + if len(actions) > 4: + new_actions = [] + prog_parts = shlex.split(self._prog) + added_options = False + for action in actions: + if action.dest == "help" or len(action.option_strings) == 0: + new_actions.append(action) + elif not added_options: + added_options = True + new_actions.append( + argparse.Action( + [ + "OPTIONS" + if len(prog_parts) == 1 + else prog_parts[-1].upper() + " OPTIONS" + ], + dest="", + ) + ) + actions = new_actions + + # Format the usage label. + if prefix is None: + prefix = str_from_rich("[bold]usage[/bold]: ") + usage = super()._format_usage( + usage, + actions, + groups, + prefix, + ) + return "\n\n" + usage diff --git a/src/tyro/_arguments.py b/src/tyro/_arguments.py new file mode 100644 index 000000000..2779131c2 --- /dev/null +++ b/src/tyro/_arguments.py @@ -0,0 +1,589 @@ +"""Rules for taking high-level field definitions and lowering them into inputs for +argparse's `add_argument()`.""" + +from __future__ import annotations + +import dataclasses +import enum +import itertools +import json +import shlex +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterable, + Mapping, + Optional, + Sequence, + Tuple, + TypeVar, + Union, + cast, +) + +import rich.markup +import shtab + +from . import _argparse as argparse +from . import _fields, _instantiators, _resolver, _strings +from ._typing import TypeForm +from .conf import _markers + +if TYPE_CHECKING: + cached_property = property +else: + try: + # Python >=3.8. + from functools import cached_property + except ImportError: + # Python 3.7. + from backports.cached_property import cached_property # type: ignore + + +_T = TypeVar("_T") + + +class BooleanOptionalAction(argparse.Action): + """Adapted from https://github.com/python/cpython/pull/27672""" + + def __init__( + self, + option_strings: Sequence[str], + dest: str, + default: _T | str | None = None, + type: Callable[[str], _T] | argparse.FileType | None = None, + choices: Iterable[_T] | None = None, + required: bool = False, + help: str | None = None, + metavar: str | tuple[str, ...] | None = None, + ) -> None: + _option_strings = [] + self._no_strings = set() + for option_string in option_strings: + _option_strings.append(option_string) + + if option_string.startswith("--"): + if "." not in option_string: + option_string = ( + "--no" + _strings.get_delimeter() + option_string[2:] + ) + else: + # Loose heuristic for where to add the no-/no_ prefix. + left, _, right = option_string.rpartition(".") + option_string = left + ".no" + _strings.get_delimeter() + right + self._no_strings.add(option_string) + + _option_strings.append(option_string) + + super().__init__( + option_strings=_option_strings, + dest=dest, + nargs=0, + default=default, + type=type, + choices=choices, + required=required, + help=help, + metavar=metavar, + ) + + def __call__(self, parser, namespace, values, option_string=None): + if option_string in self.option_strings: + assert option_string is not None + setattr(namespace, self.dest, option_string not in self._no_strings) + + # Typically only supported in Python 3.10, but we backport some functionality in + # _argparse_formatters.py + def format_usage(self): + return " | ".join(self.option_strings) + + +@dataclasses.dataclass(frozen=True) +class ArgumentDefinition: + """Structure containing everything needed to define an argument.""" + + intern_prefix: str # True prefix. (eg for the argument's dest field) + extern_prefix: str # User-facing prefix. + subcommand_prefix: str # Prefix for nesting. + field: _fields.FieldDefinition + type_from_typevar: Dict[TypeVar, TypeForm[Any]] + + def add_argument( + self, parser: Union[argparse.ArgumentParser, argparse._ArgumentGroup] + ) -> None: + """Add a defined argument to a parser.""" + + # Get keyword arguments, with None values removed. + kwargs = dict(self.lowered.__dict__) # type: ignore + kwargs.pop("instantiator") + kwargs = {k: v for k, v in kwargs.items() if v is not None} + name_or_flag = kwargs.pop("name_or_flag") + if len(name_or_flag) == 0: + name_or_flag = _strings.dummy_field_name + + # We're actually going to skip the default field: if an argument is unset, the + # MISSING value will be detected in _calling.py and the field default will + # directly be used. This helps reduce the likelihood of issues with converting + # the field default to a string format, then back to the desired type. + action = kwargs.get("action", None) + if action not in {"append", "count"}: + kwargs["default"] = _fields.MISSING_NONPROP + elif action in {BooleanOptionalAction, "count"}: + pass + else: + kwargs["default"] = [] + + # Apply overrides in our arg configuration object. + # Note that the `name` field is applied when the field object is instantiated! + if self.field.argconf.metavar is not None: + kwargs["metavar"] = self.field.argconf.metavar + + # Add argument, with aliases if available. + if self.field.argconf.aliases is not None and not self.field.is_positional(): + arg = parser.add_argument( + name_or_flag, *self.field.argconf.aliases, **kwargs + ) + else: + if self.field.argconf.aliases is not None: + import warnings + + warnings.warn( + f"Aliases were specified, but {name_or_flag} is positional. Aliases will be ignored." + ) + arg = parser.add_argument(name_or_flag, **kwargs) + + # Do our best to tab complete paths. + # There will be false positives here, but if choices is unset they should be + # harmless. + if "choices" not in kwargs: + name_suggests_dir = ( + # The conditions are intended to be conservative; if a directory path is + # registered as a normal file one that's OK, the reverse on the other + # hand will be overly restrictive. + self.field.intern_name.endswith("_dir") + or self.field.intern_name.endswith("_directory") + or self.field.intern_name.endswith("_folder") + ) + name_suggests_path = ( + self.field.intern_name.endswith("_file") + or self.field.intern_name.endswith("_path") + or self.field.intern_name.endswith("_filename") + or name_suggests_dir + ) + complete_as_path = ( + # Catch types like Path, List[Path], Tuple[Path, ...] etc. + "Path" in str(self.field.type_or_callable) + # For string types, we require more evidence. + or ("str" in str(self.field.type_or_callable) and name_suggests_path) + ) + if complete_as_path: + arg.complete = shtab.DIRECTORY if name_suggests_dir else shtab.FILE # type: ignore + + @cached_property + def lowered(self) -> LoweredArgumentDefinition: + """Lowered argument definition, generated by applying a sequence of rules.""" + # Each rule will mutate the lowered object. This is (unfortunately) + # much faster than a functional approach. + lowered = LoweredArgumentDefinition() + _rule_handle_defaults(self, lowered) + _rule_handle_boolean_flags(self, lowered) + _rule_recursive_instantiator_from_type(self, lowered) + _rule_convert_defaults_to_strings(self, lowered) + _rule_counters(self, lowered) + _rule_generate_helptext(self, lowered) + _rule_set_name_or_flag_and_dest(self, lowered) + _rule_positional_special_handling(self, lowered) + return lowered + + +@dataclasses.dataclass +class LoweredArgumentDefinition: + """Contains fields meant to be passed directly into argparse.""" + + # Action that is called on parsed arguments. This handles conversions from strings + # to our desired types. + # + # The main reason we use this instead of the standard 'type' argument is to enable + # mixed-type tuples. + instantiator: Optional[_instantiators.Instantiator] = None + + def is_fixed(self) -> bool: + """If the instantiator is set to `None`, even after all argument + transformations, it means that we don't have a valid instantiator for an + argument. We then mark the argument as 'fixed', with a value always equal to the + field default.""" + return self.instantiator is None + + # From here on out, all fields correspond 1:1 to inputs to argparse's + # add_argument() method. + name_or_flag: str = "" + default: Optional[Any] = None + dest: Optional[str] = None + required: Optional[bool] = None + action: Optional[Any] = None + nargs: Optional[Union[int, str]] = None + choices: Optional[Tuple[str, ...]] = None + # Note: unlike in vanilla argparse, our metavar is always a string. We handle + # sequences, multiple arguments, etc, manually. + metavar: Optional[str] = None + help: Optional[str] = None + + +def _rule_handle_defaults( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + """Set `required=False` if a default value is set.""" + + # Mark lowered as required if a default is set. + if ( + arg.field.default in _fields.MISSING_SINGLETONS + and _markers._OPTIONAL_GROUP not in arg.field.markers + ): + lowered.default = None + lowered.required = True + return + + lowered.default = arg.field.default + return + + +def _rule_handle_boolean_flags( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + if ( + _resolver.apply_type_from_typevar( + arg.field.type_or_callable, arg.type_from_typevar + ) + is not bool + ): + return + + if ( + arg.field.default in _fields.MISSING_SINGLETONS + or arg.field.is_positional() + or _markers.FlagConversionOff in arg.field.markers + or _markers.Fixed in arg.field.markers + ): + # Treat bools as a normal parameter. + return + elif arg.field.default in (True, False): + # Default `False` => --flag passed in flips to `True`. + lowered.action = BooleanOptionalAction + lowered.instantiator = lambda x: x # argparse will directly give us a bool! + return + + assert False, ( + f"Expected a boolean as a default for {arg.field.intern_name}, but got" + f" {lowered.default}." + ) + + +def _rule_recursive_instantiator_from_type( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + """The bulkiest bit: recursively analyze the type annotation and use it to determine + how to instantiate it given some string from the commandline. + + Important: as far as argparse is concerned, all inputs are strings. + + Conversions from strings to our desired types happen in the instantiator; this is a + bit more flexible, and lets us handle more complex types like enums and multi-type + tuples.""" + if _markers.Fixed in arg.field.markers: + lowered.instantiator = None + lowered.metavar = "{fixed}" + lowered.required = False + lowered.default = _fields.MISSING_PROP + return + if lowered.instantiator is not None: + return + try: + instantiator, metadata = _instantiators.instantiator_from_type( + arg.field.type_or_callable, + arg.type_from_typevar, + arg.field.markers, + ) + except _instantiators.UnsupportedTypeAnnotationError as e: + if arg.field.default in _fields.MISSING_SINGLETONS: + field_name = _strings.make_field_name( + [arg.extern_prefix, arg.field.intern_name] + ) + if field_name != "": + raise _instantiators.UnsupportedTypeAnnotationError( + f"Unsupported type annotation for the field {field_name}; " + f"{e.args[0]} " + "To suppress this error, assign the field either a default value or a different type." + ) from e + else: + # If the field name is empty, it means we're raising an error + # for the direct input to `tyro.cli()`. We don't need to write + # out which specific field we're complaining about. + raise e + else: + # For fields with a default, we'll get by even if there's no instantiator + # available. + lowered.metavar = "{fixed}" + lowered.required = False + lowered.default = _fields.MISSING_PROP + return + + if metadata.action == "append": + + def append_instantiator(x: Any) -> Any: + out = instantiator(x) + if arg.field.default in _fields.MISSING_SINGLETONS: + return instantiator(x) + + return type(out)(arg.field.default) + out + + lowered.instantiator = append_instantiator + lowered.default = None + lowered.choices = metadata.choices + lowered.nargs = metadata.nargs + lowered.metavar = metadata.metavar + lowered.action = metadata.action + lowered.required = False + return + else: + lowered.instantiator = instantiator + lowered.choices = metadata.choices + lowered.nargs = metadata.nargs + lowered.metavar = metadata.metavar + lowered.action = metadata.action + return + + +def _rule_convert_defaults_to_strings( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + """Sets all default values to strings, as required as input to our instantiator + functions. Special-cased for enums.""" + + def as_str(x: Any) -> Tuple[str, ...]: + if isinstance(x, str): + return (x,) + elif isinstance(x, enum.Enum): + return (x.name,) + elif isinstance(x, Mapping): + return tuple(itertools.chain(*map(as_str, itertools.chain(*x.items())))) + elif isinstance(x, Sequence): + return tuple(itertools.chain(*map(as_str, x))) + else: + return (str(x),) + + if ( + lowered.default is None + or lowered.default in _fields.MISSING_SINGLETONS + or lowered.action is not None + ): + return + else: + lowered.default = as_str(lowered.default) + return + + +# This can be turned off when we don't want rich-based formatting. (notably for +# completion scripts) +# +# TODO: the global state here is gross. Should be revisited. +USE_RICH = True + + +# TODO: this function is also called outside of _arguments.py. Should be revisited. +def _rich_tag_if_enabled(x: str, tag: str) -> str: + x = rich.markup.escape(_strings.strip_ansi_sequences(x)) + return x if not USE_RICH else f"[{tag}]{x}[/{tag}]" + + +def _rule_counters( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + """Handle counters, like -vvv for level-3 verbosity.""" + if ( + _markers.UseCounterAction in arg.field.markers + and arg.field.type_or_callable is int + and not arg.field.is_positional() + ): + lowered.metavar = None + lowered.nargs = None + lowered.action = "count" + lowered.default = 0 + lowered.required = False + lowered.instantiator = lambda x: x # argparse will directly give us an int! + return + + +def _rule_generate_helptext( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + """Generate helptext from docstring, argument name, default values.""" + + # If the suppress marker is attached, hide the argument. + if _markers.Suppress in arg.field.markers or ( + _markers.SuppressFixed in arg.field.markers and lowered.is_fixed() + ): + lowered.help = argparse.SUPPRESS + return + + help_parts = [] + + primary_help = arg.field.helptext + + if primary_help is None and _markers.Positional in arg.field.markers: + primary_help = _strings.make_field_name( + [arg.extern_prefix, arg.field.intern_name] + ) + + if primary_help is not None and primary_help != "": + help_parts.append(_rich_tag_if_enabled(primary_help, "helptext")) + + if not lowered.required: + # Get the default value. + # Note: lowered.default is the stringified version! See + # `_rule_convert_defaults_to_strings()`. + default = lowered.default + if lowered.is_fixed() or lowered.action == "append": + # Cases where we'll be missing the lowered default. Use field default instead. + assert default in _fields.MISSING_SINGLETONS or default is None + default = arg.field.default + + # Get the default value label. + if arg.field.argconf.constructor_factory is not None: + default_label = ( + str(default) + if arg.field.type_or_callable is not json.loads + else json.dumps(arg.field.default) + ) + elif hasattr(default, "__iter__"): + # For tuple types, we might have default as (0, 1, 2, 3). + # For list types, we might have default as [0, 1, 2, 3]. + # For set types, we might have default as {0, 1, 2, 3}. + # + # In all cases, we want to display (default: 0 1 2 3), for consistency with + # the format that argparse expects when we set nargs. + assert default is not None + default_label = " ".join(map(shlex.quote, map(str, default))) + else: + default_label = str(default) + + # Include default value in helptext. We intentionally don't use the % template + # because the types of all arguments are set to strings, which will cause the + # default to be casted to a string and introduce extra quotation marks. + if lowered.instantiator is None: + # Intentionally not quoted via shlex, since this can't actually be passed + # in via the commandline. + default_text = f"(fixed to: {default_label})" + elif lowered.action == "count": + # Repeatable argument. + default_text = "(repeatable)" + elif lowered.action == "append" and ( + default in _fields.MISSING_SINGLETONS or len(cast(tuple, default)) == 0 + ): + default_text = "(repeatable)" + elif lowered.action == "append" and len(cast(tuple, default)) > 0: + assert default is not None # Just for type checker. + default_text = f"(repeatable, appends to: {default_label})" + elif arg.field.default is _fields.EXCLUDE_FROM_CALL: + # ^important to use arg.field.default and not the stringified default variable. + default_text = "(unset by default)" + elif ( + _markers._OPTIONAL_GROUP in arg.field.markers + and default in _fields.MISSING_SINGLETONS + ): + # Argument in an optional group, but with no default. This is typically used + # Note: lowered.default is the stringified version! See + # `_rule_convert_defaults_to_strings()`. + # when general (non-argument, non-dataclass) object arguments are given a + # default, or when we use `tyro.conf.arg(constructor=...)`. + # + # There are some usage details that aren't communicated right now in the + # helptext. For example: all arguments within an optional group without a + # default should be passed in or none at all. + default_text = "(optional)" + elif _markers._OPTIONAL_GROUP in arg.field.markers: + # Argument in an optional group, but which also has a default. + default_text = f"(default if used: {default_label})" + else: + default_text = f"(default: {default_label})" + + help_parts.append(_rich_tag_if_enabled(default_text, "helptext_default")) + else: + help_parts.append(_rich_tag_if_enabled("(required)", "helptext_required")) + + # Note that the percent symbol needs some extra handling in argparse. + # https://stackoverflow.com/questions/21168120/python-argparse-errors-with-in-help-string + lowered.help = " ".join(help_parts).replace("%", "%%") + return + + +def _rule_set_name_or_flag_and_dest( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + name_or_flag = _strings.make_field_name( + [arg.extern_prefix, arg.field.extern_name] + if arg.field.argconf.prefix_name + and _markers.OmitArgPrefixes not in arg.field.markers + else [arg.field.extern_name] + ) + + # Prefix keyword arguments with --. + if not arg.field.is_positional(): + name_or_flag = "--" + name_or_flag + + # Strip. + if ( + # If OmitArgPrefixes was applied, then the subcommand prefix was already + # stripped. :) + _markers.OmitArgPrefixes not in arg.field.markers + and _markers.Positional not in arg.field.markers + and name_or_flag.startswith("--") + and arg.subcommand_prefix != "" + ): + # This will run even when unused because we want the assert. + strip_prefix = "--" + arg.subcommand_prefix + "." + if _markers.OmitSubcommandPrefixes in arg.field.markers: + assert name_or_flag.startswith(strip_prefix), name_or_flag + name_or_flag = "--" + name_or_flag[len(strip_prefix) :] + + lowered.name_or_flag = name_or_flag + lowered.dest = _strings.make_field_name([arg.intern_prefix, arg.field.intern_name]) + + +def _rule_positional_special_handling( + arg: ArgumentDefinition, + lowered: LoweredArgumentDefinition, +) -> None: + if not arg.field.is_positional(): + return None + + metavar = lowered.metavar + if lowered.required: + nargs = lowered.nargs + else: + if metavar is not None: + metavar = "[" + metavar + "]" + if lowered.nargs == 1: + # Optional positional arguments. Note that this needs to be special-cased in + # _calling.py. + nargs = "?" + else: + # If lowered.nargs is either + or an int. + nargs = "*" + + lowered.name_or_flag = _strings.make_field_name( + [arg.intern_prefix, arg.field.intern_name] + ) + lowered.dest = None + lowered.required = None # Can't be passed in for positionals. + lowered.metavar = metavar + lowered.nargs = nargs + return diff --git a/src/tyro/_calling.py b/src/tyro/_calling.py new file mode 100644 index 000000000..632b5bc18 --- /dev/null +++ b/src/tyro/_calling.py @@ -0,0 +1,277 @@ +"""Core functionality for calling functions with arguments specified by argparse +namespaces.""" + +from __future__ import annotations + +import dataclasses +import itertools +from functools import partial +from typing import Any, Callable, Dict, List, Set, Tuple, TypeVar, Union + +from typing_extensions import get_args + +from . import _arguments, _fields, _parsers, _resolver, _strings +from .conf import _markers + + +@dataclasses.dataclass(frozen=True) +class InstantiationError(Exception): + """Exception raised when instantiation fail; this typically means that values from + the CLI are invalid.""" + + message: str + arg: Union[_arguments.ArgumentDefinition, str] + + +T = TypeVar("T") + + +def callable_with_args( + f: Callable[..., T], + parser_definition: _parsers.ParserSpecification, + default_instance: Union[T, _fields.NonpropagatingMissingType], + value_from_prefixed_field_name: Dict[str, Any], + field_name_prefix: str, +) -> Tuple[Callable[[], T], Set[str]]: + """Populate `f` with arguments specified by a dictionary of values from argparse. + + Returns a partialed version of `f` with arguments populated, and a set of + used arguments. + + We return a `Callable[[], OutT]` instead of `T` directly for aesthetic + reasons; it lets use reduce layers in stack traces for errors from + functions passed to `tyro`. + """ + + positional_args: List[Any] = [] + kwargs: Dict[str, Any] = {} + consumed_keywords: Set[str] = set() + + def get_value_from_arg(prefixed_field_name: str) -> Any: + """Helper for getting values from `value_from_arg` + doing some extra + asserts.""" + assert ( + prefixed_field_name in value_from_prefixed_field_name + ), f"{prefixed_field_name} not in {value_from_prefixed_field_name}" + return value_from_prefixed_field_name[prefixed_field_name] + + arg_from_prefixed_field_name: Dict[str, _arguments.ArgumentDefinition] = {} + for arg in parser_definition.args: + arg_from_prefixed_field_name[ + _strings.make_field_name([arg.intern_prefix, arg.field.intern_name]) + ] = arg + + any_arguments_provided = False + + for field in parser_definition.field_list: + value: Any + prefixed_field_name = _strings.make_field_name( + [field_name_prefix, field.intern_name] + ) + + # Resolve field type. + field_type = field.type_or_callable + if prefixed_field_name in arg_from_prefixed_field_name: + assert prefixed_field_name not in consumed_keywords + + # Standard arguments. + arg = arg_from_prefixed_field_name[prefixed_field_name] + name_maybe_prefixed = prefixed_field_name + consumed_keywords.add(name_maybe_prefixed) + if not arg.lowered.is_fixed(): + value = get_value_from_arg(name_maybe_prefixed) + + if value in _fields.MISSING_SINGLETONS: + value = arg.field.default + + # Consider a function with a positional sequence argument: + # + # def f(x: tuple[int, ...], /) + # + # If we run this script with no arguments, we should interpret this + # as empty input for x. But the argparse default will be a MISSING + # value, and the field default will be inspect.Parameter.empty. + if ( + value in _fields.MISSING_SINGLETONS + and field.is_positional_call() + and arg.lowered.nargs in ("?", "*") + ): + value = [] + else: + any_arguments_provided = True + if arg.lowered.nargs == "?": + # Special case for optional positional arguments: this is the + # only time that arguments don't come back as a list. + value = [value] + + try: + assert arg.lowered.instantiator is not None + value = arg.lowered.instantiator(value) + except (ValueError, TypeError) as e: + raise InstantiationError( + e.args[0], + arg, + ) + else: + assert arg.field.default not in _fields.MISSING_SINGLETONS + value = arg.field.default + parsed_value = value_from_prefixed_field_name.get(prefixed_field_name) + if parsed_value not in _fields.MISSING_SINGLETONS: + raise InstantiationError( + f"{arg.lowered.name_or_flag} was passed in, but" + " is a fixed argument that cannot be parsed", + arg, + ) + elif prefixed_field_name in parser_definition.child_from_prefix: + # Nested callable. + if _resolver.unwrap_origin_strip_extras(field_type) is Union: + field_type = type(field.default) + get_value, consumed_keywords_child = callable_with_args( + field_type, + parser_definition.child_from_prefix[prefixed_field_name], + field.default, + value_from_prefixed_field_name, + field_name_prefix=prefixed_field_name, + ) + value = get_value() + del get_value + consumed_keywords |= consumed_keywords_child + else: + # Unions over dataclasses (subparsers). This is the only other option. + subparser_def = parser_definition.subparsers_from_intern_prefix[ + prefixed_field_name + ] + subparser_dest = _strings.make_subparser_dest(name=prefixed_field_name) + consumed_keywords.add(subparser_dest) + if subparser_dest in value_from_prefixed_field_name: + subparser_name = get_value_from_arg(subparser_dest) + else: + assert subparser_def.default_instance not in _fields.MISSING_SINGLETONS + subparser_name = None + + if subparser_name is None: + # No subparser selected -- this should only happen when we have a + # default/default_factory set. + assert ( + type(None) in get_args(field_type) + or subparser_def.default_instance is not None + ) + value = subparser_def.default_instance + else: + chosen_f = subparser_def.options[ + list(subparser_def.parser_from_name.keys()).index(subparser_name) + ] + get_value, consumed_keywords_child = callable_with_args( + chosen_f, + subparser_def.parser_from_name[subparser_name], + ( + field.default + if type(field.default) is chosen_f + else _fields.MISSING_NONPROP + ), + value_from_prefixed_field_name, + field_name_prefix=prefixed_field_name, + ) + value = get_value() + del get_value + consumed_keywords |= consumed_keywords_child + + if value is _fields.EXCLUDE_FROM_CALL: + continue + + if _markers._UnpackArgsCall in field.markers: + if len(positional_args) == 0 and len(kwargs) > 0: + positional_args.extend(kwargs.values()) + kwargs.clear() + assert isinstance(value, tuple) + positional_args.extend(value) + elif _markers._UnpackKwargsCall in field.markers: + assert isinstance(value, dict) + kwargs.update(value) + elif field.is_positional_call(): + assert len(kwargs) == 0 + positional_args.append(value) + else: + kwargs[field.call_argname] = value + + # Logic for _markers._OPTIONAL_GROUP. + is_missing_list = [ + v in _fields.MISSING_SINGLETONS + for v in itertools.chain(positional_args, kwargs.values()) + ] + if any(is_missing_list): + if not any_arguments_provided: + # No arguments were provided in this group. + return lambda: default_instance, consumed_keywords # type: ignore + + message = "either all arguments must be provided or none of them." + if len(kwargs) > 0: + missing_args: List[str] = [] + for k, v in kwargs.items(): + if v not in _fields.MISSING_SINGLETONS: + break + + # Argument is missing. + found = False + for arg in arg_from_prefixed_field_name.values(): + if arg.field.call_argname == k: + missing_args.append(arg.lowered.name_or_flag) + found = True + break + assert found, "This is likely a bug in tyro." + + if len(missing_args): + message += f" We're missing arguments {missing_args}." + raise InstantiationError( + message, + field_name_prefix, + ) + + # Note: we unwrap types both before and after narrowing. This is because narrowing + # sometimes produces types like `Tuple[T1, T2, ...]`, where we actually want just + # `tuple`. + unwrapped_f = f + unwrapped_f = _resolver.swap_type_using_confstruct(unwrapped_f) + unwrapped_f = _resolver.unwrap_origin_strip_extras(unwrapped_f) + unwrapped_f = _resolver.unwrap_newtype_and_narrow_subtypes( + unwrapped_f, default_instance + ) + unwrapped_f = _resolver.unwrap_origin_strip_extras(unwrapped_f) + + if unwrapped_f in (tuple, list, set): + if len(positional_args) > 0: + # Triggered when support_single_arg_types=True is used. + assert len(kwargs) == 0 + assert len(positional_args) == 1 + return lambda: positional_args[0], consumed_keywords # type: ignore + else: + assert len(positional_args) == 0 + return partial(unwrapped_f, kwargs.values()), consumed_keywords # type: ignore + elif unwrapped_f is dict: + if len(positional_args) > 0: + # Triggered when support_single_arg_types=True is used. + assert len(kwargs) == 0 + assert len(positional_args) == 1 + return partial(unwrapped_f, *positional_args), consumed_keywords # type: ignore + else: + assert len(positional_args) == 0 + return lambda: kwargs, consumed_keywords # type: ignore + else: + if field_name_prefix == "": + # Don't catch any errors for the "root" field. If main() in tyro.cli(main) + # raises a ValueError, this shouldn't be caught. + return partial(unwrapped_f, *positional_args, **kwargs), consumed_keywords # type: ignore + else: + # Try to catch ValueErrors raised by field constructors. + def with_instantiation_error(): + try: + return unwrapped_f(*positional_args, **kwargs) + # If unwrapped_f raises a ValueError, wrap the message with a more informative + # InstantiationError if possible. + except ValueError as e: + raise InstantiationError( + e.args[0], + field_name_prefix, + ) + + return with_instantiation_error, consumed_keywords # type: ignore diff --git a/src/tyro/_cli.py b/src/tyro/_cli.py new file mode 100644 index 000000000..4173b16f2 --- /dev/null +++ b/src/tyro/_cli.py @@ -0,0 +1,555 @@ +"""Core public API.""" + +import dataclasses +import pathlib +import sys +import warnings +from typing import ( + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + TypeVar, + Union, + cast, + overload, +) + +import shtab +from typing_extensions import Literal + +from . import _argparse as argparse +from . import ( + _argparse_formatter, + _arguments, + _calling, + _fields, + _parsers, + _strings, + _unsafe_cache, + conf, +) +from ._typing import TypeForm + +OutT = TypeVar("OutT") + + +# Note that the overload here is necessary for pyright and pylance due to special-casing +# related to using typing.Type[] as a temporary replacement for typing.TypeForm[]. +# +# https://github.com/microsoft/pyright/issues/4298 + + +@overload +def cli( + f: TypeForm[OutT], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + default: Optional[OutT] = None, + return_unknown_args: Literal[False] = False, + use_underscores: bool = False, + console_outputs: bool = True, + config: Optional[Sequence[conf._markers.Marker]] = None, +) -> OutT: ... + + +@overload +def cli( + f: TypeForm[OutT], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + default: Optional[OutT] = None, + return_unknown_args: Literal[True], + use_underscores: bool = False, + console_outputs: bool = True, + config: Optional[Sequence[conf._markers.Marker]] = None, +) -> Tuple[OutT, List[str]]: ... + + +@overload +def cli( + f: Callable[..., OutT], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + # Passing a default makes sense for things like dataclasses, but are not + # supported for general callables. These can, however, be specified in the + # signature of the callable itself. + default: None = None, + return_unknown_args: Literal[False] = False, + use_underscores: bool = False, + console_outputs: bool = True, + config: Optional[Sequence[conf._markers.Marker]] = None, +) -> OutT: ... + + +@overload +def cli( + f: Callable[..., OutT], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + # Passing a default makes sense for things like dataclasses, but are not + # supported for general callables. These can, however, be specified in the + # signature of the callable itself. + default: None = None, + return_unknown_args: Literal[True], + use_underscores: bool = False, + console_outputs: bool = True, + config: Optional[Sequence[conf._markers.Marker]] = None, +) -> Tuple[OutT, List[str]]: ... + + +def cli( + f: Union[TypeForm[OutT], Callable[..., OutT]], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + default: Optional[OutT] = None, + return_unknown_args: bool = False, + use_underscores: bool = False, + console_outputs: bool = True, + config: Optional[Sequence[conf._markers.Marker]] = None, + **deprecated_kwargs, +) -> Union[OutT, Tuple[OutT, List[str]]]: + """Call or instantiate `f`, with inputs populated from an automatically generated + CLI interface. + + `f` should have type-annotated inputs, and can be a function or type. If + `f` is a type, `tyro.cli()` returns an instance. + + The parser is generated by populating helptext from docstrings and types from + annotations; a broad range of core type annotations are supported. + - Types natively accepted by `argparse`: str, int, float, pathlib.Path, etc. + - Default values for optional parameters. + - Booleans, which are automatically converted to flags when provided a default + value. + - Enums (via `enum.Enum`). + - Various annotations from the standard typing library. Some examples: + - `typing.ClassVar[T]`. + - `typing.Optional[T]`. + - `typing.Literal[T]`. + - `typing.Sequence[T]`. + - `typing.List[T]`. + - `typing.Dict[K, V]`. + - `typing.Tuple`, such as `typing.Tuple[T1, T2, T3]` or + `typing.Tuple[T, ...]`. + - `typing.Set[T]`. + - `typing.Final[T]` and `typing.Annotated[T]`. + - `typing.Union[T1, T2]`. + - Various nested combinations of the above: `Optional[Literal[T]]`, + `Final[Optional[Sequence[T]]]`, etc. + - Hierarchical structures via nested dataclasses, TypedDict, NamedTuple, + classes. + - Simple nesting. + - Unions over nested structures (subparsers). + - Optional unions over nested structures (optional subparsers). + - Generics (including nested generics). + + Completion script generation for interactive shells is also provided. To write a + script that can be used for tab completion, pass in: + `--tyro-write-completion {bash/zsh/tcsh} {path to script to write}`. + + Args: + f: Function or type. + 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, `f`'s 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()`. + default: An instance of `OutT` to use for default values; supported if `f` is a + type like a dataclass or dictionary, but not if `f` is a general callable + like a function or standard class. Helpful for merging CLI arguments with + values loaded from elsewhere. (for example, a config object loaded from a + yaml file) + return_unknown_args: If True, return a tuple of the output of `f` and a list of + unknown arguments. Mirrors the unknown arguments returned from + `argparse.ArgumentParser.parse_known_args()`. + use_underscores: If True, use underscores as a word delimeter 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 + console_outputs: If set to `False`, parsing errors and help messages will be + supressed. This can be useful for distributed settings, where `tyro.cli()` + is called from multiple workers but we only want console outputs from the + main one. + config: Sequence of config marker objects, from `tyro.conf`. As an + alternative to using them locally in annotations + (`tyro.conf.FlagConversionOff[bool]`), we can also pass in a sequence of + them here to apply globally. + + Returns: + The output of `f(...)` or an instance `f`. If `f` is a class, the two are + equivalent. If `return_unknown_args` is True, returns a tuple of the output of + `f(...)` and a list of unknown arguments. + """ + + # Make sure we start on a clean slate. Some tests may fail without this due to + # memory address conflicts. + _unsafe_cache.clear_cache() + + if config is not None: + f = conf.configure(*config)(f) + + with _strings.delimeter_context("_" if use_underscores else "-"): + output = _cli_impl( + f, + prog=prog, + description=description, + args=args, + default=default, + return_parser=False, + return_unknown_args=return_unknown_args, + use_underscores=use_underscores, + console_outputs=console_outputs, + **deprecated_kwargs, + ) + + # Prevent unnecessary memory usage. + _unsafe_cache.clear_cache() + + if return_unknown_args: + assert isinstance(output, tuple) + run_with_args_from_cli = output[0] + return run_with_args_from_cli(), output[1] + else: + run_with_args_from_cli = cast(Callable[[], OutT], output) + return run_with_args_from_cli() + + +@overload +def get_parser( + f: TypeForm[OutT], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + default: Optional[OutT] = None, + use_underscores: bool = False, + console_outputs: bool = True, +) -> argparse.ArgumentParser: ... + + +@overload +def get_parser( + f: Callable[..., OutT], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + default: Optional[OutT] = None, + use_underscores: bool = False, + console_outputs: bool = True, +) -> argparse.ArgumentParser: ... + + +def get_parser( + f: Union[TypeForm[OutT], Callable[..., OutT]], + *, + # We have no `args` argument, since this is only used when + # parser.parse_args() is called. + prog: Optional[str] = None, + description: Optional[str] = None, + default: Optional[OutT] = None, + use_underscores: bool = False, + console_outputs: bool = True, +) -> argparse.ArgumentParser: + """Get the `argparse.ArgumentParser` object generated under-the-hood by + `tyro.cli()`. Useful for tools like `sphinx-argparse`, `argcomplete`, etc. + + For tab completion, we recommend using `tyro.cli()`'s built-in `--tyro-write-completion` + flag.""" + with _strings.delimeter_context("_" if use_underscores else "-"): + return cast( + argparse.ArgumentParser, + _cli_impl( + f, + prog=prog, + description=description, + args=None, + default=default, + return_parser=True, + return_unknown_args=False, + use_underscores=use_underscores, + console_outputs=console_outputs, + ), + ) + + +def _cli_impl( + f: Union[TypeForm[OutT], Callable[..., OutT]], + *, + prog: Optional[str] = None, + description: Optional[str], + args: Optional[Sequence[str]], + default: Optional[OutT], + return_parser: bool, + return_unknown_args: bool, + console_outputs: bool, + **deprecated_kwargs, +) -> Union[ + OutT, + argparse.ArgumentParser, + Tuple[Callable[[], OutT], List[str]], +]: + """Helper for stitching the `tyro` pipeline together.""" + if "default_instance" in deprecated_kwargs: + warnings.warn( + "`default_instance=` is deprecated! use `default=` instead.", stacklevel=2 + ) + default = deprecated_kwargs["default_instance"] + if deprecated_kwargs.get("avoid_subparsers", False): + f = conf.AvoidSubcommands[f] # type: ignore + warnings.warn( + "`avoid_subparsers=` is deprecated! use `tyro.conf.AvoidSubparsers[]`" + " instead.", + stacklevel=2, + ) + + # Internally, we distinguish between two concepts: + # - "default", which is used for individual arguments. + # - "default_instance", which is used for _fields_ (which may be broken down into + # one or many arguments, depending on various factors). + # + # This could be revisited. + default_instance_internal: Union[_fields.NonpropagatingMissingType, OutT] = ( + _fields.MISSING_NONPROP if default is None else default + ) + + # We wrap our type with a dummy dataclass if it can't be treated as a nested type. + # For example: passing in f=int will result in a dataclass with a single field + # typed as int. + if not _fields.is_nested_type(cast(type, f), default_instance_internal): + dummy_field = cast( + dataclasses.Field, + dataclasses.field(), + ) + f = dataclasses.make_dataclass( + cls_name="dummy", + fields=[(_strings.dummy_field_name, cast(type, f), dummy_field)], + frozen=True, + ) + default_instance_internal = f(default_instance_internal) # type: ignore + dummy_wrapped = True + else: + dummy_wrapped = False + + # Read and fix arguments. If the user passes in --field_name instead of + # --field-name, correct for them. + args = list(sys.argv[1:]) if args is None else list(args) + + # Fix arguments. This will modify all option-style arguments replacing + # underscores with hyphens, or vice versa if use_underscores=True. + # If two options are ambiguous, e.g., --a_b and --a-b, raise a runtime error. + modified_args: Dict[str, str] = {} + for index, arg in enumerate(args): + if not arg.startswith("--"): + continue + + if "=" in arg: + arg, _, val = arg.partition("=") + fixed = "--" + _strings.replace_delimeter_in_part(arg[2:]) + "=" + val + else: + fixed = "--" + _strings.replace_delimeter_in_part(arg[2:]) + if ( + return_unknown_args + and fixed in modified_args + and modified_args[fixed] != arg + ): + raise RuntimeError( + "Ambiguous arguments: " + modified_args[fixed] + " and " + arg + ) + modified_args[fixed] = arg + args[index] = fixed + + # If we pass in the --tyro-print-completion or --tyro-write-completion flags: turn + # formatting tags, and get the shell we want to generate a completion script for + # (bash/zsh/tcsh). + # + # shtab also offers an add_argument_to() functions that fulfills a similar goal, but + # manual parsing of argv is convenient for turning off formatting. + # + # Note: --tyro-print-completion is deprecated! --tyro-write-completion is less prone + # to errors from accidental logging, print statements, etc. + print_completion = len(args) >= 2 and args[0] == "--tyro-print-completion" + write_completion = len(args) >= 3 and args[0] == "--tyro-write-completion" + + # Note: setting USE_RICH must happen before the parser specification is generated. + # TODO: revisit this. Ideally we should be able to eliminate the global state + # changes. + completion_shell = None + completion_target_path = None + if print_completion or write_completion: + completion_shell = args[1] + if write_completion: + completion_target_path = pathlib.Path(args[2]) + if print_completion or write_completion or return_parser: + _arguments.USE_RICH = False + else: + _arguments.USE_RICH = True + + # Map a callable to the relevant CLI arguments + subparsers. + parser_spec = _parsers.ParserSpecification.from_callable_or_type( + f, + description=description, + parent_classes=set(), # Used for recursive calls. + default_instance=default_instance_internal, # Overrides for default values. + intern_prefix="", # Used for recursive calls. + extern_prefix="", # Used for recursive calls. + subcommand_prefix="", # Used for recursive calls. + ) + + # Generate parser! + with _argparse_formatter.ansi_context(): + parser = _argparse_formatter.TyroArgumentParser( + prog=prog, + formatter_class=_argparse_formatter.TyroArgparseHelpFormatter, + allow_abbrev=False, + ) + parser._parser_specification = parser_spec + parser._parsing_known_args = return_unknown_args + parser._console_outputs = console_outputs + parser._args = args + parser_spec.apply(parser) + + # Print help message when no arguments are passed in. (but arguments are + # expected) + # if len(args) == 0 and parser_spec.has_required_args: + # args = ["--help"] + + if return_parser: + _arguments.USE_RICH = True + return parser + + if print_completion or write_completion: + _arguments.USE_RICH = True + assert completion_shell in ( + "bash", + "zsh", + "tcsh", + ), ( + "Shell should be one `bash`, `zsh`, or `tcsh`, but got" + f" {completion_shell}" + ) + + if write_completion: + assert completion_target_path is not None + completion_target_path.write_text( + shtab.complete( + parser=parser, + shell=completion_shell, + root_prefix=f"tyro_{parser.prog}", + ) + ) + else: + print( + shtab.complete( + parser=parser, + shell=completion_shell, + root_prefix=f"tyro_{parser.prog}", + ) + ) + sys.exit() + + if return_unknown_args: + namespace, unknown_args = parser.parse_known_args(args=args) + else: + unknown_args = None + namespace = parser.parse_args(args=args) + value_from_prefixed_field_name = vars(namespace) + + if dummy_wrapped: + value_from_prefixed_field_name = { + k.replace(_strings.dummy_field_name, ""): v + for k, v in value_from_prefixed_field_name.items() + } + + try: + # Attempt to call `f` using whatever was passed in. + get_out, consumed_keywords = _calling.callable_with_args( + f, + parser_spec, + default_instance_internal, + value_from_prefixed_field_name, + field_name_prefix="", + ) + except _calling.InstantiationError as e: + # Print prettier errors. + # Note that this doesn't catch errors raised directly by get_out(), + # since that's called later! This is intentional, because we do less + # error handling for the root callable. Relevant: the + # `field_name_prefix == ""` condition in `callable_with_args()`! + + # Emulate argparse's error behavior when invalid arguments are passed in. + from rich.console import Console, Group, RenderableType + from rich.padding import Padding + from rich.panel import Panel + from rich.rule import Rule + from rich.style import Style + + from ._argparse_formatter import THEME + + if console_outputs: + console = Console(theme=THEME.as_rich_theme(), stderr=True) + console.print( + Panel( + Group( + "[bright_red][bold]Error parsing" + f" {e.arg.lowered.name_or_flag if isinstance(e.arg, _arguments.ArgumentDefinition) else e.arg}[/bold]:[/bright_red] {e.message}", + *cast( # Cast to appease mypy... + List[RenderableType], + ( + [] + if not isinstance(e.arg, _arguments.ArgumentDefinition) + or e.arg.lowered.help is None + else [ + Rule(style=Style(color="red")), + "Argument helptext:", + Padding( + Group( + f"{e.arg.lowered.name_or_flag} [bold]{e.arg.lowered.metavar}[/bold]", + e.arg.lowered.help, + ), + pad=(0, 0, 0, 4), + ), + Rule(style=Style(color="red")), + f"For full helptext, see [bold]{parser.prog} --help[/bold]", + ] + ), + ), + ), + title="[bold]Value error[/bold]", + title_align="left", + border_style=Style(color="red"), + ) + ) + sys.exit(2) + + assert len(value_from_prefixed_field_name.keys() - consumed_keywords) == 0, ( + f"Parsed {value_from_prefixed_field_name.keys()}, but only consumed" + f" {consumed_keywords}" + ) + + if dummy_wrapped: + get_wrapped_out = get_out + get_out = lambda: getattr(get_wrapped_out(), _strings.dummy_field_name) # noqa + + if return_unknown_args: + assert unknown_args is not None, "Should have parsed with `parse_known_args()`" + # If we're parsed unknown args, we should return the original args, not + # the fixed ones. + unknown_args = [modified_args.get(arg, arg) for arg in unknown_args] + return get_out, unknown_args # type: ignore + else: + assert unknown_args is None, "Should have parsed with `parse_args()`" + return get_out # type: ignore diff --git a/src/tyro/_deprecated.py b/src/tyro/_deprecated.py new file mode 100644 index 000000000..c879522a8 --- /dev/null +++ b/src/tyro/_deprecated.py @@ -0,0 +1,2 @@ +from ._cli import cli as parse # noqa +from .extras._serialization import from_yaml, to_yaml # noqa diff --git a/src/tyro/_docstrings.py b/src/tyro/_docstrings.py new file mode 100644 index 000000000..3d461927b --- /dev/null +++ b/src/tyro/_docstrings.py @@ -0,0 +1,349 @@ +"""Helpers for parsing docstrings. Used for helptext generation.""" + +import builtins +import collections.abc +import dataclasses +import functools +import inspect +import io +import itertools +import tokenize +from typing import Callable, Dict, Generic, Hashable, List, Optional, Set, Type, TypeVar + +import docstring_parser +from typing_extensions import get_origin, is_typeddict + +from . import _resolver, _strings, _unsafe_cache + +T = TypeVar("T", bound=Callable) + + +@dataclasses.dataclass(frozen=True) +class _Token: + token_type: int + content: str + logical_line: int + actual_line: int + + +@dataclasses.dataclass(frozen=True) +class _FieldData: + index: int + logical_line: int + actual_line: int + prev_field_logical_line: int + + +@dataclasses.dataclass(frozen=True) +class _ClassTokenization: + tokens: List[_Token] + tokens_from_logical_line: Dict[int, List[_Token]] + tokens_from_actual_line: Dict[int, List[_Token]] + field_data_from_name: Dict[str, _FieldData] + + @staticmethod + @_unsafe_cache.unsafe_cache(64) + def make(clz) -> "_ClassTokenization": + """Parse the source code of a class, and cache some tokenization information.""" + readline = io.BytesIO(inspect.getsource(clz).encode("utf-8")).readline + + tokens: List[_Token] = [] + tokens_from_logical_line: Dict[int, List[_Token]] = {1: []} + tokens_from_actual_line: Dict[int, List[_Token]] = {1: []} + field_data_from_name: Dict[str, _FieldData] = {} + + logical_line: int = 1 + actual_line: int = 1 + for toktype, tok, start, end, line in tokenize.tokenize(readline): + # Note: we only track logical line numbers, which are delimited by + # `tokenize.NEWLINE`. `tokenize.NL` tokens appear when logical lines are + # broken into multiple lines of code; these are ignored. + if toktype == tokenize.NEWLINE: + logical_line += 1 + actual_line += 1 + tokens_from_logical_line[logical_line] = [] + tokens_from_actual_line[actual_line] = [] + elif toktype == tokenize.NL: + actual_line += 1 + tokens_from_actual_line[actual_line] = [] + elif toktype is not tokenize.INDENT: + token = _Token( + token_type=toktype, + content=tok, + logical_line=logical_line, + actual_line=actual_line, + ) + tokens.append(token) + tokens_from_logical_line[logical_line].append(token) + tokens_from_actual_line[actual_line].append(token) + + prev_field_logical_line: int = 1 + for i, token in enumerate(tokens[:-1]): + if token.token_type == tokenize.NAME: + # Naive heuristic for field names. + is_first_token = True + for t in tokens_from_logical_line[token.logical_line]: + if t == token: + break + if t.token_type is not tokenize.COMMENT: + is_first_token = False + break + + if not is_first_token: + continue + + if ( + tokens[i + 1].content == ":" + and token.content not in field_data_from_name + ): + field_data_from_name[token.content] = _FieldData( + index=i, + logical_line=token.logical_line, + actual_line=token.actual_line, + prev_field_logical_line=prev_field_logical_line, + ) + prev_field_logical_line = token.logical_line + + return _ClassTokenization( + tokens=tokens, + tokens_from_logical_line=tokens_from_logical_line, + tokens_from_actual_line=tokens_from_actual_line, + field_data_from_name=field_data_from_name, + ) + + +@_unsafe_cache.unsafe_cache(1024) +def get_class_tokenization_with_field( + cls: Type, field_name: str +) -> Optional[_ClassTokenization]: + # Search for token in this class + all parents. + found_field: bool = False + classes_to_search = cls.__mro__ + tokenization = None + for search_cls in classes_to_search: + # Inherited generics seem challenging for now. + # https://github.com/python/typing/issues/777 + assert search_cls is Generic or get_origin(search_cls) is None + + try: + tokenization = _ClassTokenization.make(search_cls) # type: ignore + except OSError: + # OSError is raised when we can't read the source code. This is + # fine, we just assume there's no docstring. We can uncomment the + # assert below for debugging. + # + # assert ( + # # Dynamic dataclasses. + # "could not find class definition" in e.args[0] + # # Pydantic. + # or "source code not available" in e.args[0] + # # Third error that can be raised by inspect.py. + # or "could not get source code" in e.args[0] + # ) + return None + except TypeError as e: # pragma: no cover + # Notebooks cause “___ is a built-in class” TypeError. + assert "built-in class" in e.args[0] + return None + + # Grab field-specific tokenization data. + if field_name in tokenization.field_data_from_name: + found_field = True + break + + if dataclasses.is_dataclass(cls): + assert found_field, ( + "Docstring parsing error -- this usually means that there are multiple" + " dataclasses in the same file with the same name but different scopes." + ) + + return tokenization + + +@functools.lru_cache(maxsize=1024) +def parse_docstring_from_object(obj: object) -> Dict[str, str]: + return { + doc.arg_name: doc.description + for doc in docstring_parser.parse_from_object(obj).params + if doc.description is not None + } + + +@_unsafe_cache.unsafe_cache(1024) +def get_field_docstring(cls: Type, field_name: str) -> Optional[str]: + """Get docstring for a field in a class.""" + + # NoneType will break docstring_parser. + if cls is type(None): + return None + + # Try to parse using docstring_parser. + for cls_search in cls.__mro__: + if cls_search.__module__ == "builtins": + continue # Skip `object`, `Callable`, `tuple`, etc. + docstring = parse_docstring_from_object(cls_search).get(field_name, None) + if docstring is not None: + return _strings.dedent( + _strings.remove_single_line_breaks(docstring) + ).strip() + + # If docstring_parser failed, let's try looking for comments. + tokenization = get_class_tokenization_with_field(cls, field_name) + if tokenization is None: # Currently only happens for dynamic dataclasses. + return None + + field_data = tokenization.field_data_from_name[field_name] + + # Check for comment on the same line as the field. + final_token_on_line = tokenization.tokens_from_logical_line[ + field_data.logical_line + ][-1] + if final_token_on_line.token_type == tokenize.COMMENT: + comment: str = final_token_on_line.content + assert comment.startswith("#") + if comment.startswith("#:"): # Sphinx autodoc-style comment. + # https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#directive-autoattribute + return _strings.remove_single_line_breaks(comment[2:].strip()) + else: + return _strings.remove_single_line_breaks(comment[1:].strip()) + + # Check for comments that come before the field. This is intentionally written to + # support comments covering multiple (grouped) fields, for example: + # + # # Optimizer hyperparameters. + # learning_rate: float + # beta1: float + # beta2: float + # + # In this case, 'Optimizer hyperparameters' will be treated as the docstring for all + # 3 fields. There are tradeoffs we are making here. + # + # The exception this is Sphinx-style comments: + # + # #: The learning rate. + # learning_rate: float + # beta1: float + # beta2: float + # + # Where, by convention the comment only applies to the field that directly follows it. + + # Get first line of the class definition, excluding decorators. This logic is only + # needed for Python >= 3.9; in 3.8, we can simply use + # `tokenization.tokens[0].logical_line`. + classdef_logical_line = -1 + for token in tokenization.tokens: + if token.content == "class": + classdef_logical_line = token.logical_line + break + assert classdef_logical_line != -1 + + comments: List[str] = [] + current_actual_line = field_data.actual_line - 1 + directly_above_field = True + is_sphinx_doc_comment = False + while current_actual_line in tokenization.tokens_from_actual_line: + actual_line_tokens = tokenization.tokens_from_actual_line[current_actual_line] + + # We stop looking if we find an empty line. + if len(actual_line_tokens) == 0: + break + + # We don't look in the first logical line. This includes all comments that come + # before the end parentheses in the class definition (eg comments in the + # subclass list). + if actual_line_tokens[0].logical_line <= classdef_logical_line: + break + + # Record single comments! + if ( + len(actual_line_tokens) == 1 + and actual_line_tokens[0].token_type is tokenize.COMMENT + ): + (comment_token,) = actual_line_tokens + assert comment_token.content.startswith("#") + if comment_token.content.startswith("#:"): # Sphinx autodoc-style comment. + # https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#directive-autoattribute + comments.append(comment_token.content[2:].strip()) + is_sphinx_doc_comment = True + else: + comments.append(comment_token.content[1:].strip()) + elif len(comments) > 0: + # Comments should be contiguous. + break + else: + # This comment is not directly above the current field. + directly_above_field = False + + current_actual_line -= 1 + + if len(comments) > 0 and not (is_sphinx_doc_comment and not directly_above_field): + return _strings.remove_single_line_breaks("\n".join(reversed(comments))) + + return None + + +_callable_description_blocklist: Set[Hashable] = set( + filter( + lambda x: isinstance(x, Hashable), # type: ignore + itertools.chain( + vars(builtins).values(), + vars(collections.abc).values(), + ), + ) +) + + +@_unsafe_cache.unsafe_cache(1024) +def get_callable_description(f: Callable) -> str: + """Get description associated with a callable via docstring parsing. + + Note that the `dataclasses.dataclass` will automatically populate __doc__ based on + the fields of the class if a docstring is not specified; this helper will ignore + these docstrings.""" + + f, _unused = _resolver.resolve_generic_types(f) + f = _resolver.unwrap_origin_strip_extras(f) + if f in _callable_description_blocklist: + return "" + + # Return original docstring when used with functools.partial, not + # functools.partial's docstinrg. + if isinstance(f, functools.partial): + f = f.func + + try: + import pydantic + except ImportError: + pydantic = None # type: ignore + + # Note inspect.getdoc() causes some corner cases with TypedDicts. + docstring = f.__doc__ + if ( + docstring is None + and inspect.isclass(f) + # Ignore TypedDict's __init__ docstring, because it will just be `dict` + and not is_typeddict(f) + # Ignore NamedTuple __init__ docstring. + and not _resolver.is_namedtuple(f) + # Ignore pydantic base model constructor docstring. + and not (pydantic is not None and f.__init__ is pydantic.BaseModel.__init__) # type: ignore + ): + docstring = f.__init__.__doc__ # type: ignore + if docstring is None: + return "" + + docstring = _strings.dedent(docstring) + + if dataclasses.is_dataclass(f): + default_doc = f.__name__ + str(inspect.signature(f)).replace(" -> None", "") # type: ignore + if docstring == default_doc: + return "" + + parsed_docstring = docstring_parser.parse(docstring) + + parts: List[str] = [] + if parsed_docstring.short_description is not None: + parts.append(parsed_docstring.short_description) + if parsed_docstring.long_description is not None: + parts.append(parsed_docstring.long_description) + return "\n".join(parts) diff --git a/src/tyro/_fields.py b/src/tyro/_fields.py new file mode 100644 index 000000000..0f43ea9cb --- /dev/null +++ b/src/tyro/_fields.py @@ -0,0 +1,1194 @@ +"""Abstractions for pulling out 'field' definitions, which specify inputs, types, and # type: ignore +defaults, from general callables.""" + +from __future__ import annotations + +import builtins +import collections +import collections.abc +import dataclasses +import enum +import functools +import inspect +import itertools +import numbers +import os +import sys +import typing +import warnings +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Hashable, + Iterable, + List, + Optional, + Set, + Tuple, + TypeVar, + Union, + cast, +) + +import docstring_parser +import typing_extensions +from typing_extensions import ( + Annotated, + NotRequired, + Required, + get_args, + get_origin, + is_typeddict, +) + +from . import ( + _docstrings, + _instantiators, + _resolver, + _singleton, + _strings, + _unsafe_cache, + conf, # Avoid circular import. +) +from ._typing import TypeForm +from .conf import _confstruct, _markers + + +@dataclasses.dataclass +class FieldDefinition: + intern_name: str + extern_name: str + type_or_callable: Union[TypeForm[Any], Callable] + """Type or callable for this field. This should have all Annotated[] annotations + stripped.""" + default: Any + # We need to record whether defaults are from default instances to + # determine if they should override the default in + # tyro.conf.subcommand(default=...). + is_default_from_default_instance: bool + helptext: Optional[str] + markers: Set[Any] + custom_constructor: bool + + argconf: _confstruct._ArgConfiguration + + # Override the name in our kwargs. Useful whenever the user-facing argument name + # doesn't match the keyword expected by our callable. + call_argname: Any + + def __post_init__(self): + if ( + _markers.Fixed in self.markers or _markers.Suppress in self.markers + ) and self.default in MISSING_SINGLETONS: + raise _instantiators.UnsupportedTypeAnnotationError( + f"Field {self.intern_name} is missing a default value!" + ) + + @staticmethod + def make( + name: str, + type_or_callable: Union[TypeForm[Any], Callable], + default: Any, + is_default_from_default_instance: bool, + helptext: Optional[str], + call_argname_override: Optional[Any] = None, + *, + markers: Tuple[_markers.Marker, ...] = (), + ): + # Try to extract argconf overrides from type. + _, argconfs = _resolver.unwrap_annotated( + type_or_callable, _confstruct._ArgConfiguration + ) + argconf = _confstruct._ArgConfiguration( + None, + None, + help=None, + aliases=None, + prefix_name=True, + constructor_factory=None, + ) + for overwrite_argconf in argconfs: + # Apply any annotated argument configuration values. + argconf = dataclasses.replace( + argconf, + **{ + field.name: getattr(overwrite_argconf, field.name) + for field in dataclasses.fields(overwrite_argconf) + if getattr(overwrite_argconf, field.name) is not None + }, + ) + if argconf.help is not None: + helptext = argconf.help + + type_or_callable, inferred_markers = _resolver.unwrap_annotated( + type_or_callable, _markers._Marker + ) + return FieldDefinition( + intern_name=name, + extern_name=name if argconf.name is None else argconf.name, + type_or_callable=type_or_callable + if argconf.constructor_factory is None + else argconf.constructor_factory(), + default=default, + is_default_from_default_instance=is_default_from_default_instance, + helptext=helptext, + markers=set(inferred_markers).union(markers), + custom_constructor=argconf.constructor_factory is not None, + argconf=argconf, + call_argname=call_argname_override + if call_argname_override is not None + else name, + ) + + def add_markers(self, markers: Tuple[Any, ...]) -> FieldDefinition: + return dataclasses.replace( + self, + markers=self.markers.union(markers), + ) + + def is_positional(self) -> bool: + """Returns True if the argument should be positional in the commandline.""" + return ( + # Explicit positionals. + _markers.Positional in self.markers + # Dummy dataclasses should have a single positional field. + or self.intern_name == _strings.dummy_field_name + or ( + # Make required arguments positional. + _markers.PositionalRequiredArgs in self.markers + and self.default in MISSING_SINGLETONS + ) + ) + + def is_positional_call(self) -> bool: + """Returns True if the argument should be positional in underlying Python call.""" + return ( + # Explicit positionals. + _markers._PositionalCall in self.markers + # Dummy dataclasses should have a single positional field. + or self.intern_name == _strings.dummy_field_name + ) + + +class PropagatingMissingType(_singleton.Singleton): + pass + + +class NonpropagatingMissingType(_singleton.Singleton): + pass + + +class ExcludeFromCallType(_singleton.Singleton): + pass + + +class NotRequiredButWeDontKnowTheValueType(_singleton.Singleton): + pass + + +# We have two types of missing sentinels: a propagating missing value, which when set as +# a default will set all child values of nested structures as missing as well, and a +# nonpropagating missing sentinel, which does not override child defaults. +MISSING_PROP = PropagatingMissingType() +MISSING_NONPROP = NonpropagatingMissingType() + +# When total=False in a TypedDict, we exclude fields from the constructor by default. +NOT_REQUIRED_BUT_WE_DONT_KNOW_THE_VALUE = NotRequiredButWeDontKnowTheValueType() + + +EXCLUDE_FROM_CALL = ExcludeFromCallType() + +# Note that our "public" missing API will always be the propagating missing sentinel. +MISSING: Any = MISSING_PROP +"""Sentinel value to mark fields as missing. Can be used to mark fields passed +in via `default=` for `tyro.cli()` as required.""" + + +MISSING_SINGLETONS = [ + dataclasses.MISSING, + MISSING_PROP, + MISSING_NONPROP, + inspect.Parameter.empty, +] +try: + # Undocumented feature: support omegaconf dataclasses out of the box. + import omegaconf + + MISSING_SINGLETONS.append(omegaconf.MISSING) +except ImportError: + pass + +DEFAULT_SENTINEL_SINGLETONS = MISSING_SINGLETONS + [ + NOT_REQUIRED_BUT_WE_DONT_KNOW_THE_VALUE, + EXCLUDE_FROM_CALL, +] + + +@dataclasses.dataclass(frozen=True) +class UnsupportedNestedTypeMessage: + """Reason why a callable cannot be treated as a nested type.""" + + message: str + + +@_unsafe_cache.unsafe_cache(maxsize=1024) +def is_nested_type( + typ: Union[TypeForm[Any], Callable], default_instance: DefaultInstance +) -> bool: + """Determine whether a type should be treated as a 'nested type', where a single + type can be broken down into multiple fields (eg for nested dataclasses or + classes). + + TODO: we should come up with a better name than 'nested type', which is a little bit + misleading.""" + + return not isinstance( + _try_field_list_from_callable(typ, default_instance), + UnsupportedNestedTypeMessage, + ) + + +def field_list_from_callable( + f: Union[Callable, TypeForm[Any]], + default_instance: DefaultInstance, + support_single_arg_types: bool, +) -> Tuple[ + Union[Callable, TypeForm[Any]], Dict[TypeVar, TypeForm], List[FieldDefinition] +]: + """Generate a list of generic 'field' objects corresponding to the inputs of some + annotated callable. + + Returns: + The type that `f` is resolved as. + A type_from_typevar dict. + A list of field definitions. + """ + # Resolve generic types. + f, type_from_typevar = _resolver.resolve_generic_types(f) + f = _resolver.unwrap_newtype_and_narrow_subtypes(f, default_instance) + + # Try to generate field list. + field_list = _try_field_list_from_callable(f, default_instance) + + if isinstance(field_list, UnsupportedNestedTypeMessage): + if support_single_arg_types: + return ( + f, + type_from_typevar, + [ + FieldDefinition( + intern_name="value", + extern_name="value", # Doesn't matter. + type_or_callable=f, + default=default_instance, + is_default_from_default_instance=True, + helptext="", + custom_constructor=False, + markers={_markers.Positional, _markers._PositionalCall}, + argconf=_confstruct._ArgConfiguration( + None, None, None, None, None, None + ), + call_argname="", + ) + ], + ) + else: + raise _instantiators.UnsupportedTypeAnnotationError(field_list.message) + + # Recursively apply markers. + _, parent_markers = _resolver.unwrap_annotated(f, _markers._Marker) + field_list = list(map(lambda field: field.add_markers(parent_markers), field_list)) + + # Try to resolve types in our list of fields. + def resolve(field: FieldDefinition) -> FieldDefinition: + typ = field.type_or_callable + typ = _resolver.apply_type_from_typevar(typ, type_from_typevar) + typ = _resolver.type_from_typevar_constraints(typ) + typ = _resolver.narrow_collection_types(typ, field.default) + typ = _resolver.narrow_union_type(typ, field.default) + + # Check that the default value matches the final resolved type. + # There's some similar Union-specific logic for this in narrow_union_type(). We + # may be able to consolidate this. + if ( + # Be relatively conservative: isinstance() can be checked on non-type + # types (like unions in Python >=3.10), but we'll only consider single types + # for now. + type(typ) is type + and not isinstance(field.default, typ) # type: ignore + # If a custom constructor is set, field.type_or_callable may not be + # matched to the annotated type. + and not field.custom_constructor + and field.default not in DEFAULT_SENTINEL_SINGLETONS + # The numeric tower in Python is wacky. This logic is non-critical, so + # we'll just skip it (+the complexity) for numbers. + and not isinstance(field.default, numbers.Number) + ): + # If the default value doesn't match the resolved type, we expand the + # type. This is inspired by https://github.com/brentyi/tyro/issues/88. + warnings.warn( + f"The field {field.intern_name} is annotated with type {field.type_or_callable}, " + f"but the default value {field.default} has type {type(field.default)}. " + f"We'll try to handle this gracefully, but it may cause unexpected behavior." + ) + typ = Union[typ, type(field.default)] # type: ignore + + field = dataclasses.replace(field, type_or_callable=typ) + return field + + field_list = list(map(resolve, field_list)) + + return f, type_from_typevar, field_list + + +# Implementation details below. + + +DefaultInstance = Union[ + Any, PropagatingMissingType, NonpropagatingMissingType, ExcludeFromCallType +] + +_known_parsable_types: Set[type] = set( + filter( + lambda x: isinstance(x, Hashable), # type: ignore + itertools.chain( + vars(builtins).values(), # type: ignore + vars(typing).values(), + vars(typing_extensions).values(), + vars(collections.abc).values(), + ), + ) +) + + +def _try_field_list_from_callable( + f: Union[Callable, TypeForm[Any]], + default_instance: DefaultInstance, +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + # Apply constructor_factory override for type. + f = _resolver.swap_type_using_confstruct(f) + + # Check for default instances in subcommand configs. This is needed for + # is_nested_type() when arguments are not valid without a default, and this + # default is specified in the subcommand config. + f, found_subcommand_configs = _resolver.unwrap_annotated( + f, conf._confstruct._SubcommandConfiguration + ) + if len(found_subcommand_configs) > 0: + default_instance = found_subcommand_configs[0].default + + # Unwrap generics. + f, type_from_typevar = _resolver.resolve_generic_types(f) + f = _resolver.apply_type_from_typevar(f, type_from_typevar) + f = _resolver.unwrap_newtype_and_narrow_subtypes(f, default_instance) + f = _resolver.narrow_collection_types(f, default_instance) + f_origin = _resolver.unwrap_origin_strip_extras(cast(TypeForm, f)) + + # If `f` is a type: + # 1. Set cls to the type. + # 2. Consider `f` to be `cls.__init__`. + cls: Optional[TypeForm[Any]] = None + if inspect.isclass(f): + cls = f + if hasattr(cls, "__init__") and cls.__init__ is not object.__init__: + f = cls.__init__ # type: ignore + elif hasattr(cls, "__new__") and cls.__new__ is not object.__new__: + f = cls.__new__ + else: + return UnsupportedNestedTypeMessage( + f"Cannot instantiate class {cls} with no unique __init__ or __new__" + " method." + ) + f_origin = cls # type: ignore + + # Try field generation from class inputs. + if cls is not None: + if is_typeddict(cls): + return _field_list_from_typeddict(cls, default_instance) + if _resolver.is_namedtuple(cls): + return _field_list_from_namedtuple(cls, default_instance) + if _resolver.is_dataclass(cls): + return _field_list_from_dataclass(cls, default_instance) + if _is_attrs(cls): + return _field_list_from_attrs(cls, default_instance) + if _is_pydantic(cls): + return _field_list_from_pydantic(cls, default_instance) + + # Standard container types. These are different because they can be nested structures + # if they contain other nested types (eg Tuple[Struct, Struct]), or treated as + # single arguments otherwise (eg Tuple[int, int]). + # + # Note that f_origin will be populated if we annotate as `Tuple[..]`, and cls will + # be populated if we annotate as just `tuple`. + if f_origin is tuple or cls is tuple: + return _field_list_from_tuple(f, default_instance) + elif f_origin in (collections.abc.Mapping, dict) or cls in ( + collections.abc.Mapping, + dict, + ): + return _field_list_from_dict(f, default_instance) + elif f_origin in (list, set, typing.Sequence, collections.abc.Sequence) or cls in ( + list, + set, + typing.Sequence, + collections.abc.Sequence, + ): + return _field_list_from_nontuple_sequence_checked(f, default_instance) + + # General cases. + if ( + cls is not None and cls in _known_parsable_types + ) or _resolver.unwrap_origin_strip_extras(f) in _known_parsable_types: + return UnsupportedNestedTypeMessage(f"{f} should be parsed directly!") + elif ( + cls is not None + and issubclass(_resolver.unwrap_origin_strip_extras(cls), os.PathLike) + and _instantiators.is_type_string_converter(cls) + ): + return UnsupportedNestedTypeMessage( + f"PathLike {cls} should be parsed directly!" + ) + else: + return _try_field_list_from_general_callable(f, cls, default_instance) + + +def _field_list_from_typeddict( + cls: TypeForm[Any], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + field_list = [] + valid_default_instance = ( + default_instance not in MISSING_SINGLETONS + and default_instance is not EXCLUDE_FROM_CALL + ) + assert not valid_default_instance or isinstance(default_instance, dict) + total = getattr(cls, "__total__", True) + assert isinstance(total, bool) + assert not valid_default_instance or isinstance(default_instance, dict) + for name, typ in _resolver.get_type_hints_with_backported_syntax( + cls, include_extras=True + ).items(): + typ_origin = get_origin(typ) + is_default_from_default_instance = False + if valid_default_instance and name in cast(dict, default_instance): + default = cast(dict, default_instance)[name] + is_default_from_default_instance = True + elif typ_origin is Required and total is False: + # Support total=False. + default = MISSING_PROP + elif total is False: + # Support total=False. + default = EXCLUDE_FROM_CALL + if is_nested_type(typ, MISSING_NONPROP): + # total=False behavior is unideal for nested structures. + pass + # raise _instantiators.UnsupportedTypeAnnotationError( + # "`total=False` not supported for nested structures." + # ) + elif typ_origin is NotRequired: + # Support typing.NotRequired[]. + default = EXCLUDE_FROM_CALL + else: + default = MISSING_PROP + + # Nested types need to be populated / can't be excluded from the call. + if default is EXCLUDE_FROM_CALL and is_nested_type(typ, MISSING_NONPROP): + default = MISSING_NONPROP + + if typ_origin in (Required, NotRequired): + args = get_args(typ) + assert ( + len(args) == 1 + ), "typing.Required[] and typing.NotRequired[T] require a concrete type T." + typ = args[0] + del args + + field_list.append( + FieldDefinition.make( + name=name, + type_or_callable=typ, + default=default, + is_default_from_default_instance=is_default_from_default_instance, + helptext=_docstrings.get_field_docstring(cls, name), + ) + ) + return field_list + + +def _field_list_from_namedtuple( + cls: TypeForm[Any], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + # Handle NamedTuples. + # + # TODO: in terms of helptext, we currently do display the default NamedTuple + # helptext. But we (intentionally) don't for dataclasses; this is somewhat + # inconsistent. + field_list = [] + field_defaults = getattr(cls, "_field_defaults") + + # Note that _field_types is removed in Python 3.9. + for name, typ in _resolver.get_type_hints_with_backported_syntax( + cls, include_extras=True + ).items(): + # Get default, with priority for `default_instance`. + default = field_defaults.get(name, MISSING_NONPROP) + if hasattr(default_instance, name): + default = getattr(default_instance, name) + if default_instance is MISSING_PROP: + default = MISSING_PROP + + field_list.append( + FieldDefinition.make( + name=name, + type_or_callable=typ, + default=default, + is_default_from_default_instance=True, + helptext=_docstrings.get_field_docstring(cls, name), + ) + ) + return field_list + + +def _field_list_from_dataclass( + cls: TypeForm[Any], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + is_flax_module = False + try: + # Check if dataclass is a flax module. This is only possible if flax is already + # loaded. + # + # We generally want to avoid importing flax, since it requires a lot of heavy + # imports. + if "flax.linen" in sys.modules.keys(): + import flax.linen + + if issubclass(cls, flax.linen.Module): + is_flax_module = True + except ImportError: + pass + + # Handle dataclasses. + field_list = [] + for dc_field in filter(lambda field: field.init, _resolver.resolved_fields(cls)): + # For flax modules, we ignore the built-in "name" and "parent" fields. + if is_flax_module and dc_field.name in ("name", "parent"): + continue + + default, is_default_from_default_instance = _get_dataclass_field_default( + dc_field, default_instance + ) + + # Try to get helptext from field metadata. This is also intended to be + # compatible with HuggingFace-style config objects. + helptext = dc_field.metadata.get("help", None) + assert isinstance(helptext, (str, type(None))) + + # Try to get helptext from docstrings. Note that this can't be generated + # dynamically. + if helptext is None: + helptext = _docstrings.get_field_docstring(cls, dc_field.name) + + assert not isinstance(dc_field.type, str) + field_list.append( + FieldDefinition.make( + name=dc_field.name, + type_or_callable=dc_field.type, + default=default, + is_default_from_default_instance=is_default_from_default_instance, + helptext=helptext, + ) + ) + return field_list + + +def _is_pydantic(cls: TypeForm[Any]) -> bool: + # pydantic will already be imported if it's used. + if "pydantic" not in sys.modules.keys(): + # This is needed for the mock import test in + # test_missing_optional_packages.py to pass. + return False + + try: + import pydantic + except ImportError: + return False + + try: + if "pydantic.v1" in sys.modules.keys(): + from pydantic import v1 as pydantic_v1 + else: + pydantic_v1 = None # type: ignore + except ImportError: + if TYPE_CHECKING: + from pydantic import v1 as pydantic_v1 + else: + pydantic_v1 = None + + if issubclass(cls, pydantic.BaseModel): + return True + if pydantic_v1 is not None and issubclass(cls, pydantic_v1.BaseModel): + return True + return False + + +def _field_list_from_pydantic( + cls: TypeForm[Any], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + import pydantic + + try: + if "pydantic.v1" in sys.modules.keys(): + from pydantic import v1 as pydantic_v1 + else: + pydantic_v1 = None # type: ignore + except ImportError: + if TYPE_CHECKING: + from pydantic import v1 as pydantic_v1 + else: + pydantic_v1 = None + + # Handle pydantic models. + field_list = [] + pydantic_version = int(getattr(pydantic, "__version__", "1.0.0").partition(".")[0]) + if pydantic_version < 2 or ( + pydantic_v1 is not None and issubclass(cls, pydantic_v1.BaseModel) + ): + # Pydantic 1.xx. + # We do a conditional cast because the pydantic.v1 module won't + # actually exist in legacy versions of pydantic. + if TYPE_CHECKING: + cls_cast = cast(pydantic_v1.BaseModel, cls) # type: ignore + else: + cls_cast = cls + for pd1_field in cls_cast.__fields__.values(): + helptext = pd1_field.field_info.description + if helptext is None: + helptext = _docstrings.get_field_docstring(cls, pd1_field.name) + + default, is_default_from_default_instance = _get_pydantic_v1_field_default( + pd1_field.name, pd1_field, default_instance + ) + + field_list.append( + FieldDefinition.make( + name=pd1_field.name, + type_or_callable=pd1_field.outer_type_, + default=default, + is_default_from_default_instance=is_default_from_default_instance, + helptext=helptext, + ) + ) + else: + # Pydantic 2.xx. + for name, pd2_field in cast(pydantic.BaseModel, cls).model_fields.items(): + helptext = pd2_field.description + if helptext is None: + helptext = _docstrings.get_field_docstring(cls, name) + + default, is_default_from_default_instance = _get_pydantic_v2_field_default( + name, pd2_field, default_instance + ) + + field_list.append( + FieldDefinition.make( + name=name, + type_or_callable=Annotated.__class_getitem__( # type: ignore + (pd2_field.annotation,) + tuple(pd2_field.metadata) + ) + if len(pd2_field.metadata) > 0 + else pd2_field.annotation, + default=default, + is_default_from_default_instance=is_default_from_default_instance, + helptext=helptext, + ) + ) + return field_list + + +def _is_attrs(cls: TypeForm[Any]) -> bool: + # attr will already be imported if it's used. + if "attr" not in sys.modules.keys(): + return False + + try: + import attr + except ImportError: + # This is needed for the mock import test in + # test_missing_optional_packages.py to pass. + return False + + return attr.has(cls) + + +def _field_list_from_attrs( + cls: TypeForm[Any], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + import attr + + # Resolve forward references in-place, if any exist. + attr.resolve_types(cls) + + # Handle attr classes. + field_list = [] + for attr_field in attr.fields(cls): + # Skip fields with init=False. + if not attr_field.init: + continue + + # Default handling. + name = attr_field.name + default = attr_field.default + is_default_from_default_instance = False + if default_instance not in MISSING_SINGLETONS: + if hasattr(default_instance, name): + default = getattr(default_instance, name) + is_default_from_default_instance = True + else: + warnings.warn( + f"Could not find field {name} in default instance" + f" {default_instance}, which has" + f" type {type(default_instance)},", + stacklevel=2, + ) + elif default is attr.NOTHING: + default = MISSING_NONPROP + elif isinstance(default, attr.Factory): # type: ignore + default = default.factory() # type: ignore + + assert attr_field.type is not None, attr_field + field_list.append( + FieldDefinition.make( + name=name, + type_or_callable=attr_field.type, + default=default, + is_default_from_default_instance=is_default_from_default_instance, + helptext=_docstrings.get_field_docstring(cls, name), + ) + ) + return field_list + + +def _field_list_from_tuple( + f: Union[Callable, TypeForm[Any]], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + # Fixed-length tuples. + field_list: List[FieldDefinition] = [] + children = get_args(f) + if Ellipsis in children: + return _try_field_list_from_sequence_inner( + next(iter(set(children) - {Ellipsis})), default_instance + ) + + # Infer more specific type when tuple annotation isn't subscripted. This generally + # doesn't happen + if len(children) == 0: + if default_instance in MISSING_SINGLETONS: + return UnsupportedNestedTypeMessage( + "If contained types of a tuple are not specified in the annotation, a" + " default instance must be specified." + ) + else: + assert isinstance(default_instance, tuple) + children = tuple(type(x) for x in default_instance) + + if ( + default_instance in MISSING_SINGLETONS + # EXCLUDE_FROM_CALL indicates we're inside a TypedDict, with total=False. + or default_instance is EXCLUDE_FROM_CALL + ): + default_instance = (default_instance,) * len(children) + + for i, child in enumerate(children): + default_i = default_instance[i] # type: ignore + field_list.append( + FieldDefinition.make( + # Ideally we'd have --tuple[0] instead of --tuple.0 as the command-line + # argument, but in practice the brackets are annoying because they + # require escaping. + name=str(i), + type_or_callable=child, + default=default_i, + is_default_from_default_instance=True, + helptext="", + # This should really set the positional marker, but the CLI is more + # intuitive for mixed nested/non-nested types in tuples when we stick + # with kwargs. Tuples are special-cased in _calling.py. + ) + ) + + contains_nested = False + for field in field_list: + if get_origin(field.type_or_callable) is Union: + for option in get_args(field.type_or_callable): + # The second argument here is the default value, which can help with + # narrowing but is generall not necessary. + contains_nested |= is_nested_type(option, MISSING_NONPROP) + + contains_nested |= is_nested_type(field.type_or_callable, field.default) + + if contains_nested: + break + if not contains_nested: + # We could also check for variable length children, which can be populated when + # the tuple is interpreted as a nested field but not a directly parsed one. + return UnsupportedNestedTypeMessage( + "Tuple does not contain any nested structures." + ) + + return field_list + + +def _field_list_from_nontuple_sequence_checked( + f: Union[Callable, TypeForm[Any]], default_instance: DefaultInstance +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + """Tuples are handled differently due to variadics (tuple[T1, T2, T3, ..., Tn]), + while list[], sequence[], set[], etc only have one typevar.""" + contained_type: Any + if len(get_args(f)) == 0: + # A raw collection type (list, tuple, etc) was passed in, but narrowing also failed. + assert ( + default_instance in MISSING_SINGLETONS + or len(cast(typing.Sequence, default_instance)) == 0 + ), f"Default instance {default_instance} for type {f} was unexpected!" + return UnsupportedNestedTypeMessage( + f"Sequence type {f} needs either an explicit type or a" + " default to infer from." + ) + else: + (contained_type,) = get_args(f) + return _try_field_list_from_sequence_inner(contained_type, default_instance) + + +def _try_field_list_from_sequence_inner( + contained_type: TypeForm[Any], + default_instance: DefaultInstance, +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + # When no default instance is specified: + # If we have List[int] => this can be parsed as a single field. + # If we have List[SomeStruct] => OK. + if default_instance in MISSING_SINGLETONS and not is_nested_type( + contained_type, MISSING_NONPROP + ): + return UnsupportedNestedTypeMessage( + f"Sequence containing type {contained_type} should be parsed directly!" + ) + + # If we have a default instance: + # [int, int, int] => this can be parsed as a single field. + # [SomeStruct, int, int] => OK. + if isinstance(default_instance, Iterable) and all( + [not is_nested_type(type(x), x) for x in default_instance] + ): + return UnsupportedNestedTypeMessage( + f"Sequence with default {default_instance} should be parsed directly!" + ) + if default_instance in MISSING_SINGLETONS: + # We use the broader error type to prevent it from being caught by + # is_possibly_nested_type(). This is for sure a bad annotation! + raise _instantiators.UnsupportedTypeAnnotationError( + "tyro currently only supports fixed-length sequences of nested types. For " + "variable-length sequences over nested types, we need a default value to " + "infer length from. You can also consider a custom constructor, see here " + "for an example: https://github.com/brentyi/tyro/issues/151" + ) + + field_list = [] + for i, default_i in enumerate(default_instance): # type: ignore + field_list.append( + FieldDefinition.make( + name=str(i), + type_or_callable=contained_type, + default=default_i, + is_default_from_default_instance=True, + helptext="", + ) + ) + return field_list + + +def _field_list_from_dict( + f: Union[Callable, TypeForm[Any]], + default_instance: DefaultInstance, +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + if default_instance in MISSING_SINGLETONS or len(cast(dict, default_instance)) == 0: + return UnsupportedNestedTypeMessage( + "Nested dictionary structures must have non-empty default instance specified." + ) + field_list = [] + for k, v in cast(dict, default_instance).items(): + field_list.append( + FieldDefinition.make( + name=str(k) if not isinstance(k, enum.Enum) else k.name, + type_or_callable=type(v), + default=v, + is_default_from_default_instance=True, + helptext=None, + # Dictionary specific key: + call_argname_override=k, + ) + ) + return field_list + + +def _try_field_list_from_general_callable( + f: Union[Callable, TypeForm[Any]], + cls: Optional[TypeForm[Any]], + default_instance: DefaultInstance, +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + # Generate field list from function signature. + if not callable(f): + return UnsupportedNestedTypeMessage( + f"Cannot extract annotations from {f}, which is not a callable type." + ) + params = list(inspect.signature(f).parameters.values()) + if cls is not None: + # Ignore self parameter. + params = params[1:] + + out = _field_list_from_params(f, cls, params) + if isinstance(out, UnsupportedNestedTypeMessage): + # Return error message. + return out + + # If a default is provided: either all or none of the arguments must be passed in. + if default_instance not in MISSING_SINGLETONS: + for i, field in enumerate(out): + out[i] = field.add_markers((_markers._OPTIONAL_GROUP,)) + + return out + + +def _field_list_from_params( + f: Union[Callable, TypeForm[Any]], + cls: Optional[TypeForm[Any]], + params: List[inspect.Parameter], +) -> Union[List[FieldDefinition], UnsupportedNestedTypeMessage]: + # Unwrap functools.wraps and functools.partial. + done = False + while not done: + done = True + if hasattr(f, "__wrapped__"): + f = f.__wrapped__ # type: ignore + done = False + if isinstance(f, functools.partial): + f = f.func + done = False + + # Sometime functools.* is applied to a class. + if inspect.isclass(f): + cls = f + f = f.__init__ # type: ignore + + # Get type annotations, docstrings. + docstring = inspect.getdoc(f) + docstring_from_arg_name = {} + if docstring is not None: + for param_doc in docstring_parser.parse(docstring).params: + docstring_from_arg_name[param_doc.arg_name] = param_doc.description + del docstring + + # This will throw a type error for torch.device, typing.Dict, etc. + try: + hints = _resolver.get_type_hints_with_backported_syntax(f, include_extras=True) + except TypeError: + return UnsupportedNestedTypeMessage(f"Could not get hints for {f}!") + + field_list = [] + for param in params: + # Get default value. + default = param.default + + # Get helptext from docstring. + helptext = docstring_from_arg_name.get(param.name) + if helptext is None and cls is not None: + helptext = _docstrings.get_field_docstring(cls, param.name) + + if param.name not in hints: + out = UnsupportedNestedTypeMessage( + f"Expected fully type-annotated callable, but {f} with arguments" + f" {tuple(map(lambda p: p.name, params))} has no annotation for" + f" '{param.name}'." + ) + if param.kind is param.KEYWORD_ONLY: + # If keyword only: this can't possibly be an instantiator function + # either, so we escalate to an error. + raise _instantiators.UnsupportedTypeAnnotationError(out.message) + return out + + # Set markers for positional + variadic arguments. + markers: Tuple[Any, ...] = () + typ: Any = hints[param.name] + if param.kind is inspect.Parameter.POSITIONAL_ONLY: + markers = (_markers.Positional, _markers._PositionalCall) + elif param.kind is inspect.Parameter.VAR_POSITIONAL: + # Handle *args signatures. + # + # This will create a `--args T [T ...]` CLI argument. + markers = (_markers._UnpackArgsCall,) + typ = Tuple.__getitem__((typ, ...)) # type: ignore + default = () + elif param.kind is inspect.Parameter.VAR_KEYWORD: + # Handle *kwargs signatures. + # + # This will create a `--kwargs STR T [STR T ...]` CLI argument. + # + # Note that it would be straightforward to make both this and *args truly + # positional, omitting the --args/--kwargs prefix, but we are choosing not + # to because it would make *args and **kwargs difficult to use in + # conjunction. + markers = (_markers._UnpackKwargsCall,) + typ = Dict.__getitem__((str, typ)) # type: ignore + default = {} + + field_list.append( + FieldDefinition.make( + name=param.name, + # Note that param.annotation doesn't resolve forward references. + type_or_callable=typ, + default=default, + is_default_from_default_instance=False, + helptext=helptext, + markers=markers, + ) + ) + + return field_list + + +def _ensure_dataclass_instance_used_as_default_is_frozen( + field: dataclasses.Field, default_instance: Any +) -> None: + """Ensure that a dataclass type used directly as a default value is marked as + frozen.""" + assert dataclasses.is_dataclass(default_instance) + cls = type(default_instance) + if not cls.__dataclass_params__.frozen: # type: ignore + warnings.warn( + f"Mutable type {cls} is used as a default value for `{field.name}`. This is" + " dangerous! Consider using `dataclasses.field(default_factory=...)` or" + f" marking {cls} as frozen." + ) + + +def _get_dataclass_field_default( + field: dataclasses.Field, parent_default_instance: Any +) -> Tuple[Any, bool]: + """Helper for getting the default instance for a dataclass field.""" + # If the dataclass's parent is explicitly marked MISSING, mark this field as missing + # as well. + if parent_default_instance is MISSING_PROP: + return MISSING_PROP, False + + # Try grabbing default from parent instance. + if ( + parent_default_instance not in MISSING_SINGLETONS + and parent_default_instance is not None + ): + # Populate default from some parent, eg `default=` in `tyro.cli()`. + if hasattr(parent_default_instance, field.name): + return getattr(parent_default_instance, field.name), True + else: + warnings.warn( + f"Could not find field {field.name} in default instance" + f" {parent_default_instance}, which has" + f" type {type(parent_default_instance)},", + stacklevel=2, + ) + + # Try grabbing default from dataclass field. + if field.default not in MISSING_SINGLETONS: + default = field.default + # Note that dataclasses.is_dataclass() will also return true for dataclass + # _types_, not just instances. + if type(default) is not type and dataclasses.is_dataclass(default): + _ensure_dataclass_instance_used_as_default_is_frozen(field, default) + return default, False + + # Populate default from `dataclasses.field(default_factory=...)`. + if field.default_factory is not dataclasses.MISSING and not ( + # Special case to ignore default_factory if we write: + # `field: Dataclass = dataclasses.field(default_factory=Dataclass)`. + # + # In other words, treat it the same way as: `field: Dataclass`. + # + # The only time this matters is when we our dataclass has a `__post_init__` + # function that mutates the dataclass. We choose here to use the default values + # before this method is called. + dataclasses.is_dataclass(field.type) and field.default_factory is field.type + ): + return field.default_factory(), False + + # Otherwise, no default. This is different from MISSING, because MISSING propagates + # to children. We could revisit this design to make it clearer. + return MISSING_NONPROP, False + + +if TYPE_CHECKING: + import pydantic as pydantic + import pydantic.v1.fields as pydantic_v1_fields + + +def _get_pydantic_v1_field_default( + name: str, + field: pydantic_v1_fields.ModelField, + parent_default_instance: DefaultInstance, +) -> Tuple[Any, bool]: + """Helper for getting the default instance for a Pydantic field.""" + + # Try grabbing default from parent instance. + if ( + parent_default_instance not in MISSING_SINGLETONS + and parent_default_instance is not None + ): + # Populate default from some parent, eg `default=` in `tyro.cli()`. + if hasattr(parent_default_instance, name): + return getattr(parent_default_instance, name), True + else: + warnings.warn( + f"Could not find field {name} in default instance" + f" {parent_default_instance}, which has" + f" type {type(parent_default_instance)},", + stacklevel=2, + ) + + if not field.required: + return field.get_default(), False + + # Otherwise, no default. + return MISSING_NONPROP, False + + +def _get_pydantic_v2_field_default( + name: str, + field: pydantic.fields.FieldInfo, + parent_default_instance: DefaultInstance, +) -> Tuple[Any, bool]: + """Helper for getting the default instance for a Pydantic field.""" + + import pydantic + + # Try grabbing default from parent instance. + if ( + parent_default_instance not in MISSING_SINGLETONS + and parent_default_instance is not None + ): + # Populate default from some parent, eg `default=` in `tyro.cli()`. + if hasattr(parent_default_instance, name): + return getattr(parent_default_instance, name), True + else: + warnings.warn( + f"Could not find field {name} in default instance" + f" {parent_default_instance}, which has" + f" type {type(parent_default_instance)},", + stacklevel=2, + ) + + if not field.is_required(): + try: + return field.get_default(call_default_factory=True), False + except pydantic.ValidationError: + pass + + # Otherwise, no default. + return MISSING_NONPROP, False diff --git a/src/tyro/_instantiators.py b/src/tyro/_instantiators.py new file mode 100644 index 000000000..c5b2f50d8 --- /dev/null +++ b/src/tyro/_instantiators.py @@ -0,0 +1,750 @@ +"""Helper for using type annotations to recursively generate instantiator functions, +which map sequences of strings to the annotated type. + +Some examples of type annotations and the desired instantiators: +``` + int + + lambda strings: int(str[0]) + + List[int] + + lambda strings: list( + [int(x) for x in strings] + ) + + List[Color], where Color is an enum + + lambda strings: list( + [Color[x] for x in strings] + ) + + Tuple[int, float] + + lambda strings: tuple( + typ(x) + for typ, x in zip( + (int, float), + strings, + ) + ) +``` +""" + +import builtins +import collections.abc +import dataclasses +import datetime +import enum +import inspect +import os +import pathlib +from collections import deque +from typing import ( + Any, + Callable, + Dict, + Hashable, + Iterable, + List, + Optional, + Sequence, + Set, + Tuple, + TypeVar, + Union, + cast, + overload, +) + +from typing_extensions import Annotated, Literal, get_args, get_origin + +from . import _resolver + +# There are cases where typing.Literal doesn't match typing_extensions.Literal: +# https://github.com/python/typing_extensions/pull/148 +try: + from typing import Literal as LiteralAlternate +except ImportError: + LiteralAlternate = Literal # type: ignore + +from . import _strings +from ._typing import TypeForm +from .conf import _markers + +_StandardInstantiator = Callable[[List[str]], Any] +_AppendNargsInstantiator = Callable[[List[List[str]]], Any] +# Special case: the only time that argparse doesn't give us a string is when the +# argument action is set to `store_true` or `store_false`. In this case, we get +# a bool directly, and the field action can be a no-op. +_FlagInstantiator = Callable[[bool], bool] + +Instantiator = Union[_StandardInstantiator, _AppendNargsInstantiator, _FlagInstantiator] + +NoneType = type(None) + + +@dataclasses.dataclass +class InstantiatorMetadata: + # Unlike in vanilla argparse, we never set nargs to None. To make things simpler, we + # instead use nargs=1. + nargs: Optional[Union[int, Literal["*"]]] + # Unlike in vanilla argparse, our metavar is always a string. We handle + # sequences, multiple arguments, etc, manually. + metavar: str + choices: Optional[Tuple[str, ...]] + action: Optional[Literal["append"]] + + def check_choices(self, strings: List[str]) -> None: + if self.choices is not None and any(s not in self.choices for s in strings): + raise ValueError(f"invalid choice: {strings} (choose from {self.choices}))") + + +class UnsupportedTypeAnnotationError(Exception): + """Exception raised when an unsupported type annotation is detected.""" + + +_builtin_set: Set[Hashable] = set( + filter( + lambda x: isinstance(x, Hashable), # type: ignore + vars(builtins).values(), # type: ignore + ) +) + + +def is_type_string_converter(typ: Union[Callable, TypeForm[Any]]) -> bool: + """Check if type is a string converter, i.e., (arg: Union[str, Any]) -> T.""" + param_count = 0 + has_var_positional = False + try: + signature = inspect.signature(typ) + except ValueError: + # pybind objects might not have a parsable signature. We try to be tolerant in this case. + # One day this should be fixed with `__text_signature__`. + return True + + type_annotations = _resolver.get_type_hints_with_backported_syntax(typ) + # Some checks we can do if the signature is available! + for i, param in enumerate(signature.parameters.values()): + annotation = type_annotations.get(param.name, param.annotation) + + # Hack: apply_type_from_typevar applies shims, like UnionType => Union + # conversion. + annotation = _resolver.apply_type_from_typevar(annotation, {}) + if i == 0 and not ( + (get_origin(annotation) is Union and str in get_args(annotation)) + or annotation in (str, inspect.Parameter.empty) + ): + return False + if param.kind is inspect.Parameter.VAR_POSITIONAL: + has_var_positional = True + elif param.default is inspect.Parameter.empty and param.kind is not ( + inspect.Parameter.VAR_KEYWORD + ): + param_count += 1 + + # Raise an error if parameters look wrong. + if not (param_count == 1 or (param_count == 0 and has_var_positional)): + return False + return True + + +def instantiator_from_type( + typ: Union[TypeForm[Any], Callable], + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: + """Recursive helper for parsing type annotations. + + Returns two things: + - An instantiator function, for instantiating the type from a string or list of + strings. The latter applies when argparse's `nargs` parameter is set. + - A metadata structure, which specifies parameters for argparse. + """ + + # Handle Any. + if typ is Any: + raise UnsupportedTypeAnnotationError("`Any` is not a parsable type.") + + # Handle NoneType. + if typ is NoneType: + + def instantiator(strings: List[str]) -> None: + # Note that other inputs should be caught by `choices` before the + # instantiator runs. + assert strings == ["None"] + + return instantiator, InstantiatorMetadata( + nargs=1, + metavar="{None}", + choices=("None",), + action=None, + ) + + # Instantiate os.PathLike annotations using pathlib.Path. + # Ideally this should be implemented in a more general way, eg using + # https://github.com/brentyi/tyro/pull/30 + if typ is os.PathLike: + typ = pathlib.Path + + # Unwrap NewType + set metavar based on NewType name. + # `isinstance(x, NewType)` doesn't work because NewType isn't a class until + # Python 3.10, so we instead do a duck typing-style check. + metavar = getattr(typ, "__name__", "").upper() + typ, maybe_newtype_name = _resolver.unwrap_newtype_and_aliases(typ) + if maybe_newtype_name is not None: + metavar = maybe_newtype_name.upper() + + typ = _resolver.unwrap_annotated(typ) + + # Address container types. If a matching container is found, this will recursively + # call instantiator_from_type(). + container_out = _instantiator_from_container_type( + cast(TypeForm[Any], typ), type_from_typevar, markers + ) + if container_out is not None: + return container_out + + # At this point: Annotated[] and containers like list[T] are all handled, + # any types with a non-None origin aren't supported. + if get_origin(typ) is not None: + raise UnsupportedTypeAnnotationError( # pragma: no cover + f"Unsupported type {typ} with origin {get_origin(typ)}" + ) + + # Validate that typ is a `(arg: str) -> T` type converter, as expected by argparse. + if typ in _builtin_set: + pass + elif not callable(typ): + raise UnsupportedTypeAnnotationError( + f"Expected {typ} to be an `(arg: str) -> T` type converter, but is not" + " callable." + ) + elif not is_type_string_converter(typ): + raise UnsupportedTypeAnnotationError( + f"Expected {typ} to be an `(arg: str) -> T` type converter, but is not" + " a valid type converter." + ) + + # Use ISO 8601 standard for dates/times. + if typ in (datetime.datetime, datetime.date, datetime.time): + + def instantiate(args: List[str]) -> Any: + (arg,) = args + try: + # Type ignore is unnecessary for pyright but needed for mypy. + return typ.fromisoformat(arg) # type: ignore + except ValueError: + raise ValueError( + f"`{typ.__name__}.fromisoformat('{arg}')` failed. Dates " + "should be specified in ISO-8601 format: " + "https://en.wikipedia.org/wiki/ISO_8601" + ) + + return ( + instantiate, + InstantiatorMetadata( + nargs=1, + metavar={ + # Actually we take any ISO 8601 string here. So these + # metavars are a bit incomplete. + # datetime.datetime: "YYYY-MM-DD[THH:MM:SS]", + # datetime.date: "YYYY-MM-DD", + # datetime.time: "HH:MM[:SS]", + # + # More complete. + datetime.datetime: "YYYY-MM-DD[THH:MM:[SS[…]]]", + datetime.date: "YYYY-MM-DD", + datetime.time: "HH:MM[:SS[…]]", + # + # Even more complete! + # datetime.datetime: "YYYY-MM-DD[THH:MM[:SS[.fff]][±HH:MM|Z]]", + # datetime.date: "YYYY-MM-DD", + # datetime.time: "HH:MM[:SS[.fff]][±HH:MM|Z]", + # + # Not very informative but (1) precise and (2) succinct. + # datetime.datetime: "ISO-DATETIME", + # datetime.date: "ISO-DATE", + # datetime.time: "ISO-TIME", + }[cast(Any, typ)], # cast is for mypy, pyright works fine + choices=None, + action=None, + ), + ) + + # Special case `choices` for some types, as implemented in `instance_from_string()`. + auto_choices: Optional[Tuple[str, ...]] = None + if typ is bool: + auto_choices = ("True", "False") + elif inspect.isclass(typ) and issubclass(typ, enum.Enum): + # Using `.__members__` dict to handle aliases correctly + auto_choices = tuple(typ.__members__.keys()) + + def instantiator_base_case(strings: List[str]) -> Any: + """Given a type and and a string from the command-line, reconstruct an object. Not + intended to deal with containers. + + This is intended to replace all calls to `type(string)`, which can cause unexpected + behavior. As an example, note that the following argparse code will always print + `True`, because `bool("True") == bool("False") == bool("0") == True`. + ``` + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--flag", type=bool) + + print(parser.parse_args().flag) + ``` + """ + assert len(get_args(typ)) == 0, f"TypeForm {typ} cannot be instantiated." + (string,) = strings + if typ is bool: + return {"True": True, "False": False}[string] # type: ignore + elif isinstance(typ, type) and issubclass(typ, enum.Enum): + return typ[string] # type: ignore + elif typ is bytes: + return bytes(string, encoding="ascii") # type: ignore + else: + # We assume this base type can be called on a string to convert + # from a string, eg int("5"), float("5."), Path("/home"), etc. + return typ(string) # type: ignore + + return instantiator_base_case, InstantiatorMetadata( + nargs=1, + metavar=( + metavar + if auto_choices is None + else "{" + ",".join(map(str, auto_choices)) + "}" + ), + choices=None if auto_choices is None else auto_choices, + action=None, + ) + + +@overload +def _instantiator_from_type_inner( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + allow_sequences: Literal["fixed_length"], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: ... + + +@overload +def _instantiator_from_type_inner( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + allow_sequences: Literal[False], + markers: Set[_markers.Marker], +) -> Tuple[_StandardInstantiator, InstantiatorMetadata]: ... + + +@overload +def _instantiator_from_type_inner( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + allow_sequences: Literal[True], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: ... + + +def _instantiator_from_type_inner( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + allow_sequences: Literal["fixed_length", True, False], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: + """Thin wrapper over instantiator_from_type, with some extra asserts for catching + errors.""" + out = instantiator_from_type(typ, type_from_typevar, markers) + if out[1].nargs == "*": + # We currently only use allow_sequences=False for options in Literal types, + # which are evaluated using `type()`. It should not be possible to hit this + # condition from polling a runtime type. + assert allow_sequences + if allow_sequences == "fixed_length" and not isinstance(out[1].nargs, int): + raise UnsupportedTypeAnnotationError( + f"{typ} is a variable-length sequence, which is ambiguous when nested." + " For nesting variable-length sequences (example: List[List[int]])," + " `tyro.conf.UseAppendAction` can help resolve ambiguities." + ) + return out + + +def _instantiator_from_container_type( + typ: TypeForm[Any], + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Optional[Tuple[Instantiator, InstantiatorMetadata]]: + """Attempt to create an instantiator from a container type. Returns `None` if no + container type is found.""" + + # Default generic types to strings. + if typ in (dict, Dict): + typ = Dict[str, str] + elif typ in (tuple, Tuple): + typ = Tuple[str, ...] # type: ignore + elif typ in (list, List, collections.abc.Sequence, Sequence): + typ = List[str] + elif typ in (set, Set): + typ = Set[str] + + type_origin = get_origin(typ) + if type_origin in (None, Annotated): + return None + + for make, matched_origins in { + _instantiator_from_sequence: ( + collections.abc.Sequence, + frozenset, + list, + set, + deque, + ), + _instantiator_from_tuple: (tuple,), + _instantiator_from_dict: (dict, collections.abc.Mapping), + _instantiator_from_union: (Union,), + _instantiator_from_literal: (Literal, LiteralAlternate), + }.items(): + if type_origin in matched_origins: + return make(typ, type_from_typevar, markers) + return None + + +def _instantiator_from_tuple( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: + types = get_args(typ) + typeset = set(types) # Note that sets are unordered. + typeset_no_ellipsis = typeset - {Ellipsis} # type: ignore + + if typeset_no_ellipsis != typeset: + # Ellipsis: variable argument counts. When an ellipsis is used, tuples must + # contain only one type. + assert len(typeset_no_ellipsis) == 1 + return _instantiator_from_sequence(typ, type_from_typevar, markers) + + else: + instantiators: List[_StandardInstantiator] = [] + metas: List[InstantiatorMetadata] = [] + nargs = 0 + for t in types: + a, b = _instantiator_from_type_inner( + t, type_from_typevar, allow_sequences="fixed_length", markers=markers + ) + instantiators.append(a) # type: ignore + metas.append(b) + assert isinstance(b.nargs, int) + nargs += b.nargs + + def fixed_length_tuple_instantiator(strings: List[str]) -> Any: + assert len(strings) == nargs + + # Make tuple. + out = [] + i = 0 + for make, meta in zip(instantiators, metas): + assert isinstance(meta.nargs, int) + meta.check_choices(strings[i : i + meta.nargs]) + out.append(make(strings[i : i + meta.nargs])) + i += meta.nargs + return tuple(out) + + return fixed_length_tuple_instantiator, InstantiatorMetadata( + nargs=nargs, + metavar=" ".join(m.metavar for m in metas), + choices=None, + action=None, + ) + + +def _join_union_metavars(metavars: Iterable[str]) -> str: + """Metavar generation helper for unions. Could be revisited. + + Examples: + None, INT => NONE|INT + {0,1,2}, {3,4} => {0,1,2,3,4} + {0,1,2}, {3,4}, STR => {0,1,2,3,4}|STR + {None}, INT [INT ...] => {None}|{INT [INT ...]} + STR, INT [INT ...] => STR|{INT [INT ...]} + STR, INT INT => STR|{INT INT} + + The curly brackets are unfortunately overloaded but alternatives all interfere with + argparse internals. + """ + metavars = tuple(metavars) + merged_metavars = [metavars[0]] + for i in range(1, len(metavars)): + prev = merged_metavars[-1] + curr = metavars[i] + if ( + prev.startswith("{") + and prev.endswith("}") + and curr.startswith("{") + and curr.endswith("}") + ): + merged_metavars[-1] = prev[:-1] + "," + curr[1:] + else: + merged_metavars.append(curr) + + for i, m in enumerate(merged_metavars): + if " " in m: + merged_metavars[i] = "{" + m + "}" + + return "|".join(merged_metavars) + + +def _instantiator_from_union( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: + options = list(get_args(typ)) + if NoneType in options: + # Move `None` types to the beginning. + # If we have `Optional[str]`, we want this to be parsed as + # `Union[NoneType, str]`. + options.remove(NoneType) + options.insert(0, NoneType) + + # General unions, eg Union[int, bool]. We'll try to convert these from left to + # right. + instantiators = [] + metas = [] + choices: Tuple[str, ...] | None = () + nargs: Optional[Union[int, Literal["*"]]] = 1 + first = True + for t in options: + a, b = _instantiator_from_type_inner( + t, type_from_typevar, allow_sequences=True, markers=markers + ) + instantiators.append(a) + metas.append(b) + if b.choices is None: + choices = None + elif choices is not None: + choices = choices + b.choices + + if t is not NoneType: + # Enforce that `nargs` is the same for all child types, except for + # NoneType. + if first: + nargs = b.nargs + first = False + elif nargs != b.nargs: + # Just be as general as possible if we see inconsistencies. + nargs = "*" + + metavar: str + metavar = _join_union_metavars(map(lambda x: cast(str, x.metavar), metas)) + + def union_instantiator(strings: List[str]) -> Any: + metadata: InstantiatorMetadata + errors = [] + for i, (instantiator, metadata) in enumerate(zip(instantiators, metas)): + # Check choices. + if metadata.choices is not None and any( + x not in metadata.choices for x in strings + ): + errors.append( + f"{options[i]}: {strings} does not match choices {metadata.choices}" + ) + continue + + # Try passing input into instantiator. + if len(strings) == metadata.nargs or (metadata.nargs == "*"): + try: + return instantiator(strings) # type: ignore + except ValueError as e: + # Failed, try next instantiator. + errors.append(f"{options[i]}: {e.args[0]}") + else: + errors.append( + f"{options[i]}: input length {len(strings)} did not match expected" + f" argument count {metadata.nargs}" + ) + raise ValueError( + f"no type in {options} could be instantiated from" + f" {strings}.\n\nGot errors: \n- " + "\n- ".join(errors) + ) + + return union_instantiator, InstantiatorMetadata( + nargs=nargs, + metavar=metavar, + choices=None if choices is None else tuple(set(choices)), + action=None, + ) + + +def _instantiator_from_dict( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: + key_type, val_type = get_args(typ) + key_instantiator, key_meta = _instantiator_from_type_inner( + key_type, type_from_typevar, allow_sequences="fixed_length", markers=markers + ) + + if _markers.UseAppendAction in markers: + val_instantiator, val_meta = _instantiator_from_type_inner( + val_type, + type_from_typevar, + allow_sequences=True, + markers=markers - {_markers.UseAppendAction}, + ) + pair_metavar = f"{key_meta.metavar} {val_meta.metavar}" + key_nargs = cast(int, key_meta.nargs) # Casts needed for mypy but not pyright! + val_nargs = val_meta.nargs + assert isinstance(key_nargs, int) + + def append_dict_instantiator(strings: List[List[str]]) -> Any: + out = {} + for s in strings: + out[key_instantiator(s[:key_nargs])] = val_instantiator(s[key_nargs:]) # type: ignore + return out + + return append_dict_instantiator, InstantiatorMetadata( + nargs=key_nargs + val_nargs if isinstance(val_nargs, int) else "*", + metavar=pair_metavar, + choices=None, + action="append", + ) + else: + val_instantiator, val_meta = _instantiator_from_type_inner( + val_type, type_from_typevar, allow_sequences="fixed_length", markers=markers + ) + pair_metavar = f"{key_meta.metavar} {val_meta.metavar}" + key_nargs = cast(int, key_meta.nargs) # Casts needed for mypy but not pyright! + val_nargs = cast(int, val_meta.nargs) + assert isinstance(key_nargs, int) + assert isinstance(val_nargs, int) + pair_nargs = key_nargs + val_nargs + + def dict_instantiator(strings: List[str]) -> Any: + out = {} + if len(strings) % pair_nargs != 0: + raise ValueError("incomplete set of key value pairs!") + + index = 0 + for _ in range(len(strings) // pair_nargs): + assert isinstance(key_nargs, int) + assert isinstance(val_nargs, int) + k = strings[index : index + key_nargs] + index += key_nargs + v = strings[index : index + val_nargs] + index += val_nargs + + if key_meta.choices is not None and any( + kj not in key_meta.choices for kj in k + ): + raise ValueError( + f"invalid choice: {k} (choose from {key_meta.choices}))" + ) + if val_meta.choices is not None and any( + vj not in val_meta.choices for vj in v + ): + raise ValueError( + f"invalid choice: {v} (choose from {val_meta.choices}))" + ) + out[key_instantiator(k)] = val_instantiator(v) # type: ignore + return out + + return dict_instantiator, InstantiatorMetadata( + nargs="*", + metavar=_strings.multi_metavar_from_single(pair_metavar), + choices=None, + action=None, + ) + + +def _instantiator_from_sequence( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Tuple[Instantiator, InstantiatorMetadata]: + """Instantiator for variable-length sequences: list, sets, Tuple[T, ...], etc.""" + container_type = get_origin(typ) + assert container_type is not None + if container_type is collections.abc.Sequence: + container_type = list + + if container_type is tuple: + (contained_type, ell) = get_args(typ) + assert ell == Ellipsis + else: + (contained_type,) = get_args(typ) + + if _markers.UseAppendAction in markers: + make, inner_meta = _instantiator_from_type_inner( + contained_type, + type_from_typevar, + allow_sequences=True, + markers=markers - {_markers.UseAppendAction}, + ) + + def append_sequence_instantiator(strings: List[List[str]]) -> Any: + assert strings is not None + return container_type(cast(_StandardInstantiator, make)(s) for s in strings) + + return append_sequence_instantiator, InstantiatorMetadata( + nargs=inner_meta.nargs, + metavar=inner_meta.metavar, + choices=inner_meta.choices, + action="append", + ) + else: + make, inner_meta = _instantiator_from_type_inner( + contained_type, + type_from_typevar, + allow_sequences="fixed_length", + markers=markers, + ) + + def sequence_instantiator(strings: List[str]) -> Any: + # Validate nargs. + if ( + isinstance(inner_meta.nargs, int) + and len(strings) % inner_meta.nargs != 0 + ): + raise ValueError( + f"input {strings} is of length {len(strings)}, which is not" + f" divisible by {inner_meta.nargs}." + ) + + # Make tuple. + out = [] + step = inner_meta.nargs if isinstance(inner_meta.nargs, int) else 1 + for i in range(0, len(strings), step): + out.append(make(strings[i : i + inner_meta.nargs])) # type: ignore + assert container_type is not None + return container_type(out) + + return sequence_instantiator, InstantiatorMetadata( + nargs="*", + metavar=_strings.multi_metavar_from_single(inner_meta.metavar), + choices=inner_meta.choices, + action=None, + ) + + +def _instantiator_from_literal( + typ: TypeForm, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + markers: Set[_markers.Marker], +) -> Tuple[_StandardInstantiator, InstantiatorMetadata]: + choices = get_args(typ) + str_choices = tuple(x.name if isinstance(x, enum.Enum) else str(x) for x in choices) + return ( + # Note that if string is not in str_choices, it will be caught from setting + # `choices` below. + lambda strings: choices[str_choices.index(strings[0])], + InstantiatorMetadata( + nargs=1, + metavar="{" + ",".join(str_choices) + "}", + choices=str_choices, + action=None, + ), + ) diff --git a/src/tyro/_parsers.py b/src/tyro/_parsers.py new file mode 100644 index 000000000..39d469f9f --- /dev/null +++ b/src/tyro/_parsers.py @@ -0,0 +1,660 @@ +"""Interface for generating `argparse.ArgumentParser()` definitions from callables.""" + +from __future__ import annotations + +import dataclasses +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Set, + Tuple, + Type, + TypeVar, + Union, + cast, +) + +from typing_extensions import Annotated, get_args, get_origin + +from . import _argparse as argparse +from . import ( + _argparse_formatter, + _arguments, + _docstrings, + _fields, + _instantiators, + _resolver, + _strings, + _subcommand_matching, +) +from ._typing import TypeForm +from .conf import _confstruct, _markers + +T = TypeVar("T") + + +@dataclasses.dataclass(frozen=True) +class ParserSpecification: + """Each parser contains a list of arguments and optionally some subparsers.""" + + f: Callable + description: str + args: List[_arguments.ArgumentDefinition] + field_list: List[_fields.FieldDefinition] + child_from_prefix: Dict[str, ParserSpecification] + helptext_from_intern_prefixed_field_name: Dict[str, Optional[str]] + + # We have two mechanics for tracking subparser groups: + # - A single subparser group, which is what gets added in the tree structure built + # by the argparse parser. + subparsers: Optional[SubparsersSpecification] + # - A set of subparser groups, which reflect the tree structure built by the + # hierarchy of a nested config structure. + subparsers_from_intern_prefix: Dict[str, SubparsersSpecification] + intern_prefix: str + extern_prefix: str + has_required_args: bool + consolidate_subcommand_args: bool + + @staticmethod + def from_callable_or_type( + f: Callable[..., T], + description: Optional[str], + parent_classes: Set[Type[Any]], + default_instance: Union[ + T, _fields.PropagatingMissingType, _fields.NonpropagatingMissingType + ], + intern_prefix: str, + extern_prefix: str, + subcommand_prefix: str = "", + support_single_arg_types: bool = False, + ) -> ParserSpecification: + """Create a parser definition from a callable or type.""" + + # Consolidate subcommand types. + markers = _resolver.unwrap_annotated(f, _markers._Marker)[1] + consolidate_subcommand_args = _markers.ConsolidateSubcommandArgs in markers + + # Resolve the type of `f`, generate a field list. + f, type_from_typevar, field_list = _fields.field_list_from_callable( + f=f, + default_instance=default_instance, + support_single_arg_types=support_single_arg_types, + ) + for i in range(len(field_list)): + field_list[i].markers |= set(markers) + + # Cycle detection. + # + # Note that 'parent' here refers to in the nesting hierarchy, not the + # superclass. + if f in parent_classes and f is not dict: + raise _instantiators.UnsupportedTypeAnnotationError( + f"Found a cyclic dependency with type {f}." + ) + + # TODO: we are abusing the (minor) distinctions between types, classes, and + # callables throughout the code. This is mostly for legacy reasons, could be + # cleaned up. + parent_classes = parent_classes | {cast(Type, f)} + + has_required_args = False + args = [] + helptext_from_intern_prefixed_field_name: Dict[str, Optional[str]] = {} + + child_from_prefix: Dict[str, ParserSpecification] = {} + + subparsers = None + subparsers_from_prefix = {} + + for field in field_list: + field_out = handle_field( + field, + type_from_typevar=type_from_typevar, + parent_classes=parent_classes, + intern_prefix=intern_prefix, + extern_prefix=extern_prefix, + subcommand_prefix=subcommand_prefix, + ) + if isinstance(field_out, _arguments.ArgumentDefinition): + # Handle single arguments. + args.append(field_out) + if field_out.lowered.required: + has_required_args = True + elif isinstance(field_out, SubparsersSpecification): + # Handle subparsers. + subparsers_from_prefix[field_out.intern_prefix] = field_out + subparsers = add_subparsers_to_leaves(subparsers, field_out) + elif isinstance(field_out, ParserSpecification): + # Handle nested parsers. + nested_parser = field_out + child_from_prefix[field_out.intern_prefix] = nested_parser + + if nested_parser.has_required_args: + has_required_args = True + + # Include nested subparsers. + if nested_parser.subparsers is not None: + subparsers_from_prefix.update( + nested_parser.subparsers_from_intern_prefix + ) + subparsers = add_subparsers_to_leaves( + subparsers, nested_parser.subparsers + ) + + # Helptext for this field; used as description for grouping arguments. + class_field_name = _strings.make_field_name([field.intern_name]) + if field.helptext is not None: + helptext_from_intern_prefixed_field_name[class_field_name] = ( + field.helptext + ) + else: + helptext_from_intern_prefixed_field_name[class_field_name] = ( + _docstrings.get_callable_description(nested_parser.f) + ) + + # If arguments are in an optional group, it indicates that the default_instance + # will be used if none of the arguments are passed in. + if ( + len(nested_parser.args) >= 1 + and _markers._OPTIONAL_GROUP in nested_parser.args[0].field.markers + ): + current_helptext = helptext_from_intern_prefixed_field_name[ + class_field_name + ] + helptext_from_intern_prefixed_field_name[class_field_name] = ( + ("" if current_helptext is None else current_helptext + "\n\n") + + "Default: " + + str(field.default) + ) + + return ParserSpecification( + f=f, + description=_strings.remove_single_line_breaks( + description + if description is not None + else _docstrings.get_callable_description(f) + ), + args=args, + field_list=field_list, + child_from_prefix=child_from_prefix, + helptext_from_intern_prefixed_field_name=helptext_from_intern_prefixed_field_name, + subparsers=subparsers, + subparsers_from_intern_prefix=subparsers_from_prefix, + intern_prefix=intern_prefix, + extern_prefix=extern_prefix, + has_required_args=has_required_args, + consolidate_subcommand_args=consolidate_subcommand_args, + ) + + def apply( + self, parser: argparse.ArgumentParser + ) -> Tuple[argparse.ArgumentParser, ...]: + """Create defined arguments and subparsers.""" + + # Generate helptext. + parser.description = self.description + + # Create subparser tree. + subparser_group = None + if self.subparsers is not None: + leaves = self.subparsers.apply(parser) + subparser_group = parser._action_groups.pop() + else: + leaves = (parser,) + + # Depending on whether we want to consolidate subcommand args, we can either + # apply arguments to the intermediate parser or only on the leaves. + if self.consolidate_subcommand_args: + for leaf in leaves: + self.apply_args(leaf) + else: + self.apply_args(parser) + + if subparser_group is not None: + parser._action_groups.append(subparser_group) + + # Break some API boundaries to rename the "optional arguments" => "options". + assert parser._action_groups[1].title in ( + # python <= 3.9 + "optional arguments", + # python >= 3.10 + "options", + ) + parser._action_groups[1].title = "options" + + return leaves + + def apply_args( + self, + parser: argparse.ArgumentParser, + parent: Optional[ParserSpecification] = None, + ) -> None: + """Create defined arguments and subparsers.""" + + # Make argument groups. + def format_group_name(prefix: str) -> str: + return (prefix + " options").strip() + + group_from_prefix: Dict[str, argparse._ArgumentGroup] = { + "": parser._action_groups[1], + **{ + cast(str, group.title).partition(" ")[0]: group + for group in parser._action_groups[2:] + }, + } + positional_group = parser._action_groups[0] + assert positional_group.title == "positional arguments" + + # Add each argument group. Note that groups with only suppressed arguments won't + # be added. + for arg in self.args: + if ( + arg.lowered.help is not argparse.SUPPRESS + and arg.extern_prefix not in group_from_prefix + ): + description = ( + parent.helptext_from_intern_prefixed_field_name.get( + arg.intern_prefix + ) + if parent is not None + else None + ) + group_from_prefix[arg.extern_prefix] = parser.add_argument_group( + format_group_name(arg.extern_prefix), + description=description, + ) + + # Add each argument. + for arg in self.args: + if arg.field.is_positional(): + arg.add_argument(positional_group) + continue + + if arg.extern_prefix in group_from_prefix: + arg.add_argument(group_from_prefix[arg.extern_prefix]) + else: + # Suppressed argument: still need to add them, but they won't show up in + # the helptext so it doesn't matter which group. + assert arg.lowered.help is argparse.SUPPRESS + arg.add_argument(group_from_prefix[""]) + + for child in self.child_from_prefix.values(): + child.apply_args(parser, parent=self) + + +def handle_field( + field: _fields.FieldDefinition, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + parent_classes: Set[Type[Any]], + intern_prefix: str, + extern_prefix: str, + subcommand_prefix: str, +) -> Union[ + _arguments.ArgumentDefinition, + ParserSpecification, + SubparsersSpecification, +]: + """Determine what to do with a single field definition.""" + + if isinstance(field.type_or_callable, TypeVar): + raise _instantiators.UnsupportedTypeAnnotationError( + f"Field {field.intern_name} has an unbound TypeVar: {field.type_or_callable}." + ) + + if _markers.Fixed not in field.markers and _markers.Suppress not in field.markers: + # (1) Handle Unions over callables; these result in subparsers. + subparsers_attempt = SubparsersSpecification.from_field( + field, + type_from_typevar=type_from_typevar, + parent_classes=parent_classes, + intern_prefix=_strings.make_field_name([intern_prefix, field.intern_name]), + extern_prefix=_strings.make_field_name([extern_prefix, field.extern_name]), + ) + if subparsers_attempt is not None: + if ( + not subparsers_attempt.required + and _markers.AvoidSubcommands in field.markers + ): + # Don't make a subparser. + field = dataclasses.replace(field, type_or_callable=type(field.default)) + else: + return subparsers_attempt + + # (2) Handle nested callables. + if _fields.is_nested_type(field.type_or_callable, field.default): + field = dataclasses.replace( + field, + type_or_callable=_resolver.unwrap_newtype_and_narrow_subtypes( + field.type_or_callable, + field.default, + ), + ) + return ParserSpecification.from_callable_or_type( + ( + # Recursively apply marker types. + field.type_or_callable + if len(field.markers) == 0 + else Annotated.__class_getitem__( # type: ignore + (field.type_or_callable,) + tuple(field.markers) + ) + ), + description=None, + parent_classes=parent_classes, + default_instance=field.default, + intern_prefix=_strings.make_field_name( + [intern_prefix, field.intern_name] + ), + extern_prefix=_strings.make_field_name( + [extern_prefix, field.extern_name] + ) + if field.argconf.prefix_name in (True, None) + else field.extern_name, + subcommand_prefix=subcommand_prefix, + support_single_arg_types=False, + ) + + # (3) Handle primitive or fixed types. These produce a single argument! + return _arguments.ArgumentDefinition( + intern_prefix=intern_prefix, + extern_prefix=extern_prefix, + subcommand_prefix=subcommand_prefix, + field=field, + type_from_typevar=type_from_typevar, + ) + + +@dataclasses.dataclass(frozen=True) +class SubparsersSpecification: + """Structure for defining subparsers. Each subparser is a parser with a name.""" + + name: str + description: Optional[str] + parser_from_name: Dict[str, ParserSpecification] + intern_prefix: str + required: bool + default_instance: Any + options: Tuple[Union[TypeForm[Any], Callable], ...] + + @staticmethod + def from_field( + field: _fields.FieldDefinition, + type_from_typevar: Dict[TypeVar, TypeForm[Any]], + parent_classes: Set[Type[Any]], + intern_prefix: str, + extern_prefix: str, + ) -> Optional[SubparsersSpecification]: + # Union of classes should create subparsers. + typ = _resolver.unwrap_annotated(field.type_or_callable) + if get_origin(typ) is not Union: + return None + + # We don't use sets here to retain order of subcommands. + options: List[Union[type, Callable]] + options = [ + _resolver.apply_type_from_typevar(typ, type_from_typevar) + for typ in get_args(typ) + ] + options = [ + ( + # Cast seems unnecessary but needed in mypy... (1.4.1) + cast(Callable, none_proxy) if o is type(None) else o + ) + for o in options + ] + + # If specified, swap types using tyro.conf.subcommand(constructor=...). + for i, option in enumerate(options): + _, found_subcommand_configs = _resolver.unwrap_annotated( + option, _confstruct._SubcommandConfiguration + ) + if ( + len(found_subcommand_configs) > 0 + and found_subcommand_configs[0].constructor_factory is not None + ): + options[i] = Annotated.__class_getitem__( # type: ignore + ( + found_subcommand_configs[0].constructor_factory(), + *_resolver.unwrap_annotated(option, "all")[1], + ) + ) + + # Exit if we don't contain any nested types. + if not any( + [ + o is not none_proxy + and _fields.is_nested_type(cast(type, o), _fields.MISSING_NONPROP) + for o in options + ] + ): + return None + + # Get subcommand configurations from `tyro.conf.subcommand()`. + subcommand_config_from_name: Dict[ + str, _confstruct._SubcommandConfiguration + ] = {} + subcommand_type_from_name: Dict[str, type] = {} + for option in options: + subcommand_name = _strings.subparser_name_from_type( + "" + if _markers.OmitSubcommandPrefixes in field.markers + else extern_prefix, + type(None) if option is none_proxy else cast(type, option), + ) + option_unwrapped, found_subcommand_configs = _resolver.unwrap_annotated( + option, _confstruct._SubcommandConfiguration + ) + if len(found_subcommand_configs) != 0: + # Explicitly annotated default. + assert len(found_subcommand_configs) == 1, ( + f"Expected only one subcommand config, but {subcommand_name} has" + f" {len(found_subcommand_configs)}." + ) + subcommand_config_from_name[subcommand_name] = found_subcommand_configs[ + 0 + ] + subcommand_type_from_name[subcommand_name] = cast(type, option) + + # If a field default is provided, try to find a matching subcommand name. + if field.default is None or field.default in _fields.MISSING_SINGLETONS: + default_name = None + else: + default_name = _subcommand_matching.match_subcommand( + field.default, subcommand_config_from_name, subcommand_type_from_name + ) + + # This should never be triggered because union types are expanded when + # a bad default value is provided. + assert default_name is not None, ( + f"`{extern_prefix}` was provided a default value of type" + f" {type(field.default)} but no matching subcommand was found. A" + " type may be missing in the Union type declaration for" + f" `{extern_prefix}`, which currently expects {options}." + ) + + # Add subcommands for each option. + parser_from_name: Dict[str, ParserSpecification] = {} + for option in options: + subcommand_name = _strings.subparser_name_from_type( + "" + if _markers.OmitSubcommandPrefixes in field.markers + else extern_prefix, + type(None) if option is none_proxy else cast(type, option), + ) + + # Get a subcommand config: either pulled from the type annotations or the + # field default. + if subcommand_name in subcommand_config_from_name: + subcommand_config = subcommand_config_from_name[subcommand_name] + else: + subcommand_config = _confstruct._SubcommandConfiguration( + "unused", + description=None, + default=_fields.MISSING_NONPROP, + prefix_name=True, + constructor_factory=None, + ) + + # If names match, borrow subcommand default from field default. + if default_name == subcommand_name and ( + field.is_default_from_default_instance + or subcommand_config.default in _fields.MISSING_SINGLETONS + ): + subcommand_config = dataclasses.replace( + subcommand_config, default=field.default + ) + + # Strip the subcommand config from the option type. + # Relevant: https://github.com/brentyi/tyro/pull/117 + option_origin, annotations = _resolver.unwrap_annotated(option, "all") + annotations = tuple( + a + for a in annotations + if not isinstance(a, _confstruct._SubcommandConfiguration) + ) + if len(annotations) == 0: + option = option_origin + else: + option = Annotated.__class_getitem__( # type: ignore + (option_origin,) + annotations + ) + + subparser = ParserSpecification.from_callable_or_type( + ( + # Recursively apply markers. + Annotated.__class_getitem__((option,) + tuple(field.markers)) # type: ignore + if len(field.markers) > 0 + else option + ), + description=subcommand_config.description, + parent_classes=parent_classes, + default_instance=subcommand_config.default, + intern_prefix=intern_prefix, + extern_prefix=extern_prefix, + subcommand_prefix=intern_prefix, + support_single_arg_types=True, + ) + + # Apply prefix to helptext in nested classes in subparsers. + subparser = dataclasses.replace( + subparser, + helptext_from_intern_prefixed_field_name={ + _strings.make_field_name([intern_prefix, k]): v + for k, v in subparser.helptext_from_intern_prefixed_field_name.items() + }, + ) + parser_from_name[subcommand_name] = subparser + + # Required if a default is missing. + required = field.default in _fields.MISSING_SINGLETONS + + # Required if a default is passed in, but the default value has missing + # parameters. + if default_name is not None: + default_parser = parser_from_name[default_name] + if any(map(lambda arg: arg.lowered.required, default_parser.args)): + required = True + if ( + default_parser.subparsers is not None + and default_parser.subparsers.required + ): + required = True + + # Required if all args are pushed to the final subcommand. + if _markers.ConsolidateSubcommandArgs in field.markers: + required = True + + # Make description. + description_parts = [] + if field.helptext is not None: + description_parts.append(field.helptext) + if not required and field.default not in _fields.MISSING_SINGLETONS: + description_parts.append(f" (default: {default_name})") + description = ( + # We use `None` instead of an empty string to prevent a line break from + # being created where the description would be. + " ".join(description_parts) if len(description_parts) > 0 else None + ) + + return SubparsersSpecification( + name=field.intern_name, + # If we wanted, we could add information about the default instance + # automatically, as is done for normal fields. But for now we just rely on + # the user to include it in the docstring. + description=description, + parser_from_name=parser_from_name, + intern_prefix=intern_prefix, + required=required, + default_instance=field.default, + options=tuple(options), + ) + + def apply( + self, parent_parser: argparse.ArgumentParser + ) -> Tuple[argparse.ArgumentParser, ...]: + title = "subcommands" + metavar = "{" + ",".join(self.parser_from_name.keys()) + "}" + if not self.required: + title = "optional " + title + metavar = f"[{metavar}]" + + # Add subparsers to every node in previous level of the tree. + argparse_subparsers = parent_parser.add_subparsers( + dest=_strings.make_subparser_dest(self.intern_prefix), + description=self.description, + required=self.required, + title=title, + metavar=metavar, + ) + + subparser_tree_leaves: List[argparse.ArgumentParser] = [] + for name, subparser_def in self.parser_from_name.items(): + helptext = subparser_def.description.replace("%", "%%") + if len(helptext) > 0: + # TODO: calling a private function here. + helptext = _arguments._rich_tag_if_enabled(helptext.strip(), "helptext") + + subparser = argparse_subparsers.add_parser( + name, + formatter_class=_argparse_formatter.TyroArgparseHelpFormatter, + help=helptext, + allow_abbrev=False, + ) + + # Attributes used for error message generation. + assert isinstance(subparser, _argparse_formatter.TyroArgumentParser) + assert isinstance(parent_parser, _argparse_formatter.TyroArgumentParser) + subparser._parsing_known_args = parent_parser._parsing_known_args + subparser._parser_specification = parent_parser._parser_specification + subparser._console_outputs = parent_parser._console_outputs + subparser._args = parent_parser._args + + subparser_tree_leaves.extend(subparser_def.apply(subparser)) + + return tuple(subparser_tree_leaves) + + +def add_subparsers_to_leaves( + root: Optional[SubparsersSpecification], leaf: SubparsersSpecification +) -> SubparsersSpecification: + if root is None: + return leaf + + new_parsers_from_name = {} + for name, parser in root.parser_from_name.items(): + new_parsers_from_name[name] = dataclasses.replace( + parser, + subparsers=add_subparsers_to_leaves(parser.subparsers, leaf), + ) + return dataclasses.replace( + root, + parser_from_name=new_parsers_from_name, + required=root.required or leaf.required, + ) + + +def none_proxy() -> None: + return None diff --git a/src/tyro/_resolver.py b/src/tyro/_resolver.py new file mode 100644 index 000000000..46b610ddc --- /dev/null +++ b/src/tyro/_resolver.py @@ -0,0 +1,528 @@ +"""Utilities for resolving types and forward references.""" + +import collections.abc +import copy +import dataclasses +import inspect +import sys +import types +import warnings +from typing import ( + Any, + Callable, + ClassVar, + Dict, + FrozenSet, + List, + Optional, + Sequence, + Set, + Tuple, + TypeVar, + Union, + cast, + overload, +) + +from typing_extensions import ( + Annotated, + Final, + ForwardRef, + Literal, + Self, + TypeAliasType, + get_args, + get_origin, + get_type_hints, +) + +from . import _fields, _unsafe_cache, conf +from ._typing import TypeForm + +TypeOrCallable = TypeVar("TypeOrCallable", TypeForm[Any], Callable) + + +def unwrap_origin_strip_extras(typ: TypeOrCallable) -> TypeOrCallable: + """Returns the origin, ignoring typing.Annotated, of typ if it exists. Otherwise, + returns typ.""" + # TODO: Annotated[] handling should be revisited... + typ = unwrap_annotated(typ) + origin = get_origin(typ) + + if origin is not None: + typ = origin + + return typ + + +def is_dataclass(cls: Union[TypeForm, Callable]) -> bool: + """Same as `dataclasses.is_dataclass`, but also handles generic aliases.""" + return dataclasses.is_dataclass(unwrap_origin_strip_extras(cls)) # type: ignore + + +def resolve_generic_types( + cls: TypeOrCallable, +) -> Tuple[TypeOrCallable, Dict[TypeVar, TypeForm[Any]]]: + """If the input is a class: no-op. If it's a generic alias: returns the origin + class, and a mapping from typevars to concrete types.""" + + annotations: Tuple[Any, ...] = () + if get_origin(cls) is Annotated: + # ^We need this `if` statement for an obscure edge case: when `cls` is a + # function with `__tyro_markers__` set, we don't want/need to return + # Annotated[func, markers]. + cls, annotations = unwrap_annotated(cls, "all") + + # We'll ignore NewType when getting the origin + args for generics. + origin_cls = get_origin(unwrap_newtype_and_aliases(cls)[0]) + type_from_typevar: Dict[TypeVar, TypeForm[Any]] = {} + + # Support typing.Self. + # We'll do this by pretending that `Self` is a TypeVar... + if hasattr(cls, "__self__"): + self_type = getattr(cls, "__self__") + if inspect.isclass(self_type): + type_from_typevar[cast(TypeVar, Self)] = self_type # type: ignore + else: + type_from_typevar[cast(TypeVar, Self)] = self_type.__class__ # type: ignore + + if ( + # Apply some heuristics for generic types. Should revisit this. + origin_cls is not None + and hasattr(origin_cls, "__parameters__") + and hasattr(origin_cls.__parameters__, "__len__") + ): + typevars = origin_cls.__parameters__ + typevar_values = get_args(unwrap_newtype_and_aliases(cls)[0]) + assert len(typevars) == len(typevar_values) + cls = origin_cls + type_from_typevar.update(dict(zip(typevars, typevar_values))) + + if hasattr(cls, "__orig_bases__"): + bases = getattr(cls, "__orig_bases__") + for base in bases: + origin_base = unwrap_origin_strip_extras(base) + if origin_base is base or not hasattr(origin_base, "__parameters__"): + continue + typevars = origin_base.__parameters__ + typevar_values = get_args(base) + type_from_typevar.update(dict(zip(typevars, typevar_values))) + + if len(annotations) == 0: + return cls, type_from_typevar + else: + return ( + Annotated.__class_getitem__((cls, *annotations)), # type: ignore + type_from_typevar, + ) + + +@_unsafe_cache.unsafe_cache(maxsize=1024) +def resolved_fields(cls: TypeForm) -> List[dataclasses.Field]: + """Similar to dataclasses.fields(), but includes dataclasses.InitVar types and + resolves forward references.""" + + assert dataclasses.is_dataclass(cls) + fields = [] + annotations = get_type_hints_with_backported_syntax( + cast(Callable, cls), include_extras=True + ) + for field in getattr(cls, "__dataclass_fields__").values(): + # Avoid mutating original field. + field = copy.copy(field) + + # Resolve forward references. + field.type = annotations[field.name] + + # Skip ClassVars. + if get_origin(field.type) is ClassVar: + continue + + # Unwrap InitVar types. + if isinstance(field.type, dataclasses.InitVar): + field.type = field.type.type + + fields.append(field) + + return fields + + +def is_namedtuple(cls: TypeForm) -> bool: + return ( + hasattr(cls, "_fields") + # `_field_types` was removed in Python >=3.9. + # and hasattr(cls, "_field_types") + and hasattr(cls, "_field_defaults") + ) + + +def type_from_typevar_constraints(typ: TypeOrCallable) -> TypeOrCallable: + """Try to concretize a type from a TypeVar's bounds or constraints. Identity if + unsuccessful.""" + if isinstance(typ, TypeVar): + if typ.__bound__ is not None: + # Try to infer type from TypeVar bound. + return typ.__bound__ + elif len(typ.__constraints__) > 0: + # Try to infer type from TypeVar constraints. + return Union.__getitem__(typ.__constraints__) # type: ignore + return typ + + +TypeOrCallableOrNone = TypeVar("TypeOrCallableOrNone", Callable, TypeForm[Any], None) + + +def unwrap_newtype_and_aliases( + typ: TypeOrCallableOrNone, +) -> Tuple[TypeOrCallableOrNone, Optional[str]]: + # Handle type aliases, eg via the `type` statement in Python 3.12. + if isinstance(typ, TypeAliasType): + return unwrap_newtype_and_aliases(typ.__value__) # type: ignore + + # We'll unwrap NewType annotations here; this is needed before issubclass + # checks! + # + # `isinstance(x, NewType)` doesn't work because NewType isn't a class until + # Python 3.10, so we instead do a duck typing-style check. + return_name = None + while hasattr(typ, "__name__") and hasattr(typ, "__supertype__"): + if return_name is None: + return_name = getattr(typ, "__name__") + typ = getattr(typ, "__supertype__") + + return typ, return_name + + +@_unsafe_cache.unsafe_cache(maxsize=1024) +def unwrap_newtype_and_narrow_subtypes( + typ: TypeOrCallable, + default_instance: Any, +) -> TypeOrCallable: + """Type narrowing: if we annotate as Animal but specify a default instance of Cat, + we should parse as Cat. + + This should generally only be applied to fields used as nested structures, not + individual arguments/fields. (if a field is annotated as Union[int, str], and a + string default is passed in, we don't want to narrow the type to always be + strings!)""" + + typ, unused_name = unwrap_newtype_and_aliases(typ) + del unused_name + + try: + potential_subclass = type(default_instance) + + if potential_subclass is type: + # Don't narrow to `type`. This happens when the default instance is a class; + # it doesn't really make sense to parse this case. + return typ + + superclass = unwrap_annotated(typ) + + # For Python 3.10. + if get_origin(superclass) is Union: + return typ + + if superclass is Any or issubclass(potential_subclass, superclass): # type: ignore + if get_origin(typ) is Annotated: + return Annotated.__class_getitem__( # type: ignore + (potential_subclass,) + get_args(typ)[1:] + ) + typ = cast(TypeOrCallable, potential_subclass) + except TypeError: + # TODO: document where this TypeError can be raised, and reduce the amount of + # code in it. + pass + + return typ + + +def swap_type_using_confstruct(typ: TypeOrCallable) -> TypeOrCallable: + """Swap types using the `constructor_factory` attribute from + `tyro.conf.arg` and `tyro.conf.subcommand`. Runtime annotations are + kept, but the type is swapped.""" + # Need to swap types. + _, annotations = unwrap_annotated(typ, search_type="all") + for anno in reversed(annotations): + if ( + isinstance( + anno, + ( + conf._confstruct._ArgConfiguration, + conf._confstruct._SubcommandConfiguration, + ), + ) + and anno.constructor_factory is not None + ): + return Annotated.__class_getitem__( # type: ignore + (anno.constructor_factory(),) + annotations + ) + return typ + + +def narrow_collection_types( + typ: TypeOrCallable, default_instance: Any +) -> TypeOrCallable: + """TypeForm narrowing for containers. Infers types of container contents.""" + args = get_args(typ) + origin = get_origin(typ) + if args == (Any,) or (origin is tuple and args == (Any, Ellipsis)): + typ = origin # type: ignore + + if typ in (list, Sequence, collections.abc.Sequence) and isinstance( + default_instance, list + ): + if len(default_instance) == 0: + return typ + typ = List.__getitem__(Union.__getitem__(tuple(map(type, default_instance)))) # type: ignore + elif typ in (set, Sequence, collections.abc.Sequence) and isinstance( + default_instance, set + ): + if len(default_instance) == 0: + return typ + typ = Set.__getitem__(Union.__getitem__(tuple(map(type, default_instance)))) # type: ignore + elif typ in (tuple, Sequence, collections.abc.Sequence) and isinstance( + default_instance, tuple + ): + if len(default_instance) == 0: + return typ + typ = Tuple.__getitem__(tuple(map(type, default_instance))) # type: ignore + return cast(TypeOrCallable, typ) + + +# `Final` and `ReadOnly` types are ignored in tyro. +try: + # Can only import ReadOnly in typing_extensions>=4.9.0, which isn't + # supported by Python 3.7. + from typing_extensions import ReadOnly + + STRIP_WRAPPER_TYPES = {Final, ReadOnly} +except ImportError: + STRIP_WRAPPER_TYPES = {Final} + +MetadataType = TypeVar("MetadataType") + + +@overload +def unwrap_annotated( + typ: TypeOrCallable, + search_type: TypeForm[MetadataType], +) -> Tuple[TypeOrCallable, Tuple[MetadataType, ...]]: ... + + +@overload +def unwrap_annotated( + typ: TypeOrCallable, + search_type: Literal["all"], +) -> Tuple[TypeOrCallable, Tuple[Any, ...]]: ... + + +@overload +def unwrap_annotated( + typ: TypeOrCallable, + search_type: None = None, +) -> TypeOrCallable: ... + + +def unwrap_annotated( + typ: TypeOrCallable, + search_type: Union[TypeForm[MetadataType], Literal["all"], object, None] = None, +) -> Union[Tuple[TypeOrCallable, Tuple[MetadataType, ...]], TypeOrCallable]: + """Helper for parsing typing.Annotated types. + + Examples: + - int, int => (int, ()) + - Annotated[int, 1], int => (int, (1,)) + - Annotated[int, "1"], int => (int, ()) + """ + + # `Final` and `ReadOnly` types are ignored in tyro. + while get_origin(typ) in STRIP_WRAPPER_TYPES: + typ = get_args(typ)[0] + + # Don't search for any annotations. + if search_type is None: + if not hasattr(typ, "__metadata__"): + return typ + else: + return get_args(typ)[0] + + # Check for __tyro_markers__ from @configure. + if hasattr(typ, "__tyro_markers__"): + if search_type == "all": + targets = getattr(typ, "__tyro_markers__") + else: + targets = tuple( + x + for x in getattr(typ, "__tyro_markers__") + if isinstance(x, search_type) # type: ignore + ) + else: + targets = () + + assert isinstance(targets, tuple) + if not hasattr(typ, "__metadata__"): + return typ, targets # type: ignore + + args = get_args(typ) + assert len(args) >= 2 + + # Look through metadata for desired metadata type. + targets += tuple( + x + for x in targets + args[1:] + if search_type == "all" or isinstance(x, search_type) # type: ignore + ) + + # Check for __tyro_markers__ in unwrapped type. + if hasattr(args[0], "__tyro_markers__"): + targets += tuple( + x + for x in getattr(args[0], "__tyro_markers__") + if search_type == "all" or isinstance(x, search_type) # type: ignore + ) + return args[0], targets # type: ignore + + +def apply_type_from_typevar( + typ: TypeOrCallable, type_from_typevar: Dict[TypeVar, TypeForm[Any]] +) -> TypeOrCallable: + if typ in type_from_typevar: + return type_from_typevar[typ] # type: ignore + + origin = get_origin(typ) + args = get_args(typ) + if len(args) > 0: + if origin is Annotated: + args = args[:1] + if origin is collections.abc.Callable: + assert isinstance(args[0], list) + args = tuple(args[0]) + args[1:] + + # Convert Python 3.9 and 3.10 types to their typing library equivalents, which + # support `.copy_with()`. + if sys.version_info[:2] >= (3, 9): + shim_table = { + # PEP 585. Requires Python 3.9. + tuple: Tuple, + list: List, + dict: Dict, + set: Set, + frozenset: FrozenSet, + } + if hasattr(types, "UnionType"): # type: ignore + # PEP 604. Requires Python 3.10. + shim_table[types.UnionType] = Union # type: ignore + + for new, old in shim_table.items(): + if isinstance(typ, new) or origin is new: # type: ignore + typ = old.__getitem__(args) # type: ignore + + new_args = tuple(apply_type_from_typevar(x, type_from_typevar) for x in args) + + # Standard generic aliases have a `copy_with()`! + if hasattr(typ, "copy_with"): + return typ.copy_with(new_args) # type: ignore + else: + # `collections` types, like collections.abc.Sequence. + assert hasattr(origin, "__class_getitem__") + return origin.__class_getitem__(new_args) # type: ignore + + return typ + + +@_unsafe_cache.unsafe_cache(maxsize=1024) +def narrow_union_type(typ: TypeOrCallable, default_instance: Any) -> TypeOrCallable: + """Narrow union types. + + This is a shim for failing more gracefully when we we're given a Union type that + doesn't match the default value. + + In this case, we raise a warning, then add the type of the default value to the + union. Loosely motivated by: https://github.com/brentyi/tyro/issues/20 + """ + if get_origin(typ) is not Union: + return typ + + options = get_args(typ) + options_unwrapped = [unwrap_origin_strip_extras(o) for o in options] + + # (A) + try: + if default_instance not in _fields.MISSING_SINGLETONS and not any( + isinstance(default_instance, o) for o in options_unwrapped + ): + warnings.warn( + f"{type(default_instance)} does not match any type in Union:" + f" {options_unwrapped}" + ) + return Union.__getitem__(options + (type(default_instance),)) # type: ignore + except TypeError: + pass + + return typ + + +NoneType = type(None) + + +def get_type_hints_with_backported_syntax( + obj: Callable[..., Any], include_extras: bool = False +) -> Dict[str, Any]: + """Same as `typing.get_type_hints()`, but supports new union syntax (X | Y) + and generics (list[str]) in older versions of Python.""" + try: + out = get_type_hints(obj, include_extras=include_extras) + + # Workaround for: + # - https://github.com/brentyi/tyro/issues/156 + # - https://github.com/python/cpython/issues/90353 + # + # Which impacts Python 3.10 and earlier. + # + # It may be possible to remove this if this issue is resolved: + # - https://github.com/python/typing_extensions/issues/310 + if sys.version_info < (3, 11): + # If we see Optional[Annotated[T, ...]], we're going to flip to Annotated[Optional[T]]... + # + # It's unlikely but possible for this to have unintended side effects. + for k, v in out.items(): + origin = get_origin(v) + args = get_args(v) + if ( + origin is Union + and len(args) == 2 + and (args[0] is NoneType or args[1] is NoneType) + ): + non_none = args[1] if args[0] is NoneType else args[0] + if get_origin(non_none) is Annotated: + annotated_args = get_args(non_none) + out[k] = Annotated.__class_getitem__( # type: ignore + ( + Union.__getitem__((annotated_args[0], None)), # type: ignore + *annotated_args[1:], + ) + ) + return out + + except TypeError as e: # pragma: no cover + # Resolve new type syntax using eval_type_backport. + if hasattr(obj, "__annotations__"): + try: + from eval_type_backport import eval_type_backport + + # Get global namespace for functions. + globalns = getattr(obj, "__globals__", None) + + # Get global namespace for classes. + if globalns is None and hasattr(globalns, "__init__"): + globalns = getattr(getattr(obj, "__init__"), "__globals__", None) + + out = { + k: eval_type_backport(ForwardRef(v), globalns=globalns, localns={}) + for k, v in getattr(obj, "__annotations__").items() + } + return out + except ImportError: + pass + raise e diff --git a/src/tyro/_singleton.py b/src/tyro/_singleton.py new file mode 100644 index 000000000..f52efefd1 --- /dev/null +++ b/src/tyro/_singleton.py @@ -0,0 +1,13 @@ +class Singleton: + # Singleton pattern. + # https://www.python.org/download/releases/2.2/descrintro/#__new__ + def __new__(cls, *args, **kwds): + it = cls.__dict__.get("__it__") + if it is not None: + return it + cls.__it__ = it = object.__new__(cls) + it.init(*args, **kwds) + return it + + def init(self, *args, **kwds): + pass diff --git a/src/tyro/_strings.py b/src/tyro/_strings.py new file mode 100644 index 000000000..e2fc7c09f --- /dev/null +++ b/src/tyro/_strings.py @@ -0,0 +1,179 @@ +"""Utilities and constants for working with strings.""" + +import contextlib +import functools +import re +import textwrap +from typing import List, Sequence, Tuple, Type + +from typing_extensions import Literal, get_args, get_origin + +from . import _resolver + +dummy_field_name = "__tyro_dummy_field__" +DELIMETER: Literal["-", "_"] = "-" + + +@contextlib.contextmanager +def delimeter_context(delimeter: Literal["-", "_"]): + """Context for setting the delimeter. Determines if `field_a` is populated as + `--field-a` or `--field_a`. Not thread-safe.""" + global DELIMETER + delimeter_restore = DELIMETER + DELIMETER = delimeter + yield + DELIMETER = delimeter_restore + + +def get_delimeter() -> Literal["-", "_"]: + """Get delimeter used to separate words.""" + return DELIMETER + + +def replace_delimeter_in_part(p: str) -> str: + """Replace hyphens with underscores (or vice versa) except when at the start.""" + if get_delimeter() == "-": + stripped = p.lstrip("_") + p = p[: len(p) - len(stripped)] + stripped.replace("_", "-") + else: + p = p.replace("-", "_") + return p + + +def make_field_name(parts: Sequence[str]) -> str: + """Join parts of a field name together. Used for nesting. + + ('parent_1', 'child') => 'parent-1.child' + ('parents', '1', '_child_node') => 'parents.1._child-node' + ('parents', '1', 'middle._child_node') => 'parents.1.middle._child-node' + """ + out = ".".join(parts) + return ".".join( + replace_delimeter_in_part(part) + for part in out.split(".") + if len(part) > 0 and part != dummy_field_name + ) + + +def make_subparser_dest(name: str) -> str: + return f"{name} (positional)" + + +def dedent(text: str) -> str: + """Same as textwrap.dedent, but ignores the first line.""" + first_line, line_break, rest = text.partition("\n") + if line_break == "": + return textwrap.dedent(text) + return f"{first_line.strip()}\n{textwrap.dedent(rest)}" + + +def hyphen_separated_from_camel_case(name: str) -> str: + out = ( + re.compile("((?<=[a-z0-9])[A-Z]|(?!^)[A-Z](?=[a-z]))") + .sub(get_delimeter() + r"\1", name) + .lower() + ) + return out + + +def _subparser_name_from_type(cls: Type) -> Tuple[str, bool]: + from .conf import _confstruct # Prevent circular imports + + cls, type_from_typevar = _resolver.resolve_generic_types(cls) + cls, found_subcommand_configs = _resolver.unwrap_annotated( + cls, _confstruct._SubcommandConfiguration + ) + + # Subparser name from `tyro.metadata.subcommand()`. + found_name = None + prefix_name = True + if len(found_subcommand_configs) > 0: + found_name = found_subcommand_configs[0].name + prefix_name = found_subcommand_configs[0].prefix_name + + if found_name is not None: + return found_name, prefix_name + + # Subparser name from class name. + def get_name(cls: Type) -> str: + orig = get_origin(cls) + if orig is not None and hasattr(orig, "__name__"): + parts = [orig.__name__] # type: ignore + parts.extend(map(get_name, get_args(cls))) + parts = [hyphen_separated_from_camel_case(part) for part in parts] + return get_delimeter().join(parts) + elif hasattr(cls, "__name__"): + return hyphen_separated_from_camel_case(cls.__name__) + else: + return hyphen_separated_from_camel_case(str(cls)) + + if len(type_from_typevar) == 0: + return get_name(cls), prefix_name # type: ignore + + return ( + get_delimeter().join( + map( + lambda x: _subparser_name_from_type(x)[0], + [cls] + list(type_from_typevar.values()), + ) + ), + prefix_name, + ) + + +def subparser_name_from_type(prefix: str, cls: Type) -> str: + suffix, use_prefix = ( + _subparser_name_from_type(cls) if cls is not type(None) else ("None", True) + ) + if len(prefix) == 0 or not use_prefix: + return suffix + + if get_delimeter() == "-": + return f"{prefix}:{make_field_name(suffix.split('.'))}" + else: + assert get_delimeter() == "_" + return f"{prefix}:{suffix}" + + +@functools.lru_cache(maxsize=None) +def _get_ansi_pattern() -> re.Pattern: + # https://stackoverflow.com/a/14693789 + return re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") + + +def strip_ansi_sequences(x: str): + return _get_ansi_pattern().sub("", x) + + +def multi_metavar_from_single(single: str) -> str: + if len(strip_ansi_sequences(single)) >= 32: + # Shorten long metavars + return f"[{single} [...]]" + else: + return f"[{single} [{single} ...]]" + + +def remove_single_line_breaks(helptext: str) -> str: + lines = helptext.split("\n") + output_parts: List[str] = [] + for line in lines: + # Remove trailing whitespace. + line = line.rstrip() + + # Empty line. + if len(line) == 0: + prev_is_break = len(output_parts) >= 1 and output_parts[-1] == "\n" + if not prev_is_break: + output_parts.append("\n") + output_parts.append("\n") + + # Empty line. + else: + if not line[0].isalpha(): + output_parts.append("\n") + prev_is_break = len(output_parts) >= 1 and output_parts[-1] == "\n" + if len(output_parts) >= 1 and not prev_is_break: + output_parts.append(" ") + output_parts.append(line) + + return "".join(output_parts).rstrip() # type: ignore diff --git a/src/tyro/_subcommand_matching.py b/src/tyro/_subcommand_matching.py new file mode 100644 index 000000000..e9a45a425 --- /dev/null +++ b/src/tyro/_subcommand_matching.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import dataclasses +from typing import Any, Callable, Dict, Optional, Tuple, Union + +from typing_extensions import get_args, get_origin + +from . import _fields, _instantiators, _resolver, _typing +from .conf import _confstruct + + +def match_subcommand( + default: Any, + subcommand_config_from_name: Dict[str, _confstruct._SubcommandConfiguration], + subcommand_type_from_name: Dict[str, type], +) -> Optional[str]: + """Given a subcommand mapping and a default, return which subcommand the default + corresponds to.""" + + # It's really hard to concretize a generic type at runtime, so we just... + # don't. :-) + if hasattr(type(default), "__parameters__"): + raise _instantiators.UnsupportedTypeAnnotationError( + "Default values for generic subparsers are not supported." + ) + + # Get default subcommand name: by default hash. + default_hash = object.__hash__(default) + for subcommand_name, conf in subcommand_config_from_name.items(): + if conf.default in _fields.MISSING_SINGLETONS: + continue + if default_hash == object.__hash__(conf.default): + return subcommand_name + + # Get default subcommand name: by default value. + for subcommand_name, conf in subcommand_config_from_name.items(): + if conf.default in _fields.MISSING_SINGLETONS: + continue + equal = default == conf.default + if isinstance(equal, bool) and equal: + return subcommand_name + + # Get default subcommand name: by concrete type tree. + typetree = _TypeTree.make(type(default), default) + for subcommand_name, conf in subcommand_config_from_name.items(): + if conf.default in _fields.MISSING_SINGLETONS: + continue + if typetree == _TypeTree.make(type(conf.default), conf.default): + return subcommand_name + + # Get default subcommand name: by annotated type tree. + for subcommand_name, subcommand_type in subcommand_type_from_name.items(): + if typetree.is_subtype_of( + _TypeTree.make(subcommand_type, _fields.MISSING_NONPROP) + ): + return subcommand_name + + # Failed. This should never happen, we'll raise an error outside of this function if + # this is the case. + return None # pragma: no cover + + +@dataclasses.dataclass(frozen=True) +class _TypeTree: + typ: Any + children: Dict[str, _TypeTree] + + @staticmethod + def make( + typ: Union[Callable, _typing.TypeForm], + default_instance: _fields.DefaultInstance, + ) -> _TypeTree: + """From an object instance, return a data structure representing the types in the object.""" + try: + typ, _type_from_typevar, field_list = _fields.field_list_from_callable( + typ, default_instance=default_instance, support_single_arg_types=False + ) + except _instantiators.UnsupportedTypeAnnotationError: + return _TypeTree(typ, {}) + + return _TypeTree( + typ, + { + field.intern_name: _TypeTree.make(field.type_or_callable, field.default) + for field in field_list + }, + ) + + def is_subtype_of(self, supertype: _TypeTree) -> bool: + # Generalize to unions. + def _get_type_options(typ: _typing.TypeForm) -> Tuple[_typing.TypeForm, ...]: + return get_args(typ) if get_origin(typ) is Union else (typ,) + + self_types = _get_type_options(self.typ) + super_types = _get_type_options(supertype.typ) + + # Check against supertypes. + for self_type in self_types: + self_type = _resolver.unwrap_annotated(self_type) + self_type, _ = _resolver.unwrap_newtype_and_aliases(self_type) + ok = False + for super_type in super_types: + super_type = _resolver.unwrap_annotated(super_type) + self_type, _ = _resolver.unwrap_newtype_and_aliases(self_type) + if issubclass(self_type, super_type): + ok = True + if not ok: + return False + + return True diff --git a/src/tyro/_typing.py b/src/tyro/_typing.py new file mode 100644 index 000000000..b925e3c36 --- /dev/null +++ b/src/tyro/_typing.py @@ -0,0 +1,16 @@ +# tyro's API relies heavily on parsing type annotations, which may be both "regular" +# types like `int`, `str`, and classes, or more general forms like `int | str` or +# `Annotated[str, ...]`. +# +# To correctly annotate variables that can take these more general type forms, we need +# to wait for a typing.TypeForm PEP: +# https://github.com/python/mypy/issues/9773 +# +# This doesn't yet exist, so in the meantime we use Type[T] everywhere we would +# otherwise have TypeForm[T]. This mostly works, and fortunately is supported by pyright +# and pylance (relevant: https://github.com/microsoft/pyright/issues/4298), but should +# be switched for the correct typing.TypeForm annotation once it's available. + +from typing import Type as TypeForm + +__all__ = ["TypeForm"] diff --git a/src/tyro/_unsafe_cache.py b/src/tyro/_unsafe_cache.py new file mode 100644 index 000000000..8b3fc37c6 --- /dev/null +++ b/src/tyro/_unsafe_cache.py @@ -0,0 +1,48 @@ +import functools +from typing import Any, Callable, Dict, List, TypeVar + +CallableType = TypeVar("CallableType", bound=Callable) + + +_cache_list: List[Dict[Any, Any]] = [] + + +def clear_cache() -> None: + for c in _cache_list: + c.clear() + + +def unsafe_cache(maxsize: int) -> Callable[[CallableType], CallableType]: + """Cache decorator that relies object IDs when arguments are unhashable. Makes the + very strong assumption of not only immutability, but that unhashable types don't go + out of scope.""" + + _cache_list.append({}) + local_cache = _cache_list[-1] + + def inner(f: CallableType) -> CallableType: + @functools.wraps(f) + def wrapped_f(*args, **kwargs): + key = tuple(unsafe_hash(arg) for arg in args) + tuple( + ("__kwarg__", k, unsafe_hash(v)) for k, v in kwargs.items() + ) + + if key in local_cache: + return local_cache[key] + + out = f(*args, **kwargs) + local_cache[key] = out + if len(local_cache) > maxsize: + local_cache.pop(next(iter(local_cache))) + return out + + return wrapped_f # type: ignore + + return inner + + +def unsafe_hash(obj: Any) -> Any: + try: + return hash(obj) + except TypeError: + return id(obj) diff --git a/src/tyro/conf/__init__.py b/src/tyro/conf/__init__.py new file mode 100644 index 000000000..f58407773 --- /dev/null +++ b/src/tyro/conf/__init__.py @@ -0,0 +1,25 @@ +"""The :mod:`tyro.conf` submodule contains helpers for attaching parsing-specific +configuration metadata to types via [PEP 593](https://peps.python.org/pep-0593/) runtime +annotations. + +Configuration flags are applied recursively, and should generally be subscripted: +`Fixed[T]`, `Suppress[T]`, etc. + +Features here are supported, but generally unnecessary and should be used sparingly. +""" + +from ._confstruct import arg as arg +from ._confstruct import subcommand as subcommand +from ._markers import AvoidSubcommands as AvoidSubcommands +from ._markers import ConsolidateSubcommandArgs as ConsolidateSubcommandArgs +from ._markers import Fixed as Fixed +from ._markers import FlagConversionOff as FlagConversionOff +from ._markers import OmitArgPrefixes as OmitArgPrefixes +from ._markers import OmitSubcommandPrefixes as OmitSubcommandPrefixes +from ._markers import Positional as Positional +from ._markers import PositionalRequiredArgs as PositionalRequiredArgs +from ._markers import Suppress as Suppress +from ._markers import SuppressFixed as SuppressFixed +from ._markers import UseAppendAction as UseAppendAction +from ._markers import UseCounterAction as UseCounterAction +from ._markers import configure as configure diff --git a/src/tyro/conf/_confstruct.py b/src/tyro/conf/_confstruct.py new file mode 100644 index 000000000..d4a87b544 --- /dev/null +++ b/src/tyro/conf/_confstruct.py @@ -0,0 +1,209 @@ +import dataclasses +from typing import Any, Callable, Optional, Sequence, Tuple, Type, Union, overload + +from .._fields import MISSING_NONPROP + + +@dataclasses.dataclass(frozen=True) +class _SubcommandConfiguration: + name: Optional[str] + default: Any + description: Optional[str] + prefix_name: bool + constructor_factory: Optional[Callable[[], Union[Type, Callable]]] + + def __hash__(self) -> int: + return object.__hash__(self) + + +@overload +def subcommand( + name: Optional[str] = None, + *, + default: Any = MISSING_NONPROP, + description: Optional[str] = None, + prefix_name: bool = True, + constructor: None = None, + constructor_factory: Optional[Callable[[], Union[Type, Callable]]] = None, +) -> Any: ... + + +@overload +def subcommand( + name: Optional[str] = None, + *, + default: Any = MISSING_NONPROP, + description: Optional[str] = None, + prefix_name: bool = True, + constructor: Optional[Union[Type, Callable]] = None, + constructor_factory: None = None, +) -> Any: ... + + +def subcommand( + name: Optional[str] = None, + *, + default: Any = MISSING_NONPROP, + description: Optional[str] = None, + prefix_name: bool = True, + constructor: Optional[Union[Type, Callable]] = None, + constructor_factory: Optional[Callable[[], Union[Type, Callable]]] = None, +) -> Any: + """Returns a metadata object for configuring subcommands with `typing.Annotated`. + Useful for aesthetics. + + Consider the standard approach for creating subcommands: + + ```python + tyro.cli( + Union[NestedTypeA, NestedTypeB] + ) + ``` + + This will create two subcommands: `nested-type-a` and `nested-type-b`. + + Annotating each type with `tyro.conf.subcommand()` allows us to override for + each subcommand the (a) name, (b) defaults, (c) helptext, and (d) whether to prefix + the name or not. + + ```python + tyro.cli( + Union[ + Annotated[ + NestedTypeA, subcommand("a", ...) + ], + Annotated[ + NestedTypeB, subcommand("b", ...) + ], + ] + ) + ``` + + Arguments: + name: The name of the subcommand in the CLI. + default: A default value for the subcommand, for struct-like types. (eg + dataclasses) + description: Description of this option to use in the helptext. Defaults to + docstring. + prefix_name: Whether to prefix the name of the subcommand based on where it + is in a nested structure. + constructor: A constructor type or function. This should either be (a) a subtype + of an argument's annotated type, or (b) a function with type-annotated + inputs that returns an instance of the annotated type. This will be used in + place of the argument's type for parsing arguments. No validation is done. + constructor_factory: A function that returns a constructor type or function. + Useful when the constructor isn't immediately available. + """ + assert not ( + constructor is not None and constructor_factory is not None + ), "`constructor` and `constructor_factory` cannot both be set." + return _SubcommandConfiguration( + name, + default, + description, + prefix_name, + constructor_factory=constructor_factory + if constructor is None + else lambda: constructor, + ) + + +@dataclasses.dataclass(frozen=True) +class _ArgConfiguration: + # These are all optional by default in order to support multiple tyro.conf.arg() + # annotations. A None value means "don't overwrite the current value". + name: Optional[str] + metavar: Optional[str] + help: Optional[str] + aliases: Optional[Tuple[str, ...]] + prefix_name: Optional[bool] + constructor_factory: Optional[Callable[[], Union[Type, Callable]]] + + +@overload +def arg( + *, + name: Optional[str] = None, + metavar: Optional[str] = None, + help: Optional[str] = None, + aliases: Optional[Sequence[str]] = None, + prefix_name: Optional[bool] = None, + constructor: None = None, + constructor_factory: Optional[Callable[[], Union[Type, Callable]]] = None, +) -> Any: ... + + +@overload +def arg( + *, + name: Optional[str] = None, + metavar: Optional[str] = None, + help: Optional[str] = None, + aliases: Optional[Sequence[str]] = None, + prefix_name: Optional[bool] = None, + constructor: Optional[Union[Type, Callable]] = None, + constructor_factory: None = None, +) -> Any: ... + + +def arg( + *, + name: Optional[str] = None, + metavar: Optional[str] = None, + help: Optional[str] = None, + aliases: Optional[Sequence[str]] = None, + prefix_name: Optional[bool] = None, + constructor: Optional[Union[Type, Callable]] = None, + constructor_factory: Optional[Callable[[], Union[Type, Callable]]] = None, +) -> Any: + """Returns a metadata object for fine-grained argument configuration with + `typing.Annotated`. Should typically not be required. + + We support using `arg()` at the root of arguments. For example: + ```python + x: Annotated[int, tyro.conf.arg(...)] + ``` + + Nesting `arg()` within other types is generally not supported: + ```python + # Not supported. + x: tuple[Annotated[int, tyro.conf.arg(...)], ...] + ``` + + Arguments: + name: A new name for the argument in the CLI. + metavar: Argument name in usage messages. The type is used by default. + help: Helptext for this argument. The docstring is used by default. + aliases: Aliases for this argument. All strings in the sequence should start + with a hyphen (-). Aliases will _not_ currently be prefixed in a nested + structure, and are not supported for positional arguments. + prefix_name: Whether or not to prefix the name of the argument based on where + it is in a nested structure. Arguments are prefixed by default. + constructor: A constructor type or function. This should either be (a) a subtype + of an argument's annotated type, or (b) a function with type-annotated + inputs that returns an instance of the annotated type. This will be used in + place of the argument's type for parsing arguments. No validation is done. + constructor_factory: A function that returns a constructor type or function. + Useful when the constructor isn't immediately available. + + Returns: + Object to attach via `typing.Annotated[]`. + """ + assert not ( + constructor is not None and constructor_factory is not None + ), "`constructor` and `constructor_factory` cannot both be set." + + if aliases is not None: + for alias in aliases: + assert alias.startswith("-"), "Argument alias needs to start with a hyphen!" + + return _ArgConfiguration( + name=name, + metavar=metavar, + help=help, + aliases=tuple(aliases) if aliases is not None else None, + prefix_name=prefix_name, + constructor_factory=constructor_factory + if constructor is None + else lambda: constructor, + ) diff --git a/src/tyro/conf/_markers.py b/src/tyro/conf/_markers.py new file mode 100644 index 000000000..833b5fc3c --- /dev/null +++ b/src/tyro/conf/_markers.py @@ -0,0 +1,195 @@ +from typing import TYPE_CHECKING, Any, Callable, TypeVar + +from typing_extensions import Annotated + +from .. import _singleton + +# Current design issue: markers are applied recursively to nested structures, but can't +# be unapplied. + +# Note that all Annotated[T, None] values are just for static checkers. The real marker +# singletons are instantiated dynamically below. +# +# An alias could ideally be made, but SpecialForm aliases aren't well supported by +# type checkers. + +T = TypeVar("T") + +Positional = Annotated[T, None] +"""A type `T` can be annotated as `Positional[T]` if we want to parse it as a positional +argument.""" + +PositionalRequiredArgs = Annotated[T, None] +"""Make all arguments without defaults positional.""" + +# Private marker. For when an argument is not only positional in the CLI, but also in +# the callable. +_PositionalCall = Annotated[T, None] + +# Private markers for when arguments should be passed in via *args or **kwargs. +_UnpackArgsCall = Annotated[T, None] +_UnpackKwargsCall = Annotated[T, None] + +# Private marker. +_OPTIONAL_GROUP = Annotated[T, None] + +# TODO: the verb tenses here are inconsistent, naming could be revisited. +# Perhaps Suppress should be Suppressed? But SuppressedFixed would be weird. + +Fixed = Annotated[T, None] +"""A type `T` can be annotated as `Fixed[T]` to prevent `tyro.cli` from parsing it; a +default value should be set instead. Fields that can't be parsed with defaults will also +be marked as fixed automatically.""" + +Suppress = Annotated[T, None] +"""A type `T` can be annotated as `Suppress[T]` to prevent `tyro.cli` from parsing it, and +to prevent it from showing up in helptext.""" + +SuppressFixed = Annotated[T, None] +"""Hide fields that are either manually or automatically marked as fixed.""" + +FlagConversionOff = Annotated[T, None] +"""Turn off flag conversion for booleans with default values. Instead, types annotated +with `bool` will expect an explicit True or False. + +Can be used directly on boolean annotations, `FlagConversionOff[bool]`, or recursively +applied to nested types.""" + +AvoidSubcommands = Annotated[T, None] +"""Avoid creating subcommands when a default is provided for unions over nested types. +This simplifies CLI interfaces, but makes them less expressive. + +Can be used directly on union types, `AvoidSubcommands[Union[...]]`, or recursively +applied to nested types.""" + +ConsolidateSubcommandArgs = Annotated[T, None] +"""Consolidate arguments applied to subcommands. Makes CLI less sensitive to argument +ordering, at the cost of support for optional subcommands. + +By default, `tyro` will generate a traditional CLI interface where args are applied to +the directly preceding subcommand. When we have two subcommands `s1` and `s2`: +``` +python x.py {--root options} s1 {--s1 options} s2 {--s2 options} +``` + +This can be frustrating because the resulting CLI is sensitive to the positioning of +options. + +To consolidate subcommands, we push arguments to the end, after all subcommands: +``` +python x.py s1 s2 {--root, s1, and s2 options} +``` + +This is more robust to reordering of options, ensuring that any new options can simply +be placed at the end of the command. +""" + +OmitSubcommandPrefixes = Annotated[T, None] +"""Make CLI inputs used for subcommands shorter by omitting the subcommand-specific +portion of the prefix. + +If we have a structure with the field: + + cmd: Union[NestedTypeA, NestedTypeB] + +By default, `--cmd.arg` may be generated as a flag for each dataclass in the union. +If subcommand prefixes are omitted, we would instead have `--arg`. + +By default, `cmd:nested-type-a` and `cmd:nested-type-b` may be generated as subcommand. +If subcommand prefixes are omitted, we would instead have `nested-type-a` and +`nested-type-b`. +""" + +OmitArgPrefixes = Annotated[T, None] +"""Make flags used for keyword arguments shorter by omitting prefixes. + +If we have a structure with the field: + + cmd: NestedType + +By default, `--cmd.arg` may be generated as a flag. If prefixes are omitted, we would +instead simply have `--arg`. +""" + +UseAppendAction = Annotated[T, None] +"""Use "append" actions for variable-length arguments. + +Given an annotation like `x: list[int]`, this means that `x = [0, 1, 2]` can be set via +the CLI syntax `--x 0 --x 1 --x 2` instead of the default of `--x 0 1 2`. + +The resulting syntax may be more user-friendly; for `tyro`, it also enables support for +otherwise ambiguous annotations like `list[list[int]]`. + +Can be applied to all variable-length sequences (`list[T]`, `Sequence[T]`, +`tuple[T, ...]`, etc), including dictionaries without default values. +""" + +UseCounterAction = Annotated[T, None] +"""Use "counter" actions for integer arguments. Example usage: `verbose: UseCounterAction[int]`.""" + + +CallableType = TypeVar("CallableType", bound=Callable) + +# Dynamically generate marker singletons. +# These can be used one of two ways: +# - Marker[T] +# - Annotated[T, Marker] + + +class _Marker(_singleton.Singleton): + def __getitem__(self, key): + return Annotated.__class_getitem__((key, self)) # type: ignore + + +Marker = Any + + +def configure(*markers: Marker) -> Callable[[CallableType], CallableType]: + """Decorator for applying configuration options. + + Consider using the `config=` argument of `tyro.cli()` instead, which takes the same + config marker objects as inputs. + + Configuration markers are implemented via `typing.Annotated` and straightforward + to apply to types, for example: + + ```python + field: tyro.conf.FlagConversionOff[bool] + ``` + + This decorator makes markers applicable to general functions as well: + + ```python + # Recursively apply FlagConversionOff to all fields in `main()`. + @tyro.conf.configure(tyro.conf.FlagConversionOff) + def main(field: bool) -> None: + ... + ``` + + Args: + markers: Options to apply. + """ + + def _inner(callable: CallableType) -> CallableType: + # We'll read from __tyro_markers__ in `_resolver.unwrap_annotated()`. + callable.__tyro_markers__ = markers # type: ignore + return callable + + return _inner + + +if not TYPE_CHECKING: + + def _make_marker(description: str) -> _Marker: + class _InnerMarker(_Marker): + def __repr__(self): + return description + + return _InnerMarker() + + _dynamic_marker_types = {} + for k, v in dict(globals()).items(): + if v == Annotated[T, None]: + _dynamic_marker_types[k] = _make_marker(k) + globals().update(_dynamic_marker_types) + del _dynamic_marker_types diff --git a/src/tyro/extras/__init__.py b/src/tyro/extras/__init__.py new file mode 100644 index 000000000..0afb6b854 --- /dev/null +++ b/src/tyro/extras/__init__.py @@ -0,0 +1,16 @@ +"""The :mod:`tyro.extras` submodule contains helpers that complement :func:`tyro.cli()`. + +Compared to the core interface, APIs here are more likely to be changed or deprecated.""" + +from .._argparse_formatter import set_accent_color as set_accent_color +from .._cli import get_parser as get_parser +from ._base_configs import overridable_config_cli as overridable_config_cli +from ._base_configs import ( + subcommand_type_from_defaults as subcommand_type_from_defaults, +) +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_cli_from_dict import ( + subcommand_cli_from_dict as subcommand_cli_from_dict, +) diff --git a/src/tyro/extras/_base_configs.py b/src/tyro/extras/_base_configs.py new file mode 100644 index 000000000..33157e076 --- /dev/null +++ b/src/tyro/extras/_base_configs.py @@ -0,0 +1,148 @@ +from typing import Mapping, Optional, Sequence, Tuple, TypeVar, Union + +from typing_extensions import Annotated + +from .._typing import TypeForm + +T = TypeVar("T") + + +def overridable_config_cli( + configs: Mapping[str, Tuple[str, T]], + *, + args: Optional[Sequence[str]] = None, +) -> T: + """Helper function for creating a CLI interface that allows us to choose + between default config objects (typically dataclasses) and override values + within it. Turns off subcommand creation for any union types within the + config object. + + This is a lightweight wrapper over :func:`tyro.cli()`, with some default + arguments populated. Also see + :func:`tyro.extras.subcommand_type_from_defaults()`. + + + Example usage: + ```python + import dataclasses + + import tyro + + + @dataclasses.dataclass + class Config: + a: int + b: str + + + default_configs = { + "small": ( + "Small config", + Config(1, "small"), + ), + "big": ( + "Big config", + Config(100, "big"), + ), + } + config = tyro.extras.overridable_config_cli(default_configs) + print(config) + ``` + + Args: + configs: A dictionary of config names mapped to a tuple of + (description, config object). + args: Optional arguments to pass to the CLI. + """ + import tyro + + return tyro.cli( + tyro.extras.subcommand_type_from_defaults( + defaults={k: v[1] for k, v in configs.items()}, + descriptions={k: v[0] for k, v in configs.items()}, + ), + # Don't create subcommands for union types within the config object. + config=(tyro.conf.AvoidSubcommands,), + args=args, + ) + + +def subcommand_type_from_defaults( + defaults: Mapping[str, T], + descriptions: Mapping[str, str] = {}, + *, + prefix_names: bool = True, +) -> TypeForm[T]: + """Construct a Union type for defining subcommands that choose between defaults. + + For example, when `defaults` is set to: + + ```python + { + "small": Config(...), + "big": Config(...), + } + ``` + + We return: + + ```python + Union[ + Annotated[ + Config, + tyro.conf.subcommand("small", default=Config(...)) + ], + Annotated[ + Config, + tyro.conf.subcommand("big", default=Config(...)) + ] + ] + ``` + + Direct use of `typing.Union` and :func:`tyro.conf.subcommand()` should generally be + preferred, but this function can be helpful for succinctness. + + .. warning:: + The type returned by this function can be safely used as an input to + :func:`tyro.cli()`, but for static analysis when used for annotations we + recommend applying a `TYPE_CHECKING` guard: + + .. code-block:: python + + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + # Static type seen by language servers, type checkers, etc. + SelectableConfig = Config + else: + # Runtime type used by tyro. + SelectableConfig = subcommand_type_from_defaults(...) + + Args: + defaults: A dictionary of default subcommand instances. + descriptions: A dictionary conttaining descriptions for helptext. + prefix_names: Whether to prefix subcommand names. + + Returns: + A subcommand type, which can be passed to :func:`tyro.cli`. + """ + import tyro + + # We need to form a union type, which requires at least two elements. + assert len(defaults) >= 2, "At least two subcommands are required." + return Union.__getitem__( # type: ignore + tuple( + Annotated.__class_getitem__( # type: ignore + ( + type(v), + tyro.conf.subcommand( + k, + default=v, + description=descriptions.get(k, ""), + prefix_name=prefix_names, + ), + ) + ) + for k, v in defaults.items() + ) + ) diff --git a/src/tyro/extras/_choices_type.py b/src/tyro/extras/_choices_type.py new file mode 100644 index 000000000..57998afc2 --- /dev/null +++ b/src/tyro/extras/_choices_type.py @@ -0,0 +1,39 @@ +import enum +from typing import Iterable, TypeVar, Union + +from typing_extensions import Literal + +from .._typing import TypeForm + +T = TypeVar("T", bound=Union[int, str, bool, enum.Enum]) + + +def literal_type_from_choices(choices: Iterable[T]) -> TypeForm[T]: + """Generate a `typing.Literal[]` type that constrains values to a set of choices. + + Using `Literal[...]` directly should generally be preferred, but this function can be + helpful when choices are generated dynamically. + + .. warning:: + The type returned by this function can be safely used as an input to + :func:`tyro.cli()`, but for static analysis when used for annotations we + recommend applying a `TYPE_CHECKING` guard: + + .. code-block:: python + + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + # Static type seen by language servers, type checkers, etc. + Color = str + else: + # Runtime type used by tyro. + Color = literal_type_from_choices(["red", "green", "blue"]) + + Args: + choices: Options to choose from. + + Returns: + A type that can be passed to :func:`tyro.cli()`. + """ + return Literal.__getitem__(tuple(choices)) # type: ignore diff --git a/src/tyro/extras/_serialization.py b/src/tyro/extras/_serialization.py new file mode 100644 index 000000000..bfb5a28d2 --- /dev/null +++ b/src/tyro/extras/_serialization.py @@ -0,0 +1,232 @@ +"""Type-safe, human-readable serialization helpers for dataclasses.""" + +from __future__ import annotations + +import dataclasses +import enum +import functools +from typing import IO, TYPE_CHECKING, Any, Optional, Set, Type, TypeVar, Union + +if TYPE_CHECKING: + # Since serialization functionality is deprecated, the yaml dependency is optional. + import yaml + +from typing_extensions import get_args, get_origin + +from .. import _fields, _resolver + +ENUM_YAML_TAG_PREFIX = "!enum:" +DATACLASS_YAML_TAG_PREFIX = "!dataclass:" +MISSING_YAML_TAG_PREFIX = "!missing" + +DataclassType = TypeVar("DataclassType") + + +def _get_contained_special_types_from_type( + cls: Type[Any], + _parent_contained_dataclasses: Optional[Set[Type[Any]]] = None, +) -> Set[Type[Any]]: + """Takes a dataclass type, and recursively searches its fields for dataclass or enum + types.""" + assert _resolver.is_dataclass(cls) + parent_contained_dataclasses = ( + set() + if _parent_contained_dataclasses is None + else _parent_contained_dataclasses + ) + + cls = _resolver.unwrap_annotated(cls) + cls, type_from_typevar = _resolver.resolve_generic_types(cls) + + contained_special_types = {cls} + + def handle_type(typ: Type[Any]) -> Set[Type[Any]]: + # Handle dataclasses. + if _resolver.is_dataclass(typ) and typ not in parent_contained_dataclasses: + return _get_contained_special_types_from_type( + typ, + _parent_contained_dataclasses=contained_special_types + | parent_contained_dataclasses, + ) + + # Handle enums. + elif isinstance(typ, enum.EnumMeta): + return {typ} + + # Handle Union, Annotated, List, etc. No-op when there are no args. + return functools.reduce(set.union, map(handle_type, get_args(typ)), set()) + + # Handle generics. + for typ in type_from_typevar.values(): + contained_special_types |= handle_type(typ) + + if cls in parent_contained_dataclasses: + return contained_special_types + + # Handle fields. + for field in _resolver.resolved_fields(cls): # type: ignore + assert not isinstance(field.type, str) + contained_special_types |= handle_type(field.type) + + # Handle subclasses. + for subclass in cls.__subclasses__(): + contained_special_types |= handle_type(subclass) + + return contained_special_types + + +def _make_loader(cls: Type[Any]) -> Type[yaml.Loader]: + import yaml + + class DataclassLoader(yaml.Loader): + pass + + # Design Q: do we want to support multiple dataclass types with the same name? + # - Why yes: rudimentary support for this is easy. + # - Why no: might cause confusing error messages? Using all possible context cues + # for this is also hard. + # + # => let's just keep things simple, assert uniqueness for now. Easier to add new + # features later than remove them. + + contained_types = list(_get_contained_special_types_from_type(cls)) + contained_type_names = list(map(lambda cls: cls.__name__, contained_types)) + assert len(set(contained_type_names)) == len(contained_type_names), ( + "Contained dataclass type names must all be unique, but got" + f" {contained_type_names}" + ) + + def make_dataclass_constructor(typ: Type[Any]): + return lambda loader, node: typ(**loader.construct_mapping(node)) + + def make_enum_constructor(typ: Type[enum.Enum]): + return lambda loader, node: typ[loader.construct_python_str(node)] + + for typ, name in zip(contained_types, contained_type_names): + if dataclasses.is_dataclass(typ): + DataclassLoader.add_constructor( + tag=DATACLASS_YAML_TAG_PREFIX + name, + constructor=make_dataclass_constructor(typ), + ) + elif issubclass(typ, enum.Enum): + DataclassLoader.add_constructor( + tag=ENUM_YAML_TAG_PREFIX + name, + constructor=make_enum_constructor(typ), + ) + else: + assert False + + DataclassLoader.add_constructor( + tag=MISSING_YAML_TAG_PREFIX, + constructor=lambda *_unused: _fields.MISSING_PROP, # type: ignore + ) + + return DataclassLoader + + +def _make_dumper(instance: Any) -> Type[yaml.Dumper]: + import yaml + + class DataclassDumper(yaml.Dumper): + def ignore_aliases(self, data): + return super().ignore_aliases(data) or data is _fields.MISSING_PROP + + contained_types = list(_get_contained_special_types_from_type(type(instance))) + contained_type_names = list(map(lambda cls: cls.__name__, contained_types)) + + # Note: this is currently a stricter than necessary assert. + assert len(set(contained_type_names)) == len(contained_type_names), ( + "Contained dataclass/enum names must all be unique, but got" + f" {contained_type_names}" + ) + + def make_representer(name: str): + def representer(dumper: DataclassDumper, data: Any) -> yaml.Node: + if dataclasses.is_dataclass(data): + return dumper.represent_mapping( + tag=DATACLASS_YAML_TAG_PREFIX + name, + mapping={ + field.name: getattr(data, field.name) + for field in dataclasses.fields(data) + if field.init + }, + ) + elif isinstance(data, enum.Enum): + return dumper.represent_scalar( + tag=ENUM_YAML_TAG_PREFIX + name, value=data.name + ) + assert False + + return representer + + for typ, name in zip(contained_types, contained_type_names): + DataclassDumper.add_representer(typ, make_representer(name)) + + DataclassDumper.add_representer( + type(_fields.MISSING_PROP), + lambda dumper, data: dumper.represent_scalar( + tag=MISSING_YAML_TAG_PREFIX, value="" + ), + ) + return DataclassDumper + + +def from_yaml( + cls: Type[DataclassType], + stream: Union[str, IO[str], bytes, IO[bytes]], +) -> DataclassType: + """Re-construct a dataclass instance from a yaml-compatible string, which should be + generated from `tyro.extras.to_yaml()`. + + As a secondary feature aimed at enabling the use of :func:`tyro.cli` for general + configuration use cases, we also introduce functions for human-readable dataclass + serialization: :func:`tyro.extras.from_yaml` and :func:`tyro.extras.to_yaml` attempt + to strike a balance between flexibility and robustness — in contrast to naively + dumping or loading dataclass instances (via pickle, PyYAML, etc), explicit type + references enable custom tags that are robust against code reorganization and + refactor, while a PyYAML backend enables serialization of arbitrary Python objects. + + .. warning:: + Serialization functionality is stable but deprecated. It may be removed in a + future version of :code:`tyro`. + + Args: + cls: Type to reconstruct. + stream: YAML to read from. + + Returns: + Instantiated dataclass. + """ + import yaml + + out = yaml.load(stream, Loader=_make_loader(cls)) + origin_cls = get_origin(cls) + assert isinstance(out, origin_cls if origin_cls is not None else cls) + return out + + +def to_yaml(instance: Any) -> str: + """Serialize a dataclass; returns a yaml-compatible string that can be deserialized + via `tyro.extras.from_yaml()`. + + As a secondary feature aimed at enabling the use of :func:`tyro.cli` for general + configuration use cases, we also introduce functions for human-readable dataclass + serialization: :func:`tyro.extras.from_yaml` and :func:`tyro.extras.to_yaml` attempt + to strike a balance between flexibility and robustness — in contrast to naively + dumping or loading dataclass instances (via pickle, PyYAML, etc), explicit type + references enable custom tags that are robust against code reorganization and + refactor, while a PyYAML backend enables serialization of arbitrary Python objects. + + .. warning:: + Serialization functionality is stable but deprecated. It may be removed in a + future version of :code:`tyro`. + + Args: + instance: Dataclass instance to serialize. + + Returns: + YAML string. + """ + import yaml + + return "# tyro YAML.\n" + yaml.dump(instance, Dumper=_make_dumper(instance)) diff --git a/src/tyro/extras/_subcommand_cli_from_dict.py b/src/tyro/extras/_subcommand_cli_from_dict.py new file mode 100644 index 000000000..6710be40f --- /dev/null +++ b/src/tyro/extras/_subcommand_cli_from_dict.py @@ -0,0 +1,111 @@ +from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, Union, overload + +from typing_extensions import Annotated + +from .._cli import cli +from ..conf import subcommand + +T = TypeVar("T") + + +@overload +def subcommand_cli_from_dict( + subcommands: Dict[str, Callable[..., T]], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + use_underscores: bool = False, +) -> T: ... + + +# TODO: hack. We prefer the above signature, which Pyright understands, but as of 1.6.1 +# mypy doesn't reason about the generics properly. +@overload +def subcommand_cli_from_dict( + subcommands: Dict[str, Callable[..., Any]], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + use_underscores: bool = False, +) -> Any: ... + + +def subcommand_cli_from_dict( + subcommands: Dict[str, Callable[..., Any]], + *, + prog: Optional[str] = None, + description: Optional[str] = None, + args: Optional[Sequence[str]] = None, + use_underscores: bool = False, +) -> Any: + """Generate a subcommand CLI from a dictionary of functions. + + For an input like: + + ```python + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + } + ) + ``` + + This is internally accomplished by generating and calling: + + ```python + from typing import Annotated, Any, Union + import tyro + + tyro.cli( + Union[ + Annotated[ + Any, + tyro.conf.subcommand(name="checkout", constructor=checkout), + ], + Annotated[ + Any, + tyro.conf.subcommand(name="commit", constructor=commit), + ], + ] + ) + ``` + + Args: + subcommands: Dictionary that maps the subcommand name to function to call. + 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, `f`'s 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 delimeter 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 + """ + # We need to form a union type, which requires at least two elements. + assert len(subcommands) >= 2, "At least two subcommands are required." + return cli( + Union.__getitem__( # type: ignore + tuple( + [ + Annotated[ + # The constructor function can return any object. + Any, + # We'll instantiate this object by invoking a subcommand with + # name k, via a constructor. + subcommand(name=k, constructor=v), + ] + for k, v in subcommands.items() + ] + ) + ), + prog=prog, + description=description, + args=args, + use_underscores=use_underscores, + ) diff --git a/dcargs/py.typed b/src/tyro/py.typed similarity index 100% rename from dcargs/py.typed rename to src/tyro/py.typed diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..23d01aae4 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +import sys +from typing import List + +collect_ignore_glob: List[str] = [] + +if not sys.version_info >= (3, 8): + collect_ignore_glob.append("*_min_py38.py") + collect_ignore_glob.append("*_min_py38_generated.py") + +if not sys.version_info >= (3, 9): + collect_ignore_glob.append("*_min_py39.py") + collect_ignore_glob.append("*_min_py39_generated.py") + +if not sys.version_info >= (3, 10): + collect_ignore_glob.append("*_min_py310.py") + collect_ignore_glob.append("*_min_py310_generated.py") + +if not sys.version_info >= (3, 12): + collect_ignore_glob.append("*_min_py312.py") + collect_ignore_glob.append("*_min_py312_generated.py") + +if not sys.version_info >= (3, 11): + collect_ignore_glob.append("test_py311_generated/*.py") diff --git a/tests/helptext_utils.py b/tests/helptext_utils.py new file mode 100644 index 000000000..041f5c58b --- /dev/null +++ b/tests/helptext_utils.py @@ -0,0 +1,84 @@ +import argparse +import contextlib +import io +import os +from typing import Any, Callable, List + +import pytest + +import tyro +import tyro._arguments +import tyro._strings + + +def get_helptext_with_checks( + f: Callable, + args: List[str] = ["--help"], + use_underscores: bool = False, + default: Any = None, +) -> str: + """Get the helptext for a given tyro with input, while running various + checks along the way.""" + # Should be empty. + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli( + f, + args=args, + use_underscores=use_underscores, + default=default, + console_outputs=False, + ) + assert target.getvalue() == "" + + # Check tyro.extras.get_parser(). + parser = tyro.extras.get_parser(f, use_underscores=use_underscores) + assert isinstance(parser, argparse.ArgumentParser) + + # Returned parser should have formatting information stripped. External tools rarely + # support ANSI sequences. + unformatted_helptext = parser.format_help() + assert ( + tyro._strings.strip_ansi_sequences(unformatted_helptext) == unformatted_helptext + ) + unformatted_usage = parser.format_usage() + assert tyro._strings.strip_ansi_sequences(unformatted_usage) == unformatted_usage + + # Completion scripts; just smoke test for now. + with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")): + # `--tyro-print-completion` is deprecated! We should use `--tyro-write-completion` instead. + tyro.cli(f, default=default, args=["--tyro-print-completion", "bash"]) + with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")): + # `--tyro-print-completion` is deprecated! We should use `--tyro-write-completion` instead. + tyro.cli(f, default=default, args=["--tyro-print-completion", "zsh"]) + with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")): + tyro.cli( + f, default=default, args=["--tyro-write-completion", "bash", os.devnull] + ) + with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")): + tyro.cli( + f, default=default, args=["--tyro-write-completion", "zsh", os.devnull] + ) + + # Get the actual helptext. + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(f, args=args, use_underscores=use_underscores, default=default) + + # Check helptext with vs without formatting. This can help catch text wrapping bugs + # caused by ANSI sequences. + target2 = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target2): + tyro._arguments.USE_RICH = False + tyro.cli(f, default=default, args=args, use_underscores=use_underscores) + tyro._arguments.USE_RICH = True + + if target2.getvalue() != tyro._strings.strip_ansi_sequences(target.getvalue()): + raise AssertionError( + "Potential wrapping bug! These two strings should match:\n" + + target2.getvalue() + + "\n\n" + + tyro._strings.strip_ansi_sequences(target.getvalue()) + ) + + return target2.getvalue() diff --git a/tests/test_attrs.py b/tests/test_attrs.py new file mode 100644 index 000000000..6f44068d3 --- /dev/null +++ b/tests/test_attrs.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import contextlib +import io +import pathlib +from typing import cast + +import attr +import pytest +from attrs import define, field + +import tyro +import tyro._strings + + +def test_attrs_basic() -> None: + @attr.s + class ManyTypesA: + i: tyro.conf.Positional[int] = attr.ib() + s: str = attr.ib() + f: float = attr.ib() + p: pathlib.Path = attr.ib() + ignored: int = attr.ib(default=3, init=False) + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesA, + args=[ + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypesA(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + +def test_attrs_defaults() -> None: + @attr.s + class ManyTypesB: + i: int = attr.ib() + s: str = attr.ib() + f: float = attr.ib(default=1.0) + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesB, + args=[ + "--i", + "5", + "--s", + "5", + ], + ) == ManyTypesB(i=5, s="5", f=1.0) + + +def test_attrs_helptext() -> None: + @attr.s + class Helptext: + """This docstring should be printed as a description.""" + + x: int = attr.ib() # Documentation 1 + + # Documentation 2 + y: int = attr.ib() + + z: int = attr.ib(default=3) + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + assert tyro._strings.strip_ansi_sequences(cast(str, Helptext.__doc__)) in helptext + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_attrs_next_gen_and_factory() -> None: + @define + class Helptext: + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int = field(factory=lambda: 3) + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + assert tyro._strings.strip_ansi_sequences(cast(str, Helptext.__doc__)) in helptext + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_attrs_default_instance() -> None: + @attr.s + class ManyTypesB: + i: int = attr.ib() + s: str = attr.ib() + f: float = attr.ib(default=1.0) + k: float = attr.ib(default=1.0) + + assert tyro.cli( + ManyTypesB, + args=[ + "--i", + "5", + "--s", + "5", + ], + default=ManyTypesB(i=5, s="5", f=2.0), + ) == ManyTypesB(i=5, s="5", f=2.0) + assert tyro.cli( + ManyTypesB, + args=["--i", "5"], + default=ManyTypesB(i=5, s="5", f=2.0), + ) == ManyTypesB(i=5, s="5", f=2.0) diff --git a/tests/test_base_configs_nested.py b/tests/test_base_configs_nested.py new file mode 100644 index 000000000..f9dca8f03 --- /dev/null +++ b/tests/test_base_configs_nested.py @@ -0,0 +1,232 @@ +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Tuple + +from torch import nn +from typing_extensions import Literal + +import tyro + + +@dataclass(frozen=True) +class AdamOptimizer: + learning_rate: float = 1e-3 + betas: Tuple[float, float] = (0.9, 0.999) + + +@dataclass(frozen=True) +class SGDOptimizer: + learning_rate: float = 1e-3 + momentum: float = 0.9 + + +@dataclass(frozen=True) +class DataConfig: + test: int = 0 + + +@dataclass(frozen=True) +class ExperimentConfig: + # Dataset to run experiment on. + dataset: Literal["mnist", "imagenet-50"] + + # Optimizer parameters. + optimizer: AdamOptimizer + + # Model size. + num_layers: int + units: int + + # Batch size. + batch_size: int + + # Total number of training steps. + train_steps: int + + # Random seed. This is helpful for making sure that our experiments are all + # reproducible! + seed: int + + # Activation to use. Not specifiable via the commandline. + activation: Callable[[], nn.Module] + + +DataConfigDataParserUnion = tyro.extras.subcommand_type_from_defaults( + { + "small-data": DataConfig( + test=2221, + ), + "big-data": DataConfig( + test=2, + ), + }, + prefix_names=False, # Omit prefixes in subcommands themselves. +) + + +# Note that we could also define this library using separate YAML files (similar to +# `config_path`/`config_name` in Hydra), but staying in Python enables seamless type +# checking + IDE support. +ExperimentConfigDataParserUnion = tyro.extras.subcommand_type_from_defaults( + { + "small": ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + "big": ExperimentConfig( + dataset="imagenet-50", + optimizer=AdamOptimizer(), + batch_size=32, + num_layers=8, + units=256, + train_steps=100_000, + seed=0, + activation=nn.GELU, + ), + }, + prefix_names=False, # Omit prefixes in subcommands themselves. +) + + +if TYPE_CHECKING: + AnnotatedDataParserUnion = DataConfig + AnnotatedExperimentParserUnion = ExperimentConfig +else: + AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ + DataConfigDataParserUnion + ] # Omit prefixes of flags in subcommands. + AnnotatedExperimentParserUnion = tyro.conf.OmitSubcommandPrefixes[ + ExperimentConfigDataParserUnion + ] # Omit prefixes of flags in subcommands. + + +@dataclass +class BaseConfig: + # The name of the experiment to run. + experiment: str + + # The path to the output directory. + output_dir: str + + # The experiment configuration. + experiment_config: AnnotatedExperimentParserUnion + + # The experiment configuration. + data_config: AnnotatedDataParserUnion = DataConfig() + + +def test_base_configs_nested() -> None: + def main(cfg: BaseConfig) -> BaseConfig: + return cfg + + assert tyro.cli( + main, + args="--cfg.experiment test --cfg.output-dir test small".split(" "), + ) == BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(0), + ) + assert tyro.cli( + main, + args="--cfg.experiment test --cfg.output-dir test small small-data".split(" "), + ) == BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(2221), + ) + assert tyro.cli( + main, + args="--cfg.experiment test --cfg.output-dir test small big-data".split(" "), + ) == BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(2), + ) + + +def test_pernicious_override(): + """From: https://github.com/nerfstudio-project/nerfstudio/issues/2789 + + Situation where we: + - have a default value in the config class + - override that default value with a subcommand annotation + - override it again with a default instance + """ + assert ( + tyro.cli( + BaseConfig, + default=BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(0), + ), + args="small small-data".split(" "), + ).data_config.test + == 0 + ) + + +def test_overridable_config_helper(): + assert tyro.extras.overridable_config_cli( + { + "small-data": ( + "Small data", + DataConfig( + test=2221, + ), + ), + "big-data": ( + "Big data", + DataConfig( + test=2, + ), + ), + }, + args=["small-data", "--test", "100"], + ) == DataConfig(100) diff --git a/tests/test_boolean_optional.py b/tests/test_boolean_optional.py new file mode 100644 index 000000000..ba48bc9f3 --- /dev/null +++ b/tests/test_boolean_optional.py @@ -0,0 +1,71 @@ +import dataclasses + +import tyro + + +def test_flag_default_false() -> None: + """Test for argparse.BooleanOptionalAction-style usage.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(False), + ) == A(True) + + assert tyro.cli( + A, + args=["--no-x"], + default=A(False), + ) == A(False) + + assert tyro.cli( + A, + args=[], + default=A(False), + ) == A(False) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + + +def test_flag_default_true() -> None: + """Test for argparse.BooleanOptionalAction-style usage.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(True), + ) == A(True) + + assert tyro.cli( + A, + args=["--no-x"], + default=A(True), + ) == A(False) + + assert tyro.cli( + A, + args=[], + default=A(True), + ) == A(True) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) diff --git a/tests/test_collections.py b/tests/test_collections.py new file mode 100644 index 000000000..2f84ffefb --- /dev/null +++ b/tests/test_collections.py @@ -0,0 +1,631 @@ +import collections +import collections.abc +import contextlib +import dataclasses +import enum +import io +import sys +from typing import ( + Any, + Deque, + Dict, + FrozenSet, + List, + Optional, + Sequence, + Set, + Tuple, + Type, + Union, +) + +import pytest +from typing_extensions import Literal + +import tyro + + +def test_tuples_fixed() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, int, int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_fixed_mixed() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, str, int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, "2", 3)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_with_default() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, int, int] = (0, 1, 2) + + assert tyro.cli(A, args=[]) == A(x=(0, 1, 2)) + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + + +def test_tuple_with_literal_and_default() -> None: + @dataclasses.dataclass + class A: + x: Tuple[Literal[1, 2, 3], ...] = (1, 2) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=(1, 2)) + assert tyro.cli(A, args=["--x"]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3", "4"]) + + +def test_positional_tuple_with_literal_and_default() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.Positional[Tuple[Literal[1, 2, 3], ...]] = (1, 2) + + assert tyro.cli(A, args=["1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=(1, 2)) + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(A, args=["1", "2", "3", "4"]) + assert "invalid choice" in target.getvalue() + + +def test_tuples_fixed_multitype() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, str, float] + + assert tyro.cli(A, args=["--x", "1", "2", "3.5"]) == A(x=(1, "2", 3.5)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_fixed_bool() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, bool, bool] + + assert tyro.cli(A, args=["--x", "True", "True", "False"]) == A( + x=(True, True, False) + ) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_variable() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, ...] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=["--x"]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_variable_bool() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, ...] + + assert tyro.cli(A, args=["--x", "True", "True", "False"]) == A( + x=(True, True, False) + ) + assert tyro.cli(A, args=["--x"]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_variable_optional() -> None: + @dataclasses.dataclass + class A: + x: Optional[Tuple[int, ...]] = None + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=["--x"]) == A(x=()) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_sequences() -> None: + @dataclasses.dataclass + class A: + x: Sequence[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_sequences_narrow() -> None: + @dataclasses.dataclass + class A: + x: Sequence = dataclasses.field(default_factory=lambda: [0]) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +def test_sequences_narrow_any() -> None: + @dataclasses.dataclass + class A: + x: Sequence[Any] = dataclasses.field(default_factory=lambda: [0]) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +if sys.version_info >= (3, 9): + + def test_abc_sequences() -> None: + @dataclasses.dataclass + class A: + x: collections.abc.Sequence[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_abc_sequences_narrow() -> None: + @dataclasses.dataclass + class A: + x: collections.abc.Sequence = dataclasses.field(default_factory=lambda: [0]) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +if sys.version_info >= (3, 9): + + def test_abc_sequences_narrow_any() -> None: + @dataclasses.dataclass + class A: + x: collections.abc.Sequence[Any] = dataclasses.field( + default_factory=lambda: [0] + ) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +def test_lists() -> None: + @dataclasses.dataclass + class A: + x: List[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_list_with_literal() -> None: + @dataclasses.dataclass + class A: + x: List[Literal[1, 2, 3]] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3", "4"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_list_with_enums() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class A: + x: List[Color] + + assert tyro.cli(A, args=["--x", "RED", "RED", "BLUE"]) == A( + x=[Color.RED, Color.RED, Color.BLUE] + ) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "RED", "RED", "YELLOW"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_lists_with_default() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class A: + x: List[Color] = dataclasses.field( + default_factory=[Color.RED, Color.GREEN].copy + ) + + assert tyro.cli(A, args=[]) == A(x=[Color.RED, Color.GREEN]) + assert tyro.cli(A, args=["--x", "RED", "GREEN", "BLUE"]) == A( + x=[Color.RED, Color.GREEN, Color.BLUE] + ) + + +def test_lists_bool() -> None: + @dataclasses.dataclass + class A: + x: List[bool] + + assert tyro.cli(A, args=["--x", "True", "False", "True"]) == A( + x=[True, False, True] + ) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_sets() -> None: + @dataclasses.dataclass + class A: + x: Set[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A(x={1, 2, 3}) + assert tyro.cli(A, args=["--x"]) == A(set()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_frozen_sets() -> None: + @dataclasses.dataclass + class A: + x: FrozenSet[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A(x=frozenset({1, 2, 3})) + assert tyro.cli(A, args=["--x"]) == A(x=frozenset()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_deque() -> None: + @dataclasses.dataclass + class A: + x: Deque[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A( + x=collections.deque([1, 2, 3, 3]) + ) + assert tyro.cli(A, args=["--x"]) == A(collections.deque()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_sets_with_default() -> None: + @dataclasses.dataclass + class A: + x: Set[int] = dataclasses.field(default_factory={0, 1, 2}.copy) + + assert tyro.cli(A, args=[]) == A(x={0, 1, 2}) + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A(x={1, 2, 3}) + assert tyro.cli(A, args=["--x"]) == A(x=set()) + + +def test_optional_sequences() -> None: + @dataclasses.dataclass + class A: + x: Optional[Sequence[int]] = None + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + assert tyro.cli(A, args=["--x", "None"]) == A(x=None) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_optional_lists() -> None: + @dataclasses.dataclass + class A: + x: Optional[List[int]] = None + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + assert tyro.cli(A, args=["--x", "None"]) == A(x=None) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_nested_optional_types() -> None: + """We support "None" as a special-case keyword. (note: this is a bit weird because + Optional[str] might interpret "None" as either a string or an actual `None` + value)""" + + @dataclasses.dataclass + class A: + x: Tuple[Optional[int], ...] + + assert tyro.cli(A, args=["--x", "0", "1"]) == A((0, 1)) + assert tyro.cli(A, args=["--x", "0", "None", "1"]) == A((0, None, 1)) + + +def test_union_over_collections() -> None: + def main(a: Union[Tuple[float, ...], Tuple[int, ...]]) -> Any: + return a + + assert tyro.cli(main, args="--a 3.3 3.3 7.0".split(" ")) == (3.3, 3.3, 7.0) + assert tyro.cli(main, args="--a 3 3 7".split(" ")) == (3, 3, 7) + + +def test_union_over_collections_2() -> None: + def main(a: Union[Tuple[str, float, str], Tuple[str, str, float]]) -> Any: + return a + + assert tyro.cli(main, args="--a 3.3 hey 7.0".split(" ")) == ("3.3", "hey", 7.0) + assert tyro.cli(main, args="--a 3 3 hey".split(" ")) == ("3", 3.0, "hey") + + +def test_union_over_collections_3() -> None: + def main(a: Union[Tuple[int, int], Tuple[int, int, int]]) -> Tuple[int, ...]: + return a + + assert tyro.cli(main, args=["--a", "5", "5"]) == (5, 5) + assert tyro.cli(main, args=["--a", "1", "2", "3"]) == (1, 2, 3) + + with pytest.raises(SystemExit): + tyro.cli(main, args=["--a", "5", "5", "2", "1"]) + + with pytest.raises(SystemExit): + tyro.cli(main, args=["--a"]) + with pytest.raises(SystemExit): + tyro.cli(main, args=[]) + + +def test_choices_in_tuples_0() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, bool] + + assert tyro.cli(A, args=["--x", "True", "False"]) == A((True, False)) + + +def test_choices_in_tuples_1() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, Literal["True", "False"]] + + assert tyro.cli(A, args=["--x", "True", "False"]) == A((True, "False")) + + +def test_choices_in_tuples_2() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, Literal["True", "False", "None"]] + + assert tyro.cli(A, args=["--x", "True", "False"]).x == (True, "False") + assert tyro.cli(A, args=["--x", "False", "None"]).x == (False, "None") + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "None", "False"]) + + +def test_nested_tuple_types() -> None: + @dataclasses.dataclass + class A: + x: Tuple[Tuple[int, int], Tuple[str, str]] + + assert tyro.cli(A, args="--x 5 5 5 5".split(" ")).x == ((5, 5), ("5", "5")) + + +def test_variable_nested_tuple() -> None: + def main(x: Tuple[Tuple[int, str], ...]) -> tuple: + return x + + assert tyro.cli(main, args="--x 1 1 2 2".split(" ")) == ((1, "1"), (2, "2")) + with pytest.raises(SystemExit): + tyro.cli(main, args="--x 1 1 2".split(" ")) + + +def test_super_nested() -> None: + def main( + x: Optional[ + List[ + Tuple[ + Optional[int], + Literal[3, 4], + Union[Tuple[int, int], Tuple[str, str]], + ] + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_dict_no_annotation() -> None: + def main(x: Dict[str, Any] = {"int": 5, "str": "5"}): + return x + + assert tyro.cli(main, args=[]) == {"int": 5, "str": "5"} + assert tyro.cli(main, args="--x.int 3 --x.str 7".split(" ")) == { + "int": 3, + "str": "7", + } + + +def test_double_dict_no_annotation() -> None: + def main( + x: Dict[str, Any] = { + "wow": {"int": 5, "str": "5"}, + }, + ): + return x + + assert tyro.cli(main, args=[]) == {"wow": {"int": 5, "str": "5"}} + assert tyro.cli(main, args="--x.wow.int 3 --x.wow.str 7".split(" ")) == { + "wow": { + "int": 3, + "str": "7", + } + } + + +def test_list_narrowing() -> None: + def main(x: list = [0, 1, 2, "hello"]) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", 5] + + +def test_list_narrowing_any() -> None: + def main(x: List[Any] = [0, 1, 2, "hello"]) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", 5] + + +def test_list_narrowing_empty() -> None: + def main(x: list = []) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", "5"] + + +def test_list_narrowing_empty_any() -> None: + def main(x: List[Any] = []) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", "5"] + + +def test_set_narrowing() -> None: + def main(x: set = {0, 1, 2, "hello"}) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", 5} + + +def test_set_narrowing_any() -> None: + def main(x: Set[Any] = {0, 1, 2, "hello"}) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", 5} + + +def test_set_narrowing_empty() -> None: + def main(x: set = set()) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", "5"} + + +def test_set_narrowing_any_empty() -> None: + def main(x: Set[Any] = set()) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", "5"} + + +def test_tuple_narrowing() -> None: + def main(x: tuple = (0, 1, 2, "hello")) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == (0, 1, 2, "3") + + +def test_tuple_narrowing_any() -> None: + def main(x: Tuple[Any, ...] = (0, 1, 2, "hello")) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == (0, 1, 2, "3") + + +def test_tuple_narrowing_empty() -> None: + def main(x: tuple = ()) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == ("0", "1", "2", "3") + + +def test_tuple_narrowing_empty_any() -> None: + def main(x: Tuple[Any, ...] = ()) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == ("0", "1", "2", "3") + + +def test_tuple_narrowing_empty_default() -> None: + def main(x: tuple = ()) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == ("0", "1", "2", "3") + + +def test_narrowing_edge_case() -> None: + """https://github.com/brentyi/tyro/issues/136""" + + @dataclasses.dataclass + class Config: + _target: Type = dataclasses.field(default_factory=lambda: MyClass) + + class MyClass: + def __len__(self): + return 0 + + assert tyro.cli(Config, args=[]) == Config() + + +def test_no_type_collections(): + assert tyro.cli(dict, args="a b c d".split(" ")) == {"a": "b", "c": "d"} + assert tyro.cli(list, args="a b c d".split(" ")) == ["a", "b", "c", "d"] + assert tyro.cli(tuple, args="a b c d".split(" ")) == ("a", "b", "c", "d") + assert tyro.cli(set, args="a b c d".split(" ")) == {"a", "b", "c", "d"} + + +def test_list_narrowing_alt() -> None: + def main(x: list = [1, "1"]) -> list: + return x + + assert tyro.cli(main, args="--x 3 four 5".split(" ")) == [3, "four", 5] + + +def test_list_narrowing_direct() -> None: + assert tyro.cli(list, default=[1, "2"], args="3 four 5".split(" ")) == [ + 3, + "four", + 5, + ] + + +def test_tuple_direct() -> None: + assert tyro.cli(Tuple[int, ...], args="1 2".split(" ")) == (1, 2) # type: ignore + assert tyro.cli(Tuple[int, int], args="1 2".split(" ")) == (1, 2) # type: ignore diff --git a/tests/test_completion.py b/tests/test_completion.py new file mode 100644 index 000000000..44d9e741e --- /dev/null +++ b/tests/test_completion.py @@ -0,0 +1,75 @@ +import contextlib +import dataclasses +import io +from typing import Union + +import pytest +from typing_extensions import Annotated, Literal, Optional + +import tyro + + +# https://github.com/brentyi/tyro/issues/9 +@dataclasses.dataclass(frozen=True) +class Subtype: + data: int = 1 + + +@dataclasses.dataclass(frozen=True) +class TypeA: + subtype: Subtype = Subtype(1) + + +@dataclasses.dataclass(frozen=True) +class TypeB: + subtype: Subtype = Subtype(2) + + +@dataclasses.dataclass(frozen=True) +class Wrapper: + supertype: Union[TypeA, TypeB] = TypeA() + + +def test_bash(): + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(Wrapper, args=["--tyro-print-completion", "bash"]) + assert "# AUTOMATICALLY GENERATED by `shtab`" in target.getvalue() + + +def test_zsh(): + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(Wrapper, args=["--tyro-print-completion", "zsh"]) + assert "# AUTOMATICALLY GENERATED by `shtab`" in target.getvalue() + + +def test_completion_zsh(): + """https://github.com/brentyi/tyro/issues/158""" + + def start_device( + preset: Annotated[ + Optional[Literal["rgb", "depth", "ir"]], tyro.conf.arg(aliases=["-p"]) + ] = None, + frame: Annotated[ + Literal["world", "base"], tyro.conf.arg(aliases=["-f"]) + ] = "world", + ) -> None: + """ + Start device with the given preset. + + :param preset: device preset to use. + :param frame: coordinate frame to use. + """ + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(start_device, args=["--tyro-print-completion", "bash"]) + + completion_script = target.getvalue() + print(completion_script) + assert "# AUTOMATICALLY GENERATED by `shtab`" in completion_script + assert "preset_choices=(" in completion_script + assert "p_choices=(" in completion_script + assert "frame_choices=(" in completion_script + assert "f_choices=(" in completion_script diff --git a/tests/test_conf.py b/tests/test_conf.py new file mode 100644 index 000000000..6b72da9af --- /dev/null +++ b/tests/test_conf.py @@ -0,0 +1,1473 @@ +import argparse +import contextlib +import dataclasses +import io +import json as json_ +import shlex +import sys +from typing import Any, Dict, Generic, List, Tuple, Type, TypeVar, Union + +import pytest +from helptext_utils import get_helptext_with_checks +from typing_extensions import Annotated + +import tyro + + +def test_suppress_subcommand() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + bc: tyro.conf.Suppress[ + Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + ] = dataclasses.field(default_factory=DefaultInstanceHTTPServer) + + assert "bc" not in get_helptext_with_checks(DefaultInstanceSubparser) + + +def test_omit_subcommand_prefix() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + bc: tyro.conf.OmitSubcommandPrefixes[ + Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + ] + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=[ + "--x", + "1", + "default-instance-http-server", + "--y", + "5", + "--no-flag", + ], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "default-instance-http-server", "--y", "5"], + default=DefaultInstanceSubparser( + x=1, bc=DefaultInstanceHTTPServer(y=3, flag=False) + ), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5, flag=False)) + ) + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "default-instance-http-server", "--y", "8"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "default-instance-http-server", "--y", "8"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + +def test_avoid_subparser_with_default() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + bc: tyro.conf.AvoidSubcommands[ + Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + ] + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "5"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "--bc.y", "5"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=3)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)) + ) + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "8"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--bc.y", "8"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + +def test_avoid_subparser_with_default_recursive() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "5"], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.AvoidSubcommands[DefaultInstanceSubparser], + args=["--x", "1", "--bc.y", "5"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=3)), # type: ignore + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultInstanceSubparser, + args=["bc:default-instance-smtp-server", "--bc.z", "3"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceSMTPServer(z=3)) + assert ( + tyro.cli( + tyro.conf.AvoidSubcommands[DefaultInstanceSubparser], + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "8"], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.AvoidSubcommands[DefaultInstanceSubparser], + args=["--bc.y", "8"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), # type: ignore + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + +def test_subparser_in_nested_with_metadata() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + @dataclasses.dataclass + class Nested2: + subcommand: Union[ + Annotated[A, tyro.conf.subcommand("command-a", default=A(7))], + Annotated[B, tyro.conf.subcommand("command-b", default=B(9))], + ] + + @dataclasses.dataclass + class Nested1: + nested2: Nested2 + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_suppressed() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: tyro.conf.Suppress[int] + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + @dataclasses.dataclass + class Nested2: + subcommand: Union[ + Annotated[A, tyro.conf.subcommand("command-a", default=A(7))], + Annotated[B, tyro.conf.subcommand("command-b", default=B(9))], + ] + + @dataclasses.dataclass + class Nested1: + nested2: Nested2 + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_generic() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + T = TypeVar("T") + + @dataclasses.dataclass + class Nested2(Generic[T]): + subcommand: T + + @dataclasses.dataclass + class Nested1: + nested2: Nested2[ + Union[ + Annotated[A, tyro.conf.subcommand("command-a", default=A(7))], + Annotated[B, tyro.conf.subcommand("command-b", default=B(9))], + ] + ] + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_generic_alt() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + T = TypeVar("T") + + @dataclasses.dataclass + class Nested2(Generic[T]): + subcommand: Union[ + Annotated[T, tyro.conf.subcommand("command-a", default=A(7))], + Annotated[B, tyro.conf.subcommand("command-b", default=B(9))], + ] + + @dataclasses.dataclass + class Nested1: + nested2: Nested2[A] + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_default_matching() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + default_one = B(3) + default_three = B(9) + + @dataclasses.dataclass + class Nested: + subcommand: Union[ + Annotated[B, tyro.conf.subcommand("one", default=default_one)], + Annotated[B, tyro.conf.subcommand("two")], + Annotated[B, tyro.conf.subcommand("three", default=default_three)], + ] + + # Match by hash. + def main_one(x: Nested = Nested(default_one)) -> None: + pass + + assert "default: x.subcommand:one" in get_helptext_with_checks(main_one) + + # Match by value. + def main_two(x: Nested = Nested(B(9))) -> None: + pass + + assert "default: x.subcommand:three" in get_helptext_with_checks(main_two) + + # Match by type. + def main_three(x: Nested = Nested(B(15))) -> None: + pass + + assert "default: x.subcommand:one" in get_helptext_with_checks(main_three) + + +def test_flag() -> None: + """When boolean flags have no default value, they must be explicitly specified.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(False), + ) == A(True) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + assert tyro.cli( + A, + args=["--x", "True"], + default=A(False), + config=(tyro.conf.FlagConversionOff,), + ) == A(True) + + +def test_fixed() -> None: + """When an argument is fixed, we shouldn't be able to override it from the CLI.""" + + @dataclasses.dataclass + class A: + x: tyro.conf.Fixed[bool] + + assert tyro.cli( + A, + args=[], + default=A(True), + ) == A(True) + + with pytest.raises(SystemExit): + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + + with pytest.raises(SystemExit): + assert tyro.cli( + A, + args=["--x", "True"], + default=A(False), # type: ignore + config=(tyro.conf.FlagConversionOff,), + ) == A(True) + + +def test_fixed_recursive() -> None: + """When an argument is fixed, we shouldn't be able to override it from the CLI.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(False), + ) == A(True) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + with pytest.raises(SystemExit): + assert tyro.cli( + tyro.conf.Fixed[tyro.conf.FlagConversionOff[A]], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + + +def test_suppressed_group() -> None: + """Reproduction of https://github.com/nerfstudio-project/nerfstudio/issues/882.""" + + @dataclasses.dataclass + class Inner: + a: int + b: int + + def main( + value: int, + inner: tyro.conf.Suppress[Inner] = Inner(1, 2), + ) -> int: + return value + inner.a + inner.b + + assert tyro.cli(main, args=["--value", "5"]) == 8 + + +def test_fixed_group() -> None: + """Inspired by https://github.com/nerfstudio-project/nerfstudio/issues/882.""" + + @dataclasses.dataclass + class Inner: + a: int + b: int + + def main( + value: int, + inner: tyro.conf.Fixed[Inner] = Inner(1, 2), + ) -> int: + return value + inner.a + inner.b + + assert tyro.cli(main, args=["--value", "5"]) == 8 + + +def test_fixed_suppressed_group() -> None: + """Reproduction of https://github.com/nerfstudio-project/nerfstudio/issues/882.""" + + @dataclasses.dataclass + class Inner: + a: int + b: int + + def main( + value: int, + inner: tyro.conf.Fixed[Inner] = Inner(1, 2), + ) -> int: + return value + inner.a + inner.b + + assert tyro.cli(main, args=["--value", "5"]) == 8 + + +def test_suppressed() -> None: + @dataclasses.dataclass + class Struct: + a: int = 5 + b: tyro.conf.Suppress[str] = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" in helptext + assert "--x.b" not in helptext + + +def test_suppress_manual_fixed() -> None: + @dataclasses.dataclass + class Struct: + a: int = 5 + b: tyro.conf.SuppressFixed[tyro.conf.Fixed[str]] = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" in helptext + assert "--x.b" not in helptext + + +def test_suppress_manual_fixed_one_arg_only() -> None: + @dataclasses.dataclass + class Struct: + b: tyro.conf.SuppressFixed[tyro.conf.Fixed[str]] = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" not in helptext + assert "--x.b" not in helptext + + +def test_suppress_auto_fixed() -> None: + @dataclasses.dataclass + class Struct: + a: int = 5 + + def b(self, x): + return 5 + + def main(x: tyro.conf.SuppressFixed[Any] = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" in helptext + assert "--x.b" not in helptext + + +def test_argconf_help() -> None: + @dataclasses.dataclass + class Struct: + a: Annotated[ + int, tyro.conf.arg(name="nice", help="Hello world", metavar="NUMBER") + ] = 5 + b: tyro.conf.Suppress[str] = "7" + + def main(x: Any = Struct()) -> int: + return x.a + + helptext = get_helptext_with_checks(main) + assert "Hello world" in helptext + assert "INT" not in helptext + assert "NUMBER" in helptext + assert "--x.a" not in helptext + assert "--x.nice" in helptext + assert "--x.b" not in helptext + + assert tyro.cli(main, args=[]) == 5 + assert tyro.cli(main, args=["--x.nice", "3"]) == 3 + + +def test_argconf_no_prefix_help() -> None: + @dataclasses.dataclass + class Struct: + a: Annotated[ + int, + tyro.conf.arg( + name="nice", help="Hello world", metavar="NUMBER", prefix_name=False + ), + ] = 5 + b: tyro.conf.Suppress[str] = "7" + + def main(x: Any = Struct()) -> int: + return x.a + + helptext = get_helptext_with_checks(main) + assert "Hello world" in helptext + assert "INT" not in helptext + assert "NUMBER" in helptext + assert "--x.a" not in helptext + assert "--x.nice" not in helptext + assert "--nice" in helptext + assert "--x.b" not in helptext + + assert tyro.cli(main, args=[]) == 5 + with pytest.raises(SystemExit): + assert tyro.cli(main, args=["--x.nice", "3"]) == 3 + assert tyro.cli(main, args=["--nice", "3"]) == 3 + + +def test_positional() -> None: + def main(x: tyro.conf.Positional[int], y: int) -> int: + return x + y + + assert tyro.cli(main, args="5 --y 3".split(" ")) == 8 + assert tyro.cli(main, args="--y 3 5".split(" ")) == 8 + + +def test_positional_required_args() -> None: + @dataclasses.dataclass + class Args: + x: int + y: int = 3 + + assert tyro.cli( + tyro.conf.PositionalRequiredArgs[Args], args="5 --y 3".split(" ") + ) == Args(5, 3) + assert tyro.cli( + tyro.conf.PositionalRequiredArgs[Args], args="--y 3 5".split(" ") + ) == Args(5, 3) + + +def test_positional_order_swap() -> None: + def main(x: int, y: tyro.conf.Positional[int]) -> int: + return x + y + + assert tyro.cli(main, args="5 --x 3".split(" ")) == 8 + assert tyro.cli(main, args="--x 3 5".split(" ")) == 8 + + +def test_omit_subcommand_prefix_and_consolidate_subcommand_args() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + bc: tyro.conf.OmitSubcommandPrefixes[ + Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + ] = dataclasses.field(default_factory=DefaultInstanceHTTPServer) + + assert ( + tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "5", + "--no-flag", + ], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "5", + ], + default=DefaultInstanceSubparser( # type: ignore + x=1, bc=DefaultInstanceHTTPServer(y=3, flag=False) + ), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5, flag=False)) + ) + assert ( + tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "8", + ], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "8", + ], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), # type: ignore + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + # Despite all defaults being set, a subcommand should be required. + with pytest.raises(SystemExit): + tyro.cli(tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], args=[]) + + +def test_omit_subcommand_prefix_and_consolidate_subcommand_args_in_function() -> None: + @tyro.conf.configure(tyro.conf.subcommand(name="http-server")) + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + + @tyro.conf.configure( + tyro.conf.OmitSubcommandPrefixes, + tyro.conf.ConsolidateSubcommandArgs, + ) + def func(parent: DefaultInstanceSubparser) -> DefaultInstanceSubparser: + return parent + + assert tyro.cli( + func, + args=[ + "http-server", + "--parent.x", + "1", + # --y and --no-flag are in a subcommand with prefix omission. + "--y", + "5", + "--no-flag", + ], + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5, flag=False)) + assert tyro.cli( + func, + args=[ + "http-server", + "--parent.x", + "1", + # --y is in a subcommand with prefix omission. + "--y", + "8", + ], + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + + +def test_append_lists() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[List[int]] + + assert tyro.cli(A, args=[]) == A(x=[]) + assert tyro.cli(A, args="--x 1 --x 2 --x 3".split(" ")) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_tuple() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[int, ...]] + + assert tyro.cli(A, args=[]) == A(x=()) + assert tyro.cli(A, args="--x 1 --x 2 --x 3".split(" ")) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_tuple_with_default() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[int, ...]] = (1, 2) + + assert tyro.cli(A, args="--x 1 --x 2 --x 3".split(" ")) == A(x=(1, 2, 1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=(1, 2)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_nested_tuple_fixed_length() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, int], ...]] + + assert tyro.cli(A, args="--x 1 1 --x 2 2 --x 3 3".split(" ")) == A( + x=(("1", 1), ("2", 2), ("3", 3)) + ) + assert tyro.cli(A, args=[]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_nested_tuple_with_default_fixed_length() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, int], ...]] = (("1", 1), ("2", 2)) + + assert tyro.cli(A, args="--x 1 1 --x 2 2 --x 3 3".split(" ")) == A( + x=(("1", 1), ("2", 2), ("1", 1), ("2", 2), ("3", 3)) + ) + assert tyro.cli(A, args=[]) == A(x=(("1", 1), ("2", 2))) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_dict() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Dict[str, int]] + + assert tyro.cli(A, args="--x 1 1 --x 2 2 --x 3 3".split(" ")) == A( + x={"1": 1, "2": 2, "3": 3} + ) + assert tyro.cli(A, args=[]) == A(x={}) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_dict_with_default() -> None: + """Append has no impact when a dictionary has a default value.""" + + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Dict[str, int]] = dataclasses.field( + default_factory=lambda: {"1": 1} + ) + + assert tyro.cli(A, args=[]) == A(x={"1": 1}) + assert tyro.cli(A, args=["--x.1", "2"]) == A(x={"1": 2}) + with pytest.raises(SystemExit): + assert tyro.cli(A, args="--x 2 2 --x 3 3".split(" ")) == A( + x={"1": 1, "2": 2, "3": 3} + ) + + +def test_append_nested_tuple() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, ...], ...]] + + assert tyro.cli(A, args="--x 1 2 3 --x 4 5".split(" ")) == A( + x=(("1", "2", "3"), ("4", "5")) + ) + assert tyro.cli(A, args=[]) == A(x=()) + + +def test_append_nested_tuple_with_default() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, ...], ...]] = (("1", "2"),) + + assert tyro.cli(A, args="--x 1 2 3 --x 4 5".split(" ")) == A( + x=(("1", "2"), ("1", "2", "3"), ("4", "5")) + ) + assert tyro.cli(A, args=[]) == A(x=(("1", "2"),)) + + +def test_append_nested_list() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[List[List[int]]] + + assert tyro.cli(A, args="--x 1 2 3 --x 4 5".split(" ")) == A(x=[[1, 2, 3], [4, 5]]) + assert tyro.cli(A, args=[]) == A(x=[]) + + +def test_append_nested_dict() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[List[Dict[str, int]]] + + assert tyro.cli(A, args="--x 1 2 3 4 --x 4 5".split(" ")) == A( + x=[{"1": 2, "3": 4}, {"4": 5}] + ) + assert tyro.cli(A, args=[]) == A(x=[]) + + +def test_append_nested_dict_double() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Dict[str, Dict[str, int]]] + + assert tyro.cli(A, args="--x 0 1 2 3 4 --x 4 5 6".split(" ")) == A( + x={"0": {"1": 2, "3": 4}, "4": {"5": 6}} + ) + assert tyro.cli(A, args=[]) == A(x={}) + + +def test_duplicated_arg() -> None: + # Loosely inspired by: https://github.com/brentyi/tyro/issues/49 + @dataclasses.dataclass + class ModelConfig: + num_slots: Annotated[int, tyro.conf.arg(prefix_name=False)] + + @dataclasses.dataclass + class TrainConfig: + num_slots: int + model: ModelConfig + + with pytest.raises(argparse.ArgumentError): + tyro.cli(TrainConfig, args="--num-slots 3".split(" ")) + + +def test_omit_arg_prefixes() -> None: + # Loosely inspired by: https://github.com/brentyi/tyro/issues/49 + @dataclasses.dataclass + class ModelConfig: + num_slots: int + + @dataclasses.dataclass + class TrainConfig: + model: ModelConfig + + assert tyro.cli( + tyro.conf.OmitSubcommandPrefixes[TrainConfig], + args="--model.num-slots 3".split(" "), + ) == TrainConfig(ModelConfig(num_slots=3)) + + assert tyro.cli( + tyro.conf.OmitArgPrefixes[TrainConfig], args="--num-slots 3".split(" ") + ) == TrainConfig(ModelConfig(num_slots=3)) + + +def test_custom_constructor_0() -> None: + def times_two(n: str) -> int: + return int(n) * 2 + + @dataclasses.dataclass + class Config: + x: Annotated[int, tyro.conf.arg(name="x-renamed", constructor=times_two)] + + assert tyro.cli(Config, args="--x-renamed.n 5".split(" ")) == Config(x=10) + + +def test_custom_constructor_1() -> None: + def times_two(n: int) -> int: + return int(n) * 2 + + @dataclasses.dataclass + class Config: + x: Annotated[int, tyro.conf.arg(constructor=times_two)] + + assert tyro.cli(Config, args="--x.n 5".split(" ")) == Config(x=10) + + +def test_custom_constructor_2() -> None: + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=int)] + + assert tyro.cli(Config, args="--x 5".split(" ")) == Config(x=5) + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x 5.23".split(" ")) + + +def test_custom_constructor_3() -> None: + def dict_from_json(json: str) -> dict: + out = json_.loads(json) + if not isinstance(out, dict): + raise ValueError(f"{json} is not a dict!") + return out + + @dataclasses.dataclass + class Config: + x: Annotated[ + dict, + tyro.conf.arg( + metavar="JSON", + constructor=dict_from_json, + ), + ] + + assert tyro.cli( + Config, args=shlex.split('--x.json \'{"hello": "world"}\'') + ) == Config(x={"hello": "world"}) + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.json 5".split(" ")) + + error = target.getvalue() + assert "Error parsing x: 5 is not a dict!" in error + + +def test_custom_constructor_4() -> None: + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=int)] = 3.23 + + assert tyro.cli(Config, args="--x 5".split(" ")) == Config(x=5) + assert tyro.cli(Config, args=[]) == Config(x=3.23) + + +def test_custom_constructor_5() -> None: + def make_float(a: float, b: float, c: float = 3) -> float: + return a * b * c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli(Config, args="--x.a 5 --x.b 2 --x.c 3".split(" ")) == Config(x=30) + assert tyro.cli(Config, args="--x.a 5 --x.b 2".split(" ")) == Config(x=30) + + # --x.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.a 5".split(" ")) + + # --x.a and --x.b are required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.c 5".split(" ")) + + +def test_custom_constructor_6() -> None: + def make_float( + a: tyro.conf.Positional[float], + b2: Annotated[float, tyro.conf.arg(name="b")], + c: float = 3, + ) -> float: + return a * b2 * c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli(Config, args="--x.b 2 --x.c 3 5".split(" ")) == Config(x=30) + assert tyro.cli(Config, args="--x.b 2 5".split(" ")) == Config(x=30) + + # --x.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="5".split(" ")) + + # --x.a and --x.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.c 5".split(" ")) + error = target.getvalue() + assert "either all arguments must be provided" in error + assert "or none of them" in error + assert "We're missing arguments" in error + + +def test_custom_constructor_7() -> None: + @dataclasses.dataclass + class Struct: + a: int + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli( + Config, args="--x.struct.a 5 --x.struct.b 2 --x.struct.c 3".split(" ") + ) == Config(x=30) + assert tyro.cli(Config, args="--x.struct.a 5 --x.struct.b 2".split(" ")) == Config( + x=30 + ) + + # --x.struct.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.struct.a 5".split(" ")) + + # --x.struct.a and --x.struct.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.struct.c 5".split(" ")) + error = target.getvalue() + assert "We're missing arguments" in error + assert "'--x.struct.b'" in error + assert "'--x.struct.a'" in error # The 5 is parsed into `a`. + + +def test_custom_constructor_8() -> None: + @dataclasses.dataclass + class Struct: + a: tyro.conf.Positional[int] + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 5".split(" ") + ) == Config(x=30) + assert tyro.cli(Config, args="--x.struct.b 2 5".split(" ")) == Config(x=30) + + # --x.struct.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="5".split(" ")) + + # --x.struct.a and --x.struct.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.struct.b 5".split(" ")) + error = target.getvalue() + assert "We're missing arguments" in error + assert "'x.struct.a'" in error + assert "'--x.struct.b'" not in error + + +def test_custom_constructor_9() -> None: + def commit(branch: str) -> int: + """Commit""" + print(f"commit branch={branch}") + return 3 + + assert ( + tyro.cli( + Annotated[Any, tyro.conf.arg(constructor=commit)], # type: ignore + args="--branch 5".split(" "), + ) + == 3 + ) + + +def test_custom_constructor_10() -> None: + def commit(branch: str) -> int: + """Commit""" + print(f"commit branch={branch}") + return 3 + + def inner(x: Annotated[Any, tyro.conf.arg(constructor=commit)]) -> None: + return x + + def inner_no_prefix( + x: Annotated[Any, tyro.conf.arg(constructor=commit, prefix_name=False)], + ) -> None: + return x + + def outer(x: Annotated[Any, tyro.conf.arg(constructor=inner)]) -> None: + return x + + def outer_no_prefix( + x: Annotated[Any, tyro.conf.arg(constructor=inner_no_prefix)], + ) -> None: + return x + + assert ( + tyro.cli( + outer, + args="--x.x.branch 5".split(" "), + ) + == 3 + ) + assert ( + tyro.cli( + outer_no_prefix, + args="--x.branch 5".split(" "), + ) + == 3 + ) + + +def test_alias() -> None: + """Arguments with aliases.""" + + @dataclasses.dataclass + class Struct: + a: Annotated[int, tyro.conf.arg(aliases=["--all", "-d"])] + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 --x.struct.a 5".split(" ") + ) == Config(x=30) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 -d 5".split(" ") + ) == Config(x=30) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 --all 5".split(" ") + ) == Config(x=30) + assert tyro.cli(Config, args="--x.struct.b 2 --x.struct.a 5".split(" ")) == Config( + x=30 + ) + + # --x.struct.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.struct.a 5".split(" ")) + + # --x.struct.a and --x.struct.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.struct.b 5".split(" ")) + error = target.getvalue() + assert "We're missing arguments" in error + assert "'--x.struct.a'" in error + assert "'--x.struct.b'" not in error + + assert "--x.struct.a INT, --all INT, -d INT" in get_helptext_with_checks(Config) + + +def test_positional_alias() -> None: + """Positional arguments with aliases (which will be ignored).""" + + @dataclasses.dataclass + class Struct: + a: Annotated[tyro.conf.Positional[int], tyro.conf.arg(aliases=["--all", "-d"])] + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + with pytest.warns(UserWarning): + assert tyro.cli(Config, args=[]) == Config(x=3.23) + with pytest.warns(UserWarning): + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 5".split(" ") + ) == Config(x=30) + + with pytest.raises(SystemExit), pytest.warns(UserWarning): + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 -d 5".split(" ") + ) == Config(x=30) + with pytest.raises(SystemExit), pytest.warns(UserWarning): + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 --all 5".split(" ") + ) == Config(x=30) + + +def test_flag_alias() -> None: + @dataclasses.dataclass + class Struct: + flag: Annotated[bool, tyro.conf.arg(aliases=["-f", "--flg"])] = False + + assert tyro.cli(Struct, args=[]).flag is False + assert tyro.cli(Struct, args="--flag".split(" ")).flag is True + assert tyro.cli(Struct, args="--no-flag".split(" ")).flag is False + assert tyro.cli(Struct, args="--flg".split(" ")).flag is True + assert tyro.cli(Struct, args="--no-flg".split(" ")).flag is False + assert tyro.cli(Struct, args="-f".split(" ")).flag is True + + # BooleanOptionalAction will ignore arguments that aren't prefixed with --. + with pytest.raises(SystemExit): + tyro.cli(Struct, args="-no-f".split(" ")) + + +def test_subcommand_constructor_mix() -> None: + """https://github.com/brentyi/tyro/issues/89""" + + def checkout(branch: str) -> str: + """Check out a branch.""" + return branch + + def commit(message: str, all: bool = False) -> str: + """Make a commit.""" + return f"{message} {all}" + + @dataclasses.dataclass + class Arg: + foo: int = 1 + + t: Any = Annotated[ + Union[ + Annotated[ + Any, + tyro.conf.subcommand(name="checkout-renamed", constructor=checkout), + ], + Annotated[ + Any, + tyro.conf.subcommand(name="commit", constructor=commit), + tyro.conf.FlagConversionOff, + ], + Arg, + ], + tyro.conf.OmitArgPrefixes, # Should do nothing. + ] + + assert tyro.cli(t, args=["arg"]) == Arg() + assert tyro.cli(t, args=["checkout-renamed", "--branch", "main"]) == "main" + assert tyro.cli(t, args=["commit", "--message", "hi", "--all", "True"]) == "hi True" + + +def test_merge() -> None: + """Test effect of tyro.conf.arg() on nested structures by approximating an + HfArgumentParser-style API.""" + + T = TypeVar("T") + + # We could add a lot overloads here if we were doing this for real. :) + def instantiate_dataclasses( + classes: Tuple[Type[T], ...], args: List[str] + ) -> Tuple[T, ...]: + return tyro.cli( + tyro.conf.OmitArgPrefixes[ # type: ignore + # Convert (type1, type2) into Tuple[type1, type2] + Tuple.__getitem__( # type: ignore + tuple(Annotated[c, tyro.conf.arg(name=c.__name__)] for c in classes) + ) + ], + args=args, + ) + + @dataclasses.dataclass(frozen=True) + class OptimizerConfig: + lr: float = 1e-4 + weight: int = 10 + + @dataclasses.dataclass(frozen=True) + class DatasetConfig: + batch_size: int = 1 + shuffle: bool = False + + assert instantiate_dataclasses( + (OptimizerConfig, DatasetConfig), args=["--lr", "1e-3"] + ) == ( + OptimizerConfig(1e-3), + DatasetConfig(), + ) + assert instantiate_dataclasses( + (OptimizerConfig, DatasetConfig), args=["--lr", "1e-3", "--shuffle"] + ) == ( + OptimizerConfig(1e-3), + DatasetConfig(shuffle=True), + ) + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + instantiate_dataclasses((OptimizerConfig, DatasetConfig), args=["--help"]) + helptext = target.getvalue() + assert "OptimizerConfig options" in helptext + assert "DatasetConfig options" in helptext + + +def test_counter_action() -> None: + def main( + verbosity: tyro.conf.UseCounterAction[int], + aliased_verbosity: Annotated[ + tyro.conf.UseCounterAction[int], tyro.conf.arg(aliases=["-v"]) + ], + ) -> Tuple[int, int]: + """Example showing how to use counter actions. + Args: + verbosity: Verbosity level. + aliased_verbosity: Same as above, but can also be specified with -v, -vv, -vvv, etc. + """ + return verbosity, aliased_verbosity + + assert tyro.cli(main, args=[]) == (0, 0) + assert tyro.cli(main, args="--verbosity --verbosity".split(" ")) == (2, 0) + assert tyro.cli(main, args="--verbosity --verbosity -v".split(" ")) == (2, 1) + if sys.version_info >= (3, 8): + # Doesn't work in Python 3.7 because of argparse limitations. + assert tyro.cli(main, args="--verbosity --verbosity -vv".split(" ")) == (2, 2) + assert tyro.cli(main, args="--verbosity --verbosity -vvv".split(" ")) == (2, 3) diff --git a/tests/test_dcargs.py b/tests/test_dcargs.py index 5d877c078..723733a1d 100644 --- a/tests/test_dcargs.py +++ b/tests/test_dcargs.py @@ -1,15 +1,40 @@ +import argparse +import copy import dataclasses +import datetime import enum +import os import pathlib -from typing import ClassVar, List, Optional, Sequence, Tuple, Union +from typing import ( + Any, + AnyStr, + Callable, + ClassVar, + Dict, + List, + Optional, + Tuple, + TypeVar, + Union, +) import pytest -from typing_extensions import Annotated, Final, Literal # Backward compatibility. +import torch +from typing_extensions import Annotated, Final, Literal, TypeAlias -import dcargs +import tyro -def test_basic(): +def test_no_args() -> None: + def main() -> int: + return 5 + + assert tyro.cli(main, args=[]) == 5 + with pytest.raises(SystemExit): + tyro.cli(main, args=["3"]) + + +def test_basic() -> None: @dataclasses.dataclass class ManyTypes: i: int @@ -17,9 +42,85 @@ class ManyTypes: f: float p: pathlib.Path - assert ( - dcargs.parse( - ManyTypes, + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypes, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + # We can directly pass a function to `tyro.cli()`: + def function(i: int, s: str, f: float, p: pathlib.Path) -> ManyTypes: + return ManyTypes(i=i, s=s, f=f, p=p) + + assert tyro.cli( + function, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + # We can directly pass a generic class to `tyro.cli()`: + class Wrapper: + def __init__(self, i: int, s: str, f: float, p: pathlib.Path): + self.inner = ManyTypes(i=i, s=s, f=f, p=p) + + assert tyro.cli( + Wrapper, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ).inner == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + +def test_init_false() -> None: + @dataclasses.dataclass + class InitFalseDataclass: + i: int + s: str + f: float + dir: pathlib.Path + ignored: str = dataclasses.field(default="hello", init=False) + + assert tyro.cli( + InitFalseDataclass, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--dir", + "~", + ], + ) == InitFalseDataclass(i=5, s="5", f=5.0, dir=pathlib.Path("~")) + + with pytest.raises(SystemExit): + tyro.cli( + InitFalseDataclass, args=[ "--i", "5", @@ -27,24 +128,24 @@ class ManyTypes: "5", "--f", "5", - "--p", + "--dir", "~", + "--ignored", + "blah", ], ) - == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) - ) -def test_required(): +def test_required() -> None: @dataclasses.dataclass class A: x: int with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) + tyro.cli(A, args=[]) -def test_flag(): +def test_flag() -> None: """When boolean flags have no default value, they must be explicitly specified.""" @dataclasses.dataclass @@ -52,40 +153,44 @@ class A: x: bool with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) + tyro.cli(A, args=[]) - assert dcargs.parse(A, args=["--x", "1"]) == A(True) - assert dcargs.parse(A, args=["--x", "true"]) == A(True) - assert dcargs.parse(A, args=["--x", "True"]) == A(True) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "true"]) + assert tyro.cli(A, args=["--x", "True"]) == A(True) - assert dcargs.parse(A, args=["--x", "0"]) == A(False) - assert dcargs.parse(A, args=["--x", "false"]) == A(False) - assert dcargs.parse(A, args=["--x", "False"]) == A(False) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "0"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "false"]) + assert tyro.cli(A, args=["--x", "False"]) == A(False) -def test_flag_default_false(): +def test_flag_default_false() -> None: """When boolean flags default to False, a --flag-name flag must be passed in to flip it to True.""" @dataclasses.dataclass class A: x: bool = False - assert dcargs.parse(A, args=[]) == A(False) - assert dcargs.parse(A, args=["--x"]) == A(True) + assert tyro.cli(A, args=[]) == A(False) + assert tyro.cli(A, args=["--x"]) == A(True) -def test_flag_default_true(): +def test_flag_default_true() -> None: """When boolean flags default to True, a --no-flag-name flag must be passed in to flip it to False.""" @dataclasses.dataclass class A: x: bool = True - assert dcargs.parse(A, args=[]) == A(True) - assert dcargs.parse(A, args=["--no-x"]) == A(False) + assert tyro.cli(A, args=[]) == A(True) + assert tyro.cli(A, args=["--no-x"]) == A(False) -def test_flag_default_true_nested(): +def test_flag_default_true_nested() -> None: """When boolean flags default to True, a --no-flag-name flag must be passed in to flip it to False.""" @dataclasses.dataclass @@ -96,200 +201,257 @@ class NestedDefaultTrue: class A: x: NestedDefaultTrue - assert dcargs.parse(A, args=[]) == A(NestedDefaultTrue(True)) - assert dcargs.parse(A, args=["--x.no-x"]) == A(NestedDefaultTrue(False)) + assert tyro.cli(A, args=[]) == A(NestedDefaultTrue(True)) + assert tyro.cli(A, args=["--x.no-x"]) == A(NestedDefaultTrue(False)) -def test_default(): +def test_default() -> None: @dataclasses.dataclass class A: x: int = 5 - assert dcargs.parse(A, args=[]) == A() + assert tyro.cli(A, args=[]) == A() -def test_optional(): +def test_default_factory() -> None: @dataclasses.dataclass class A: - x: Optional[int] + x: int = dataclasses.field(default_factory=lambda: 5) - assert dcargs.parse(A, args=[]) == A(x=None) + assert tyro.cli(A, args=[]) == A() -def test_sequences(): +def test_optional() -> None: @dataclasses.dataclass class A: - x: Sequence[int] + x: Optional[int] = None - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) - with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_union_basic() -> None: + def main(x: Union[int, str]) -> Union[int, str]: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_union_with_list() -> None: + def main(x: Union[int, str, List[bool]]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + assert tyro.cli(main, args=["--x", "True"]) == "True" + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_union_literal() -> None: + def main(x: Union[Literal[1, 2], Literal[3, 4, 5], str]) -> Union[int, str]: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == "6" + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_func_typevar() -> None: + T = TypeVar("T", int, str) + + def main(x: T) -> T: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_func_typevar_bound() -> None: + T = TypeVar("T", bound=int) + + def main(x: T) -> T: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) + tyro.cli(main, args=["--x", "five"]) + + +def test_enum() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + @dataclasses.dataclass + class EnumClassA: + color: Color -def test_lists(): @dataclasses.dataclass - class A: - x: List[int] + class EnumClassB: + color: Color = Color.GREEN - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) - with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) - with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) + assert tyro.cli(EnumClassA, args=["--color", "RED"]) == EnumClassA(color=Color.RED) + assert tyro.cli(EnumClassB, args=[]) == EnumClassB() -def test_optional_sequences(): +def test_enum_alias() -> None: + class Color(enum.Enum): + RED = 1 + ROUGE = 1 + @dataclasses.dataclass class A: - x: Optional[Sequence[int]] + color: Color - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) - with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) - assert dcargs.parse(A, args=[]) == A(x=None) + assert tyro.cli(A, args=["--color", "RED"]) == A(color=Color.RED) + assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.ROUGE) + assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.RED) -def test_optional_lists(): +def test_literal() -> None: @dataclasses.dataclass class A: - x: Optional[List[int]] + x: Literal[0, 1, 2] - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) - assert dcargs.parse(A, args=[]) == A(x=None) + assert tyro.cli(A, args=["--x", "3"]) -def test_tuples_fixed(): +def test_literal_none() -> None: @dataclasses.dataclass class A: - x: Tuple[int, int, int] + x: Literal[0, 1, None, 2] - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) + assert tyro.cli(A, args=["--x", "None"]) == A(x=None) with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) - with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) + assert tyro.cli(A, args=["--x", "3"]) + +Choices = int +Choices = tyro.extras.literal_type_from_choices([0, 1, 2]) # type: ignore -def test_tuples_variable(): + +def test_dynamic_literal() -> None: @dataclasses.dataclass class A: - x: Tuple[int, ...] + x: Choices # type: ignore - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) - with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) + assert tyro.cli(A, args=["--x", "3"]) -def test_tuples_variable_optional(): - @dataclasses.dataclass - class A: - x: Optional[Tuple[int, ...]] +def test_literal_bool() -> None: + def main(x: Literal[True]) -> bool: + return x + + assert tyro.cli(main, args=["--x", "True"]) is True + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "False"]) - assert dcargs.parse(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + def main2(x: Literal[True, False]) -> bool: + return x + + assert tyro.cli(main2, args=["--x", "True"]) is True + assert tyro.cli(main2, args=["--x", "False"]) is False with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x"]) - assert dcargs.parse(A, args=[]) == A(x=None) + tyro.cli(main2, args=["--x", "Tru"]) -def test_enum(): +def test_literal_enum() -> None: class Color(enum.Enum): RED = enum.auto() GREEN = enum.auto() BLUE = enum.auto() @dataclasses.dataclass - class EnumClassA: - color: Color - - @dataclasses.dataclass - class EnumClassB: - color: Color = Color.GREEN + class A: + x: Literal[Color.RED, Color.GREEN] - assert dcargs.parse(EnumClassA, args=["--color", "RED"]) == EnumClassA( - color=Color.RED - ) - assert dcargs.parse(EnumClassB) == EnumClassB() + assert tyro.cli(A, args=["--x", "RED"]) == A(x=Color.RED) + assert tyro.cli(A, args=["--x", "GREEN"]) == A(x=Color.GREEN) + with pytest.raises(SystemExit): + assert tyro.cli(A, args=["--x", "BLUE"]) -def test_literal(): +def test_optional_literal() -> None: @dataclasses.dataclass class A: - x: Literal[0, 1, 2] + x: Optional[Literal[0, 1, 2]] = None - assert dcargs.parse(A, args=["--x", "1"]) == A(x=1) + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) with pytest.raises(SystemExit): - assert dcargs.parse(A, args=["--x", "3"]) + assert tyro.cli(A, args=["--x", "3"]) + assert tyro.cli(A, args=[]) == A(x=None) -def test_optional_literal(): - @dataclasses.dataclass - class A: - x: Optional[Literal[0, 1, 2]] +def test_multitype_literal() -> None: + def main(x: Literal[0, "5"]) -> Any: + return x - assert dcargs.parse(A, args=["--x", "1"]) == A(x=1) + assert tyro.cli(main, args=["--x", "0"]) == 0 + assert tyro.cli(main, args=["--x", "5"]) == "5" with pytest.raises(SystemExit): - assert dcargs.parse(A, args=["--x", "3"]) - assert dcargs.parse(A, args=[]) == A(x=None) + tyro.cli(main, args=["--x", "6"]) -def test_annotated(): +def test_annotated() -> None: """Annotated[] is a no-op.""" @dataclasses.dataclass class A: x: Annotated[int, "some label"] = 3 - assert dcargs.parse(A, args=["--x", "5"]) == A(x=5) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) -def test_annotated_optional(): +def test_annotated_optional() -> None: """Annotated[] is a no-op.""" @dataclasses.dataclass class A: x: Annotated[Optional[int], "some label"] = 3 - assert dcargs.parse(A, args=[]) == A(x=3) - assert dcargs.parse(A, args=["--x", "5"]) == A(x=5) + assert tyro.cli(A, args=[]) == A(x=3) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) -def test_optional_annotated(): +def test_optional_annotated() -> None: """Annotated[] is a no-op.""" @dataclasses.dataclass class A: x: Optional[Annotated[int, "some label"]] = 3 - assert dcargs.parse(A, args=[]) == A(x=3) - assert dcargs.parse(A, args=["--x", "5"]) == A(x=5) + assert tyro.cli(A, args=[]) == A(x=3) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) -def test_final(): +def test_final() -> None: """Final[] is a no-op.""" @dataclasses.dataclass class A: x: Final[int] = 3 - assert dcargs.parse(A, args=["--x", "5"]) == A(x=5) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) -def test_final_optional(): +def test_final_optional() -> None: @dataclasses.dataclass class A: x: Final[Optional[int]] = 3 - assert dcargs.parse(A, args=[]) == A(x=3) - assert dcargs.parse(A, args=["--x", "5"]) == A(x=5) + assert tyro.cli(A, args=[]) == A(x=3) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) -def test_classvar(): +def test_classvar() -> None: """ClassVar[] types should be skipped.""" @dataclasses.dataclass @@ -297,104 +459,419 @@ class A: x: ClassVar[int] = 5 with pytest.raises(SystemExit): - dcargs.parse(A, args=["--x", "1"]) - assert dcargs.parse(A, args=[]) == A() + tyro.cli(A, args=["--x", "1"]) + assert tyro.cli(A, args=[]) == A() -def test_nested(): - @dataclasses.dataclass - class B: - y: int +def test_parse_empty_description() -> None: + """If the file has no dosctring, it should be treated as an empty string.""" @dataclasses.dataclass - class Nested: - x: int - b: B + class A: + x: int = 0 - assert dcargs.parse(Nested, args=["--x", "1", "--b.y", "3"]) == Nested( - x=1, b=B(y=3) - ) + assert tyro.cli(A, description=None, args=[]) == A(x=0) + + +SomeTypeAlias: TypeAlias = int + + +def test_type_alias() -> None: + def add(a: SomeTypeAlias, b: SomeTypeAlias) -> SomeTypeAlias: + return a + b + + assert tyro.cli(add, args=["--a", "5", "--b", "7"]) == 12 + + +def test_any() -> None: + def main(x: Any = 5) -> Any: + return x + + assert tyro.cli(main, args=[]) == 5 + + +def test_bytes() -> None: + def main(x: bytes) -> bytes: + return x + + assert tyro.cli(main, args=["--x", "hello"]) == b"hello" + + +def test_any_str() -> None: + def main(x: AnyStr) -> AnyStr: + return x + + # Use bytes when provided ascii-compatible inputs. + assert tyro.cli(main, args=["--x", "hello"]) == b"hello" + assert tyro.cli(main, args=["--x", "hello„"]) == "hello„" + + +def test_fixed() -> None: + def main(x: Callable[[int], int] = lambda x: x * 2) -> Callable[[int], int]: + return x + + assert tyro.cli(main, args=[])(3) == 6 + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "something"]) + + +def test_fixed_dataclass_type() -> None: + def dummy(): + return 5 # noqa + + def main(x: Callable = dummy) -> Callable: + return x + + assert tyro.cli(main, args=[]) is dummy with pytest.raises(SystemExit): - dcargs.parse(Nested, args=["--x", "1"]) + tyro.cli(main, args=["--x", "something"]) -# TODO: implement this! -# def test_optional_nested(): -# @dataclasses.dataclass -# class OptionalNestedChild: -# y: int -# z: int -# -# @dataclasses.dataclass -# class OptionalNested: -# x: int -# b: Optional[OptionalNestedChild] -# -# assert dcargs.parse(OptionalNested, args=["--x", "1"]) == OptionalNested( -# x=1, b=None -# ) -# with pytest.raises(SystemExit): -# dcargs.parse(OptionalNested, args=["--x", "1", "--b.y", "3"]) -# with pytest.raises(SystemExit): -# dcargs.parse(OptionalNested, args=["--x", "1", "--b.z", "3"]) -# -# assert dcargs.parse( -# OptionalNested, args=["--x", "1", "--b.y", "2", "--b.z", "3"] -# ) == OptionalNested(x=1, b=OptionalNestedChild(y=2, z=3)) +def test_missing_singleton() -> None: + assert tyro.MISSING is copy.deepcopy(tyro.MISSING) -def test_subparser(): +def test_torch_device() -> None: + def main(device: torch.device) -> torch.device: + return device + + assert tyro.cli(main, args=["--device", "cpu"]) == torch.device("cpu") + + +def test_supports_inference_mode_decorator() -> None: + @torch.inference_mode() + def main(x: int, device: str) -> Tuple[int, str]: + return x, device + + assert tyro.cli(main, args="--x 3 --device cuda".split(" ")) == (3, "cuda") + + +def test_torch_device_2() -> None: + assert tyro.cli(torch.device, args=["cpu"]) == torch.device("cpu") + + +def test_just_int() -> None: + assert tyro.cli(int, args=["123"]) == 123 + + +def test_just_dict() -> None: + assert tyro.cli(Dict[str, str], args="key value key2 value2".split(" ")) == { + "key": "value", + "key2": "value2", + } + + +def test_just_list() -> None: + assert tyro.cli(List[int], args="1 2 3 4".split(" ")) == [1, 2, 3, 4] + + +def test_just_tuple() -> None: + # Need a type: ignore for mypy. Seems like a mypy bug. + assert tyro.cli(Tuple[int, int, int, int], args="1 2 3 4".split(" ")) == ( # type: ignore + 1, + 2, + 3, + 4, + ) + + +def test_return_parser() -> None: + def main() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + return parser + + assert isinstance(tyro.cli(main, args=[]), argparse.ArgumentParser) + + +def test_pathlike_custom_class() -> None: + class CustomPath(pathlib.PurePosixPath): + def __new__(cls, *args: Union[str, os.PathLike]): + return super().__new__(cls, *args) + + def main(a: CustomPath) -> CustomPath: + return a + + assert tyro.cli(main, args=["--a", "/dev/null"]) == CustomPath("/dev/null") + + +def test_class_with_new_and_no_init() -> None: + class A(object): + def __new__(cls, x: int = 5): + return cls._custom_initializer(x) + + @classmethod + def _custom_initializer(cls, x: int = 5): + self = object.__new__(cls) + self.x = x # type: ignore + return self + + def __eq__(self, other) -> bool: + return self.x == other.x # type: ignore + + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_return_unknown_args() -> None: @dataclasses.dataclass - class HTTPServer: - y: int + class A: + x: int = 0 + + a, unknown_args = tyro.cli( + A, args=["positional", "--x", "5", "--y", "7"], return_unknown_args=True + ) + assert a == A(x=5) + assert unknown_args == ["positional", "--y", "7"] + +def test_unknown_args_with_arg_fixing() -> None: @dataclasses.dataclass - class SMTPServer: - z: int + class A: + x: int = 0 + + a, unknown_args = tyro.cli( + A, + args=["--x", "5", "--a_b", "--a-c"], + return_unknown_args=True, + ) + assert a == A(x=5) + # Should return the unfixed arguments + assert unknown_args == ["--a_b", "--a-c"] + +def test_allow_ambiguous_args_when_not_returning_unknown_args() -> None: @dataclasses.dataclass - class Subparser: - x: int - bc: Union[HTTPServer, SMTPServer] + class A: + a_b: List[int] = dataclasses.field(default_factory=list) + + a = tyro.cli( + A, + args=["--a_b", "5", "--a-b", "7"], + ) + assert a == A(a_b=[7]) - assert dcargs.parse( - Subparser, args=["--x", "1", "http-server", "--y", "3"] - ) == Subparser(x=1, bc=HTTPServer(y=3)) - assert dcargs.parse( - Subparser, args=["--x", "1", "smtp-server", "--z", "3"] - ) == Subparser(x=1, bc=SMTPServer(z=3)) - with pytest.raises(SystemExit): - dcargs.parse(Subparser, args=["--x", "1", "b", "--z", "3"]) - with pytest.raises(SystemExit): - dcargs.parse(Subparser, args=["--x", "1", "c", "--y", "3"]) +def test_disallow_ambiguous_args_when_returning_unknown_args() -> None: + @dataclasses.dataclass + class A: + x: int = 0 + + # If there's an argument that's ambiguous then we should raise an error when we're + # returning unknown args. + with pytest.raises(RuntimeError, match="Ambiguous .* --a_b and --a-b"): + tyro.cli( + A, + args=["--x", "5", "--a_b", "--a-b"], + return_unknown_args=True, + ) -def test_optional_subparser(): +def test_unknown_args_with_consistent_duplicates() -> None: @dataclasses.dataclass - class OptionalHTTPServer: - y: int + class A: + a_b: List[int] = dataclasses.field(default_factory=list) + c_d: List[int] = dataclasses.field(default_factory=list) + + # Tests logic for consistent duplicate arguments when performing argument fixing. + # i.e., we can fix arguments if the separator is consistent (all _'s or all -'s). + a, unknown_args = tyro.cli( + A, + args=[ + "--a-b", + "5", + "--a-b", + "7", + "--c_d", + "5", + "--c_d", + "7", + "--e-f", + "--e-f", + "--g_h", + "--g_h", + ], + return_unknown_args=True, + ) + assert a == A(a_b=[7], c_d=[7]) + assert unknown_args == ["--e-f", "--e-f", "--g_h", "--g_h"] + +def test_pathlike(): + def main(x: os.PathLike) -> os.PathLike: + return x + + assert tyro.cli(main, args=["--x", "/dev/null"]) == pathlib.Path("/dev/null") + + +def test_unpack() -> None: + def main(*args: int, **kwargs: float) -> Tuple[Tuple[int, ...], Dict[str, float]]: + return args, kwargs + + assert tyro.cli(main, args=[]) == ((), {}) + assert tyro.cli(main, args="--args --kwargs".split(" ")) == ((), {}) + assert tyro.cli( + main, args="--args 1 2 3 --kwargs learning_rate 1e-4 beta1 0.99".split(" ") + ) == ((1, 2, 3), {"learning_rate": 1e-4, "beta1": 0.99}) + + +def test_unpack_with_other_args() -> None: + def main( + other1: str, other2: str = "3", *args: int, **kwargs: float + ) -> Tuple[str, str, Tuple[int, ...], Dict[str, float]]: + return other1, other2, args, kwargs + + assert tyro.cli(main, args=["--other1", "hi"]) == ("hi", "3", (), {}) + assert tyro.cli(main, args="--other1 ok --args --kwargs".split(" ")) == ( + "ok", + "3", + (), + {}, + ) + assert tyro.cli( + main, + args="--args 1 2 3 --kwargs learning_rate 1e-4 beta1 0.99 --other1 hello".split( + " " + ), + ) == ("hello", "3", (1, 2, 3), {"learning_rate": 1e-4, "beta1": 0.99}) + + +def test_empty_container() -> None: @dataclasses.dataclass - class OptionalSMTPServer: - z: int + class A: + x: Tuple[int, ...] = (1, 2, 3) + y: Union[int, str, List[bool]] = dataclasses.field( + default_factory=lambda: [False, False, True] + ) + assert tyro.cli(A, args="--x".split(" ")).x == () + assert tyro.cli(A, args="--y".split(" ")).y == [] + + +def test_unknown_args_with_consistent_duplicates_use_underscores() -> None: @dataclasses.dataclass - class OptionalSubparser: - x: int - bc: Optional[Union[OptionalHTTPServer, OptionalSMTPServer]] - - assert dcargs.parse( - OptionalSubparser, args=["--x", "1", "optional-http-server", "--y", "3"] - ) == OptionalSubparser(x=1, bc=OptionalHTTPServer(y=3)) - assert dcargs.parse( - OptionalSubparser, args=["--x", "1", "optional-smtp-server", "--z", "3"] - ) == OptionalSubparser(x=1, bc=OptionalSMTPServer(z=3)) - assert dcargs.parse(OptionalSubparser, args=["--x", "1"]) == OptionalSubparser( - x=1, bc=None + class A: + a_b: List[int] = dataclasses.field(default_factory=list) + c_d: List[int] = dataclasses.field(default_factory=list) + + # Tests logic for consistent duplicate arguments when performing argument fixing. + # i.e., we can fix arguments if the separator is consistent (all _'s or all -'s). + a, unknown_args = tyro.cli( + A, + args=[ + "--a-b", + "5", + "--a-b", + "7", + "--c_d", + "5", + "--c_d", + "7", + "--e-f", + "--e-f", + "--g_h", + "--g_h", + ], + return_unknown_args=True, + use_underscores=True, + ) + assert a == A(a_b=[7], c_d=[7]) + assert unknown_args == ["--e-f", "--e-f", "--g_h", "--g_h"] + + +def test_datetime_parsing() -> None: + assert tyro.cli( + datetime.datetime, args=["2024-08-25T14:30:00"] + ) == datetime.datetime(2024, 8, 25, 14, 30, 0) + assert tyro.cli( + datetime.datetime, args=["2024-08-25T14:30:00.123456"] + ) == datetime.datetime(2024, 8, 25, 14, 30, 0, 123456) + assert tyro.cli( + datetime.datetime, args=["2024-08-25T14:30:00+05:30"] + ) == datetime.datetime( + 2024, 8, 25, 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_date_parsing() -> None: + assert tyro.cli(datetime.date, args=["2024-08-25"]) == datetime.date(2024, 8, 25) + assert tyro.cli(datetime.date, args=["1999-12-31"]) == datetime.date(1999, 12, 31) + + +def test_time_parsing() -> None: + assert tyro.cli(datetime.time, args=["14:30:00"]) == datetime.time(14, 30, 0) + assert tyro.cli(datetime.time, args=["14:30:00.123456"]) == datetime.time( + 14, 30, 0, 123456 + ) + assert tyro.cli(datetime.time, args=["14:30:00+05:30"]) == datetime.time( + 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_datetime_parsing_with_main(): + def main(dt: datetime.datetime) -> datetime.datetime: + return dt + + assert tyro.cli(main, args=["--dt", "2024-08-25T14:30:00"]) == datetime.datetime( + 2024, 8, 25, 14, 30, 0 + ) + assert tyro.cli( + main, args=["--dt", "2024-08-25T14:30:00.123456"] + ) == datetime.datetime(2024, 8, 25, 14, 30, 0, 123456) + assert tyro.cli( + main, args=["--dt", "2024-08-25T14:30:00+05:30"] + ) == datetime.datetime( + 2024, 8, 25, 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_date_parsing_with_main(): + def main(dt: datetime.date) -> datetime.date: + return dt + + assert tyro.cli(main, args=["--dt", "2024-08-25"]) == datetime.date(2024, 8, 25) + assert tyro.cli(main, args=["--dt", "1999-12-31"]) == datetime.date(1999, 12, 31) + + +def test_time_parsing_with_main(): + def main(dt: datetime.time) -> datetime.time: + return dt + + assert tyro.cli(main, args=["--dt", "14:30:00"]) == datetime.time(14, 30, 0) + assert tyro.cli(main, args=["--dt", "14:30:00.123456"]) == datetime.time( + 14, 30, 0, 123456 + ) + assert tyro.cli(main, args=["--dt", "14:30:00+05:30"]) == datetime.time( + 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) ) + +def test_date_parsing_harder_format(): + def main(dt: datetime.date) -> datetime.date: + return dt + + # Wrong separator '/'. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "2024/08/25"]) + + # Day-Month-Year format instead of Year-Month-Day. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "25-08-2024"]) + + +def test_time_parsing_harder_format(): + def main(dt: datetime.time) -> datetime.time: + return dt + + # Wrong separator '/'. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "14/30"]) + + # Missing seconds. + assert tyro.cli(main, args=["--dt", "14:30"]) == datetime.time(14, 30, 0) + + # Invalid minute value. with pytest.raises(SystemExit): - dcargs.parse(OptionalSubparser, args=["--x", "1", "B", "--z", "3"]) + tyro.cli(main, args=["--dt", "14:60:00"]) + + # Invalid hour value. with pytest.raises(SystemExit): - dcargs.parse(OptionalSubparser, args=["--x", "1", "C", "--y", "3"]) + tyro.cli(main, args=["--dt", "25:00:00"]) diff --git a/tests/test_dict_namedtuple.py b/tests/test_dict_namedtuple.py new file mode 100644 index 000000000..6ddc804f9 --- /dev/null +++ b/tests/test_dict_namedtuple.py @@ -0,0 +1,636 @@ +import contextlib +import copy +import dataclasses +import io +import pathlib +from typing import Any, Dict, Mapping, NamedTuple, Tuple, Union, cast + +import pytest +from typing_extensions import Literal, NotRequired, Required, TypedDict + +import tyro +import tyro._strings + + +def test_basic_dict() -> None: + def main(params: Dict[str, int]) -> Dict[str, int]: + return params + + assert tyro.cli(main, args="--params hey 5 hello 2".split(" ")) == { + "hey": 5, + "hello": 2, + } + assert tyro.cli(main, args="--params hey 5 hello 2".split(" ")) == { + "hey": 5, + "hello": 2, + } + assert tyro.cli(main, args="--params".split(" ")) == {} + with pytest.raises(SystemExit): + tyro.cli(main, args="--params hey 5 hello hey".split(" ")) + with pytest.raises(SystemExit): + tyro.cli(main, args="--params hey 5 hello".split(" ")) + + +def test_dict_with_default() -> None: + def main(params: Mapping[Literal[1, 3, 5, 7], bool] = {5: False, 1: True}) -> Any: + return params + + assert tyro.cli(main, args=[]) == {5: False, 1: True} + assert tyro.cli(main, args="--params.5 --params.no-1".split(" ")) == { + 5: True, + 1: False, + } + with pytest.raises(SystemExit): + tyro.cli(main, args="--params".split(" ")) + + +def test_tuple_in_dict() -> None: + def main(x: Dict[Union[Tuple[int, int], Tuple[str, str]], Tuple[int, int]]) -> dict: + return x + + assert tyro.cli(main, args="--x 1 1 2 2 3 3 4 4".split(" ")) == { + (1, 1): (2, 2), + (3, 3): (4, 4), + } + + +def test_basic_typeddict() -> None: + class ManyTypesTypedDict(TypedDict): + i: int + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) + + +def test_positional_in_typeddict() -> None: + class ManyTypesTypedDict(TypedDict): + i: tyro.conf.Positional[int] + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="5".split(" ")) + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) + + +def test_total_false_typeddict() -> None: + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + assert tyro.cli(ManyTypesTypedDict, args=[]) == dict() + assert tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) == dict(i=5) + assert tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) == dict(s="5") + + +def test_total_false_required_typeddict() -> None: + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: Required[str] + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + # `s` is Required[]. + assert tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) == dict(i=5) + assert tyro.cli( + ManyTypesTypedDict, args="--i 5".split(" "), default={"s": "5"} + ) == dict(i=5, s="5") + assert tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) == dict(s="5") + + +def test_total_true_not_required_typeddict() -> None: + class ManyTypesTypedDict(TypedDict, total=True): + i: NotRequired[int] + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + # `s` is Required[]. + assert tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) == dict(i=5) + assert tyro.cli( + ManyTypesTypedDict, args="--i 5".split(" "), default={"s": "5"} + ) == dict(i=5, s="5") + assert tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) == dict(s="5") + + +def test_total_false_nested_typeddict() -> None: + class ChildTypedDict(TypedDict, total=False): + i: int + s: str + + class ParentTypedDict(TypedDict, total=False): + child: ChildTypedDict + + assert tyro.cli( + ParentTypedDict, + args="--child.i 5 --child.s 5".split(" "), + ) == {"child": {"i": 5, "s": "5"}} + + # total=False is ~ignored on the parent. + assert tyro.cli( + ParentTypedDict, + args=[], + ) == {"child": {}} + + +def test_total_false_typeddict_with_nested() -> None: + @dataclasses.dataclass + class Inner: + j: float + + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: Inner + + # --s.j is (unfortunately) still required. + with pytest.raises(SystemExit): + tyro.cli( + ManyTypesTypedDict, + args="".split(" "), + ) + + assert tyro.cli( + ManyTypesTypedDict, + args="--s.j 5".split(" "), + ) == {"s": Inner(5.0)} + + +def test_total_false_typeddict_with_tuple() -> None: + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: Tuple[str, str] + + assert ( + tyro.cli( + ManyTypesTypedDict, + args=[], + ) + == dict() + ) + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5 5".split(" "), + ) == dict(i=5, s=("5", "5")) + + +def test_nested_typeddict() -> None: + class ChildTypedDict(TypedDict): + y: int + + class NestedTypedDict(TypedDict): + x: int + b: ChildTypedDict + + assert tyro.cli(NestedTypedDict, args=["--x", "1", "--b.y", "3"]) == dict( + x=1, b=dict(y=3) + ) + with pytest.raises(SystemExit): + tyro.cli(NestedTypedDict, args=["--x", "1"]) + + +def test_helptext_and_default_typeddict() -> None: + class HelptextTypedDict(TypedDict): + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(HelptextTypedDict, default={"z": 3}, args=["--help"]) + helptext = tyro._strings.strip_ansi_sequences(f.getvalue()) + assert cast(str, HelptextTypedDict.__doc__) in helptext + assert "--x INT" in helptext + assert "--y INT" in helptext + assert "--z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_basic_namedtuple() -> None: + class ManyTypesNamedTuple(NamedTuple): + i: int + s: str + f: float + p: pathlib.Path + + assert tyro.cli( + ManyTypesNamedTuple, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypesNamedTuple(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + +def test_nested_namedtuple() -> None: + class ChildNamedTuple(NamedTuple): + y: int + + class NestedNamedTuple(NamedTuple): + x: int + b: ChildNamedTuple + + assert tyro.cli( + NestedNamedTuple, args=["--x", "1", "--b.y", "3"] + ) == NestedNamedTuple(x=1, b=ChildNamedTuple(y=3)) + with pytest.raises(SystemExit): + tyro.cli(NestedNamedTuple, args=["--x", "1"]) + + +def test_helptext_and_default_namedtuple() -> None: + class HelptextNamedTupleDefault(NamedTuple): + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int = 3 + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(HelptextNamedTupleDefault, args=["--help"]) + helptext = tyro._strings.strip_ansi_sequences(f.getvalue()) + assert cast(str, HelptextNamedTupleDefault.__doc__) in helptext + assert "--x INT" in helptext + assert "--y INT" in helptext + assert "--z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_and_default_namedtuple_alternate() -> None: + class HelptextNamedTuple(NamedTuple): + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int + """Documentation 3""" + + with pytest.raises(SystemExit): + tyro.cli( + HelptextNamedTuple, + default=tyro.MISSING, + args=[], + ) + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli( + HelptextNamedTuple, + default=HelptextNamedTuple( + x=tyro.MISSING, + y=tyro.MISSING, + z=3, + ), + args=["--help"], + ) + helptext = tyro._strings.strip_ansi_sequences(f.getvalue()) + assert cast(str, HelptextNamedTuple.__doc__) in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3" in helptext + assert "(default: 3)" in helptext + + +def test_nested_dict() -> None: + loaded_config = { + "batch_size": 32, + "optimizer": { + "learning_rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule_type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule_type", + "exponential", + ], + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch_size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule_type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_use_underscores() -> None: + loaded_config = { + "batch_size": 32, + "optimizer": { + "learning_rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule_type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule-type", + "exponential", + ], + use_underscores=True, + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch_size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule_type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_hyphen() -> None: + # We do a lot of underscore <=> conversion in the code; this is just to make sure it + # doesn't break anything! + loaded_config = { + "batch-size": 32, + "optimizer": { + "learning-rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule-type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch-size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_hyphen_use_underscores() -> None: + # We do a lot of underscore <=> conversion in the code; this is just to make sure it + # doesn't break anything! + loaded_config = { + "batch-size": 32, + "optimizer": { + "learning-rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule-type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule-type", + "exponential", + ], + use_underscores=True, + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch-size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch_size", + "16", + "--optimizer.scheduler.schedule_type", + "exponential", + ], + use_underscores=True, + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch-size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_annotations() -> None: + loaded_config = { + "optimizer": { + "scheduler": {"schedule-type": "constant"}, + }, + } + + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + del overrided_config + + overrided_config = tyro.cli( + Dict[str, Dict], + default=loaded_config, + args=[ + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + del overrided_config + + overrided_config = tyro.cli( + Dict[str, Dict[str, Dict]], + default=loaded_config, + args=[ + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + del overrided_config + + +def test_functional_typeddict(): + """Source: https://github.com/brentyi/tyro/issues/87""" + NerfMLPHiddenLayers_0 = TypedDict( + "NerfMLPHiddenLayers_0", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + ) + NerfMLPHiddenLayers_1 = TypedDict( + "NerfMLPHiddenLayers_1", + { + "hidden_layers.0": NotRequired[int], + "hidden_layers.1": NotRequired[int], + "hidden_layers.2": NotRequired[int], + "hidden_layers.3": NotRequired[int], + "hidden_layers.4": NotRequired[int], + "hidden_layers.5": NotRequired[int], + "hidden_layers.6": NotRequired[int], + "hidden_layers.7": NotRequired[int], + }, + ) + NerfMLPHiddenLayers_2 = TypedDict( + "NerfMLPHiddenLayers_2", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + total=False, + ) + with pytest.raises(SystemExit): + tyro.cli(NerfMLPHiddenLayers_0, args=["--hidden_layers.0", "3"]) + assert tyro.cli(NerfMLPHiddenLayers_1, args=["--hidden_layers.0", "3"]) == { + "hidden_layers.0": 3 + } + assert tyro.cli(NerfMLPHiddenLayers_2, args=["--hidden_layers.0", "3"]) == { + "hidden_layers.0": 3 + } + + +def test_functional_typeddict_with_default(): + """Source: https://github.com/brentyi/tyro/issues/87""" + NerfMLPHiddenLayers_0 = TypedDict( + "NerfMLPHiddenLayers_0", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + ) + NerfMLPHiddenLayers_1 = TypedDict( + "NerfMLPHiddenLayers_1", + { + "hidden_layers.0": NotRequired[int], + "hidden_layers.1": NotRequired[int], + "hidden_layers.2": NotRequired[int], + "hidden_layers.3": NotRequired[int], + "hidden_layers.4": NotRequired[int], + "hidden_layers.5": NotRequired[int], + "hidden_layers.6": NotRequired[int], + "hidden_layers.7": NotRequired[int], + }, + ) + NerfMLPHiddenLayers_2 = TypedDict( + "NerfMLPHiddenLayers_2", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + total=False, + ) + with pytest.raises(SystemExit): + tyro.cli(NerfMLPHiddenLayers_0, args=["--hidden_layers.0", "3"], default={}) + assert tyro.cli( + NerfMLPHiddenLayers_1, args=["--hidden_layers.0", "3"], default={} + ) == {"hidden_layers.0": 3} + assert tyro.cli( + NerfMLPHiddenLayers_2, args=["--hidden_layers.0", "3"], default={} + ) == {"hidden_layers.0": 3} diff --git a/tests/test_docstrings.py b/tests/test_docstrings.py deleted file mode 100644 index a5a1f6cd7..000000000 --- a/tests/test_docstrings.py +++ /dev/null @@ -1,57 +0,0 @@ -import contextlib -import dataclasses -import io - -import pytest - -import dcargs - - -def test_helptext(): - @dataclasses.dataclass - class Helptext: - x: int # Documentation 1 - - # Documentation 2 - y: int - - z: int = 3 - """Documentation 3""" - - f = io.StringIO() - with pytest.raises(SystemExit): - with contextlib.redirect_stdout(f): - dcargs.parse(Helptext, args=["--help"]) - helptext = f.getvalue() - assert "--x INT Documentation 1\n" in helptext - assert "--y INT Documentation 2\n" in helptext - assert "--z INT Documentation 3 (default: 3)\n" in helptext - - -def test_multiline_helptext(): - @dataclasses.dataclass - class HelptextMultiline: - x: int # Documentation 1 - - # Documentation 2 - # Next line of documentation 2 - y: int - - z: int = 3 - """Documentation 3 - Next line of documentation 3""" - - f = io.StringIO() - with pytest.raises(SystemExit): - with contextlib.redirect_stdout(f): - dcargs.parse(HelptextMultiline, args=["--help"]) - helptext = f.getvalue() - assert " --x INT Documentation 1\n" in helptext - assert ( - " --y INT Documentation 2\n Next line of documentation 2\n" - in helptext - ) - assert ( - " --z INT Documentation 3\n Next line of documentation 3 (default: 3)\n" - in helptext - ) diff --git a/tests/test_dynamic_dataclasses.py b/tests/test_dynamic_dataclasses.py index 639f5b230..b0f30091a 100644 --- a/tests/test_dynamic_dataclasses.py +++ b/tests/test_dynamic_dataclasses.py @@ -2,7 +2,7 @@ import pytest -import dcargs +import tyro def test_dynamic(): @@ -10,5 +10,5 @@ def test_dynamic(): A = make_dataclass("A", [("b", B, field())]) with pytest.raises(SystemExit): - dcargs.parse(A, args=[]) - assert dcargs.parse(A, args=["--b.c", "5"]) == A(b=B(c=5)) + tyro.cli(A, args=[]) + assert tyro.cli(A, args=["--b.c", "5"]) == A(b=B(c=5)) diff --git a/tests/test_errors.py b/tests/test_errors.py new file mode 100644 index 000000000..839947968 --- /dev/null +++ b/tests/test_errors.py @@ -0,0 +1,587 @@ +import contextlib +import dataclasses +import io +from typing import Dict, List, Tuple, TypeVar, Union + +import pytest +from typing_extensions import Literal + +import tyro + + +def test_ambiguous_collection_0() -> None: + @dataclasses.dataclass + class A: + x: Tuple[Tuple[int, ...], ...] + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(A, args=["--x", "0", "1"]) + + +def test_ambiguous_collection_1() -> None: + @dataclasses.dataclass + class A: + x: List[List[int]] + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(A, args=["--x", "0", "1"]) + + +def test_ambiguous_collection_2() -> None: + def main(x: Tuple[List[str], List[str]]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError) as e: + tyro.cli(main, args=["--help"]) + assert "Unsupported type annotation for the field x" in e.value.args[0] + + +def test_ambiguous_collection_3() -> None: + def main(x: List[Union[Tuple[int, int], Tuple[int, int, int]]]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError) as e: + tyro.cli(main, args=["--help"]) + assert "Unsupported type annotation for the field x" in e.value.args[0] + + +def test_ambiguous_collection_4() -> None: + X = List[Union[Tuple[int, int], Tuple[int, int, int]]] + with pytest.raises(tyro.UnsupportedTypeAnnotationError) as e: + tyro.cli(X, args=["--help"]) + assert "Unsupported type annotation for the field" not in e.value.args[0] + + +# Must be global. +@dataclasses.dataclass +class _CycleDataclass: + x: "_CycleDataclass" + + +def test_cycle() -> None: + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(_CycleDataclass, args=[]) + + +def test_uncallable_annotation() -> None: + def main(arg: 5) -> None: # type: ignore + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_nested_annotation() -> None: + @dataclasses.dataclass + class OneIntArg: + x: int + + def main(arg: List[OneIntArg]) -> List[OneIntArg]: # type: ignore + return arg + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + @dataclasses.dataclass + class OneStringArg: + x: str + + def main(arg: List[OneStringArg]) -> List[OneStringArg]: # type: ignore + return arg + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--arg", "0", "1", "2"]) + + @dataclasses.dataclass + class TwoStringArg: + x: str + y: str + + def main2(arg: List[TwoStringArg]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main2, args=[]) + + +def test_missing_annotation_1() -> None: + def main(a, b) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_missing_annotation_2() -> None: + def main(*, a) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_tuple_needs_default() -> None: + def main(arg: tuple) -> None: # type: ignore + pass + + # This formerly raised an error, but now defaults to Tuple[str, ...]. + # with pytest.raises(tyro.UnsupportedTypeAnnotationError): + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_unbound_typevar() -> None: + T = TypeVar("T") + + def main(arg: T) -> None: # type: ignore + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_missing_default_fixed() -> None: + def main(value: tyro.conf.SuppressFixed[tyro.conf.Fixed[int]]) -> int: + return value + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_missing_default_suppressed() -> None: + def main(value: tyro.conf.Suppress[int]) -> int: + return value + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_ambiguous_sequence() -> None: + def main(value: list) -> None: + return None + + # This formerly raised an error, but now defaults to List[str]. + # with pytest.raises(tyro.UnsupportedTypeAnnotationError): + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_similar_arguments_basic() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class Class: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Class, args="--reward.trac".split(" ")) + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + + assert error.count("--reward.track") == 1 + assert error.count("--help") == 1 + + +def test_suppress_console_outputs() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class Class: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Class, args="--reward.trac".split(" "), console_outputs=False) + + error = target.getvalue() + assert error == "" + + +def test_similar_arguments_subcommands() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Union[ClassA, ClassB], args="--reward.trac".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track") == 1 + assert error.count("--help") == 3 + + +def test_similar_arguments_subcommands_multiple() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + trace: int + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Union[ClassA, ClassB], args="--fjdkslaj --reward.trac".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Arguments similar to --reward.trac" in error + assert error.count("--reward.track {True,False}") == 1 + assert error.count("--reward.trace INT") == 1 + assert error.count("--help") == 5 + + +def test_similar_arguments_subcommands_multiple_contains_match() -> None: + @dataclasses.dataclass + class RewardConfigA: + track: bool + trace: int + + @dataclasses.dataclass + class RewardConfigB: ... + + @dataclasses.dataclass + class ClassA: + reward: RewardConfigA + + @dataclasses.dataclass + class ClassB: + reward: RewardConfigB + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( + Union[ClassA, ClassB], + args="class-b --reward.track True --reward.trace 7".split(" "), + ) # type: ignore + + error = target.getvalue() + assert "Unrecognized or misplaced" in error + assert ( + "(applied to " in error and "class-b)" in error + ) # (applied to {root prog} class-b) + assert "Arguments similar to" in error + assert error.count("--reward.track {True,False}") == 2 + assert error.count("--reward.trace INT") == 2 + assert error.count("--help") == 5 # 2 subcommands * 2 arguments + usage hint. + + +def test_similar_arguments_subcommands_multiple_contains_match_alt() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + trace: int + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Union[ClassA, ClassB], args="--track".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track {True,False}") == 1 + assert ( + error.count("--help") == 3 + ) # Should show two possible subcommands + usage hint. + + +def test_similar_arguments_subcommands_overflow_different() -> None: + @dataclasses.dataclass + class RewardConfig: + track0: bool + track1: bool + track2: bool + track3: bool + track4: bool + track5: bool + track6: bool + track7: bool + track8: bool + track9: bool + track10: bool + track11: bool + track12: bool + track13: bool + track14: bool + track15: bool + track16: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Union[ClassA, ClassB], args="--track".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track") == 10 + assert "[...]" not in error + assert error.count("--help") == 21 + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + # --tracked is intentionally between 0.8 ~ 0.9 similarity to track{i} for test + # coverage. + tyro.cli(RewardConfig, args="--tracked".split(" ")) # type: ignore + + # Usage print should be clipped. + error = target.getvalue() + assert "For full helptext, run" in error + + +def test_similar_arguments_subcommands_overflow_same() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + @dataclasses.dataclass + class ClassC: + reward: RewardConfig + + @dataclasses.dataclass + class ClassD: + reward: RewardConfig + + @dataclasses.dataclass + class ClassE: + reward: RewardConfig + + @dataclasses.dataclass + class ClassF: + reward: RewardConfig + + @dataclasses.dataclass + class ClassG: + reward: RewardConfig + + @dataclasses.dataclass + class ClassH: + reward: RewardConfig + + @dataclasses.dataclass + class ClassI: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( # type: ignore + Union[ + ClassA, ClassB, ClassC, ClassD, ClassE, ClassF, ClassG, ClassH, ClassI + ], + args="--track".split(" "), + ) + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track") == 1 + assert "[...]" in error + assert error.count("--help") == 5 + + +def test_similar_arguments_subcommands_overflow_same_startswith_multiple() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + @dataclasses.dataclass + class ClassC: + reward: RewardConfig + + @dataclasses.dataclass + class ClassD: + reward: RewardConfig + + @dataclasses.dataclass + class ClassE: + reward: RewardConfig + rewarde: tyro.conf.Suppress[int] = 10 + + @dataclasses.dataclass + class ClassF: + reward: RewardConfig + + @dataclasses.dataclass + class ClassG: + reward: RewardConfig + + @dataclasses.dataclass + class ClassH: + reward: RewardConfig + + @dataclasses.dataclass + class ClassI: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( # type: ignore + Union[ + ClassA, ClassB, ClassC, ClassD, ClassE, ClassF, ClassG, ClassH, ClassI + ], + args="--track --ffff".split(" "), + ) + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Arguments similar to --track" in error + assert error.count("--rewar") == 1 + assert "rewarde" not in error + assert "[...]" in error + assert error.count("--help") == 5 + + +def test_similar_flag() -> None: + @dataclasses.dataclass + class Args: + flag: bool = False + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( + Args, + args="--lag".split(" "), + ) + + error = target.getvalue() + + # We don't print usage text anymore. + assert error.count("--flag | --no-flag") == 0 + + # Printed in the similar argument list. + assert error.count("--flag, --no-flag") == 1 + + +def test_value_error() -> None: + """https://github.com/brentyi/tyro/issues/86""" + + def main() -> None: + raise ValueError("This shouldn't be caught by tyro") + + with pytest.raises(ValueError): + tyro.cli(main, args=[]) + + +def test_value_error_subcommand() -> None: + """https://github.com/brentyi/tyro/issues/86""" + + def main1() -> None: + raise ValueError("This shouldn't be caught by tyro") + + def main2() -> None: + raise ValueError("This shouldn't be caught by tyro") + + with pytest.raises(ValueError): + tyro.extras.subcommand_cli_from_dict( + {"main1": main1, "main2": main2}, args=["main1"] + ) + + +def test_wrong_annotation() -> None: + @dataclasses.dataclass + class Args: + x: dict = None # type: ignore + + with pytest.warns(UserWarning): + assert tyro.cli(Args, args=[]).x is None + with pytest.warns(UserWarning): + assert tyro.cli(Args, args="--x key value k v".split(" ")).x == { + "key": "value", + "k": "v", + } + + +def test_bad_dict_control() -> None: + assert tyro.cli( + Dict[Literal["a", "b"], Literal["c", "d"]], args="a c b d".split(" ") + ) == {"a": "c", "b": "d"} + + +def test_bad_dict_key() -> None: + with pytest.raises(SystemExit): + tyro.cli(Dict[Literal["a", "b"], Literal["c", "d"]], args="c c d d".split(" ")) + + +def test_bad_dict_val() -> None: + assert tyro.cli( + Dict[Literal["a", "b"], Literal["c", "d"]], args="a c b d".split(" ") + ) == {"a": "c", "b": "d"} + with pytest.raises(SystemExit): + tyro.cli(Dict[Literal["a", "b"], Literal["c", "d"]], args="a a b b".split(" ")) + + +def test_invalid_subcommand_default() -> None: + @dataclasses.dataclass + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + + @dataclasses.dataclass + class C: + c: int + + assert tyro.cli(Union[A, B], args="a --a 5".split(" ")) == A(5) # type: ignore + assert tyro.cli(Union[A, B], args="b --b 5".split(" ")) == B(5) # type: ignore + + with pytest.warns(UserWarning): + assert tyro.cli(Union[A, B], default=C(3), args=[]) == C(3) # type: ignore + + with pytest.warns(UserWarning): + assert tyro.cli(Union[A, B], default=C(3), args="b --b 5".split(" ")) == B(5) # type: ignore diff --git a/tests/test_flax_min_py38.py b/tests/test_flax_min_py38.py new file mode 100644 index 000000000..ab450c77a --- /dev/null +++ b/tests/test_flax_min_py38.py @@ -0,0 +1,68 @@ +"""Tests initializing flax modules directly via tyro.""" + +from typing import cast + +import jax +import pytest +from flax import linen as nn +from helptext_utils import get_helptext_with_checks +from jax import numpy as jnp + +import tyro + + +class Classifier(nn.Module): + layers: int + units: int + output_dim: int + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # type: ignore + for i in range(self.layers - 1): + x = nn.Dense( + self.units, + kernel_init=nn.initializers.kaiming_normal(), + )(x) + x = nn.relu(x) + + x = nn.Dense( + self.output_dim, + kernel_init=nn.initializers.xavier_normal(), + )(x) + x = nn.sigmoid(x) + return x + + +def test_ok(): + network = tyro.cli( + Classifier, + args=[ + "--layers", + "3", + "--units", + "8", + "--output-dim", + "3", + ], + ) + + x = jnp.zeros((10, 4)) + params = network.init(jax.random.PRNGKey(0), x) + assert cast(jax.Array, network.apply(params, x)).shape == (10, 3) + + helptext = get_helptext_with_checks(Classifier) + assert "parent" not in helptext + assert "name" not in helptext + + +def test_missing(): + with pytest.raises(SystemExit): + tyro.cli( + Classifier, + args=[ + "--layers", + "3", + "--units", + "8", + ], + ) diff --git a/tests/test_forward_ref.py b/tests/test_forward_ref.py index d794ec183..ccd29863e 100644 --- a/tests/test_forward_ref.py +++ b/tests/test_forward_ref.py @@ -3,7 +3,7 @@ import pytest -import dcargs +import tyro @dataclasses.dataclass @@ -29,22 +29,20 @@ class C: def test_forward_ref_1(): - - assert dcargs.parse(A1, args=["--x", "1", "b", "--y", "3"]) == A1(x=1, bc=B(y=3)) - assert dcargs.parse(A1, args=["--x", "1", "c", "--z", "3"]) == A1(x=1, bc=C(z=3)) + assert tyro.cli(A1, args=["--x", "1", "bc:b", "--bc.y", "3"]) == A1(x=1, bc=B(y=3)) + assert tyro.cli(A1, args=["--x", "1", "bc:c", "--bc.z", "3"]) == A1(x=1, bc=C(z=3)) with pytest.raises(SystemExit): - dcargs.parse(A1, args=["--x", "1", "b", "--z", "3"]) + tyro.cli(A1, args=["--x", "1", "bc:b", "--bc.z", "3"]) with pytest.raises(SystemExit): - dcargs.parse(A1, args=["--x", "1", "c", "--y", "3"]) + tyro.cli(A1, args=["--x", "1", "bc:c", "--bc.y", "3"]) def test_forward_ref_2(): - - assert dcargs.parse(A2, args=["--x", "1", "b", "--y", "3"]) == A2(x=1, bc=B(y=3)) - assert dcargs.parse(A2, args=["--x", "1", "c", "--z", "3"]) == A2(x=1, bc=C(z=3)) + assert tyro.cli(A2, args=["--x", "1", "bc:b", "--bc.y", "3"]) == A2(x=1, bc=B(y=3)) + assert tyro.cli(A2, args=["--x", "1", "bc:c", "--bc.z", "3"]) == A2(x=1, bc=C(z=3)) with pytest.raises(SystemExit): - dcargs.parse(A2, args=["--x", "1", "b", "--z", "3"]) + tyro.cli(A2, args=["--x", "1", "bc:b", "--bc.z", "3"]) with pytest.raises(SystemExit): - dcargs.parse(A2, args=["--x", "1", "c", "--y", "3"]) + tyro.cli(A2, args=["--x", "1", "bc:c", "--bc.y", "3"]) diff --git a/tests/test_functools.py b/tests/test_functools.py new file mode 100644 index 000000000..f1a0aaf06 --- /dev/null +++ b/tests/test_functools.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import dataclasses +import functools + +import pytest +from helptext_utils import get_helptext_with_checks + +import tyro + + +def test_partial_func() -> None: + def main(a: int, b: str) -> str: + return b * a + + assert tyro.cli(functools.partial(main, a=3), args=["--b", "hi"]) == "hihihi" + + +def test_partial_class() -> None: + class Main: + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + assert tyro.cli(functools.partial(Main, a=3), args=["--b", "hi"]).inner == "hihihi" + + +def test_partial_helptext_func() -> None: + def main(a: int, b: str) -> str: + """Hello!""" + return b * a + + helptext = get_helptext_with_checks(functools.partial(main, b="hello world")) + assert "partial" not in helptext + assert "Hello!" in helptext + assert "hello world" in helptext + + +def test_partial_helptext_class() -> None: + class Main: + """Hello!""" + + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + helptext = get_helptext_with_checks(functools.partial(Main, b="3")) + assert "partial" not in helptext + assert "Hello!" in helptext + + +def test_wraps_func() -> None: + def main(a: int, b: str) -> str: + return b * a + + @functools.wraps(main) + def wrapper(*args, **kwargs) -> int: + return kwargs["a"] + + assert tyro.cli(wrapper, args=["--a", "3", "--b", "hi"]) == 3 + with pytest.raises(SystemExit): + tyro.cli(wrapper, args=["--a", "3"]) + + +def test_wraps_partial_func_helptext() -> None: + def main(a: int, b: str) -> str: + """Hello! + + Args: + a: Argument. + b: Argument. + """ + return b * a + + @functools.wraps(main) + def wrapper(*args, **kwargs) -> int: + return kwargs["a"] + + assert tyro.cli(functools.partial(wrapper, a=3), args=["--b", "hi"]) == 3 + + helptext = get_helptext_with_checks(functools.partial(wrapper, b="3")) + assert "wraps" not in helptext + assert "Hello!" in helptext + assert "Argument." in helptext + + +def test_wraps_partial_class_helptext() -> None: + class Main: + """Hello!""" + + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + @functools.wraps(Main) + def wrapper(*args, **kwargs) -> int: + return kwargs["a"] + + assert tyro.cli(functools.partial(wrapper, a=3), args=["--b", "hi"]) == 3 + + helptext = get_helptext_with_checks(functools.partial(wrapper, b="3")) + assert "wraps" not in helptext + assert "Hello!" in helptext + + +@dataclasses.dataclass +class WrappedDataclass: + """Hello!""" + + a: int + b: str + """Second field.""" + + +def test_wraps_partial_dataclass() -> None: + @functools.wraps(WrappedDataclass) + def wrapper(*args, **kwargs) -> str: + return kwargs["a"] * kwargs["b"] + + assert tyro.cli(functools.partial(wrapper, a=3), args=["--b", "hi"]) == "hihihi" + assert ( + tyro.cli( + functools.partial(functools.partial(wrapper, a=3), b="hello"), + args=["--b", "hi"], + ) + == "hihihi" + ) + assert ( + tyro.cli(functools.partial(functools.partial(wrapper, a=3), b="hello"), args=[]) + == "hellohellohello" + ) + + helptext = get_helptext_with_checks(functools.partial(wrapper, b="3")) + assert "wraps" not in helptext + assert "Hello!" in helptext + assert "Second field." in helptext diff --git a/tests/test_generics.py b/tests/test_generics.py deleted file mode 100644 index 5d275dcd6..000000000 --- a/tests/test_generics.py +++ /dev/null @@ -1,114 +0,0 @@ -import dataclasses -from typing import Generic, TypeVar - -import pytest - -import dcargs - -ScalarType = TypeVar("ScalarType") - - -@dataclasses.dataclass -class Point3(Generic[ScalarType]): - x: ScalarType - y: ScalarType - z: ScalarType - frame_id: str - - -def test_simple_generic(): - @dataclasses.dataclass - class SimpleGeneric: - point_continuous: Point3[float] - point_discrete: Point3[int] - - assert ( - dcargs.parse( - SimpleGeneric, - args=[ - "--point-continuous.x", - "1.2", - "--point-continuous.y", - "2.2", - "--point-continuous.z", - "3.2", - "--point-continuous.frame-id", - "world", - "--point-discrete.x", - "1", - "--point-discrete.y", - "2", - "--point-discrete.z", - "3", - "--point-discrete.frame-id", - "world", - ], - ) - == SimpleGeneric(Point3(1.2, 2.2, 3.2, "world"), Point3(1, 2, 3, "world")) - ) - - with pytest.raises(SystemExit): - # Accidentally pass in floats instead of ints for discrete - dcargs.parse( - SimpleGeneric, - args=[ - "--point-continuous.x", - "1.2", - "--point-continuous.y", - "2.2", - "--point-continuous.z", - "3.2", - "--point-continuous.frame-id", - "world", - "--point-discrete.x", - "1.5", - "--point-discrete.y", - "2.5", - "--point-discrete.z", - "3.5", - "--point-discrete.frame-id", - "world", - ], - ) - - -def test_multilevel_generic(): - @dataclasses.dataclass - class Triangle(Generic[ScalarType]): - a: Point3[ScalarType] - b: Point3[ScalarType] - c: Point3[ScalarType] - - dcargs.parse( - Triangle[float], - args=[ - "--a.x", - "1.0", - "--a.y", - "1.2", - "--a.z", - "1.3", - "--a.frame-id", - "world", - "--b.x", - "1.0", - "--b.y", - "1.2", - "--b.z", - "1.3", - "--b.frame-id", - "world", - "--c.x", - "1.0", - "--c.y", - "1.2", - "--c.z", - "1.3", - "--c.frame-id", - "world", - ], - ) == Triangle( - Point3(1.0, 1.2, 1.3, "world"), - Point3(1.0, 1.2, 1.3, "world"), - Point3(1.0, 1.2, 1.3, "world"), - ) diff --git a/tests/test_generics_and_serialization.py b/tests/test_generics_and_serialization.py new file mode 100644 index 000000000..467a73424 --- /dev/null +++ b/tests/test_generics_and_serialization.py @@ -0,0 +1,460 @@ +import contextlib +import dataclasses +import enum +import io +from typing import Generic, List, NewType, Tuple, Type, TypeVar, Union + +import pytest +import yaml +from typing_extensions import Annotated + +import tyro + +T = TypeVar("T") + + +def _check_serialization_identity(cls: Type[T], instance: T) -> None: + assert tyro.extras.from_yaml(cls, tyro.extras.to_yaml(instance)) == instance + + +ScalarType = TypeVar("ScalarType") + + +def test_tuple_generic_variable() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + assert tyro.cli( + TupleGenericVariable[int], args=["--xyz", "1", "2", "3"] + ) == TupleGenericVariable((1, 2, 3)) + + +def test_tuple_generic_variable_newtype() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + SpecialInt = NewType("SpecialInt", int) + assert tyro.cli( + TupleGenericVariable[SpecialInt], args=["--xyz", "1", "2", "3"] + ) == TupleGenericVariable((SpecialInt(1), SpecialInt(2), SpecialInt(3))) + + +def test_tuple_generic_variable_newtype_container() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + SpecialInt = NewType("SpecialInt", Tuple[int]) + assert tyro.cli( + TupleGenericVariable[SpecialInt], args=["--xyz", "1", "2", "3"] + ) == TupleGenericVariable((SpecialInt((1,)), SpecialInt((2,)), SpecialInt((3,)))) + + +def test_tuple_generic_variable_more_newtype() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + SpecialInt = NewType("SpecialInt", int) + SpecialTuple = NewType("SpecialTuple", TupleGenericVariable[SpecialInt]) + assert tyro.cli(SpecialTuple, args=["--xyz", "1", "2", "3"]) == SpecialTuple( + TupleGenericVariable((SpecialInt(1), SpecialInt(2), SpecialInt(3))) + ) + + +def test_tuple_generic_helptext() -> None: + @dataclasses.dataclass + class TupleGenericVariableHelptext(Generic[ScalarType]): + """Helptext!""" + + xyz: Tuple[ScalarType, ...] + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(TupleGenericVariableHelptext[int], args=["--help"]) + helptext = f.getvalue() + assert "Helptext!" in helptext + + +def test_tuple_generic_no_helptext() -> None: + @dataclasses.dataclass + class TupleGenericVariableNoHelptext(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(TupleGenericVariableNoHelptext[int], args=["--help"]) + helptext = f.getvalue() + assert "Helptext!" not in helptext + + # Check that we don't accidentally grab docstrings from the generic alias! + assert "The central part of internal API" not in helptext + + +def test_tuple_generic_fixed() -> None: + @dataclasses.dataclass + class TupleGenericFixed(Generic[ScalarType]): + xyz: Tuple[ScalarType, ScalarType, ScalarType] + + assert tyro.cli( + TupleGenericFixed[int], args=["--xyz", "1", "2", "3"] + ) == TupleGenericFixed((1, 2, 3)) + + +class CoordinateFrame(enum.Enum): + WORLD = enum.auto() + CAMERA = enum.auto() + + +@dataclasses.dataclass +class Point3(Generic[ScalarType]): + x: ScalarType + y: ScalarType + z: ScalarType + frame: CoordinateFrame + + +def test_simple_generic() -> None: + @dataclasses.dataclass + class SimpleGeneric: + point_continuous: Point3[float] + point_discrete: Point3[int] + + parsed_instance = tyro.cli( + SimpleGeneric, + args=[ + "--point-continuous.x", + "1.2", + "--point-continuous.y", + "2.2", + "--point-continuous.z", + "3.2", + "--point-continuous.frame", + "WORLD", + "--point-discrete.x", + "1", + "--point-discrete.y", + "2", + "--point-discrete.z", + "3", + "--point-discrete.frame", + "WORLD", + ], + ) + assert parsed_instance == SimpleGeneric( + Point3(1.2, 2.2, 3.2, CoordinateFrame.WORLD), + Point3(1, 2, 3, CoordinateFrame.WORLD), + ) + _check_serialization_identity(SimpleGeneric, parsed_instance) + + with pytest.raises(SystemExit): + # Accidentally pass in floats instead of ints for discrete + tyro.cli( + SimpleGeneric, + args=[ + "--point-continuous.x", + "1.2", + "--point-continuous.y", + "2.2", + "--point-continuous.z", + "3.2", + "--point-continuous.frame", + "WORLD", + "--point-discrete.x", + "1.5", + "--point-discrete.y", + "2.5", + "--point-discrete.z", + "3.5", + "--point-discrete.frame", + "WORLD", + ], + ) + + +def test_multilevel_generic() -> None: + @dataclasses.dataclass + class Triangle(Generic[ScalarType]): + a: Point3[ScalarType] + b: Point3[ScalarType] + c: Point3[ScalarType] + + parsed_instance = tyro.cli( + Triangle[float], + args=[ + "--a.x", + "1.0", + "--a.y", + "1.2", + "--a.z", + "1.3", + "--a.frame", + "WORLD", + "--b.x", + "1.0", + "--b.y", + "1.2", + "--b.z", + "1.3", + "--b.frame", + "WORLD", + "--c.x", + "1.0", + "--c.y", + "1.2", + "--c.z", + "1.3", + "--c.frame", + "WORLD", + ], + ) + assert parsed_instance == Triangle( + Point3(1.0, 1.2, 1.3, CoordinateFrame.WORLD), + Point3(1.0, 1.2, 1.3, CoordinateFrame.WORLD), + Point3(1.0, 1.2, 1.3, CoordinateFrame.WORLD), + ) + _check_serialization_identity(Triangle[float], parsed_instance) + + +def test_multilevel_generic_no_helptext() -> None: + @dataclasses.dataclass + class LineSegment(Generic[ScalarType]): + a: Point3[ScalarType] + b: Point3[ScalarType] + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(LineSegment[int], args=["--help"]) + helptext = f.getvalue() + + # Check that we don't accidentally grab docstrings from the generic alias! + assert "The central part of internal API" not in helptext + + +def test_generic_nested_dataclass() -> None: + @dataclasses.dataclass + class Child: + a: int + b: int + + T = TypeVar("T") + + @dataclasses.dataclass + class DataclassGeneric(Generic[T]): + child: T + + parsed_instance = tyro.cli( + DataclassGeneric[Child], args=["--child.a", "5", "--child.b", "7"] + ) + assert parsed_instance == DataclassGeneric(Child(5, 7)) + + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity(DataclassGeneric[Child], parsed_instance) + + +def test_generic_nested_dataclass_helptext() -> None: + @dataclasses.dataclass + class Child: + a: int + b: int + + T = TypeVar("T") + + @dataclasses.dataclass + class DataclassGeneric(Generic[T]): + child: T + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(DataclassGeneric[Child], args=["--help"]) + helptext = f.getvalue() + + # Check that we don't accidentally grab docstrings from the generic alias! + assert "The central part of internal API" not in helptext + + +def test_generic_subparsers() -> None: + @dataclasses.dataclass + class CommandOne: + a: int + + @dataclasses.dataclass + class CommandTwo: + b: int + + T1 = TypeVar("T1") + T2 = TypeVar("T2") + + @dataclasses.dataclass + class Subparser(Generic[T1, T2]): + command: Union[T1, T2] + + parsed_instance = tyro.cli( + Subparser[CommandOne, CommandTwo], + args="command:command-one --command.a 5".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(CommandOne(5)) # type: ignore + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[CommandOne, CommandTwo], parsed_instance + ) + + parsed_instance = tyro.cli( + Subparser[CommandOne, CommandTwo], + args="command:command-two --command.b 7".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(CommandTwo(7)) # type: ignore + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[CommandOne, CommandTwo], parsed_instance + ) + + +def test_generic_subparsers_in_container() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class Command(Generic[T]): + a: List[T] + + T1 = TypeVar("T1") + T2 = TypeVar("T2") + + @dataclasses.dataclass + class Subparser(Generic[T1, T2]): + command: Union[T1, T2] + + parsed_instance = tyro.cli( + Subparser[Command[int], Command[float]], + args="command:command-int --command.a 5 3".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(Command([5, 3])) and isinstance( # type: ignore + parsed_instance.command.a[0], int + ) + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[Command[int], Command[float]], parsed_instance + ) + + parsed_instance = tyro.cli( + Subparser[Command[int], Command[float]], + args="command:command-float --command.a 7 2".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(Command([7.0, 2.0])) and isinstance( # type: ignore + parsed_instance.command.a[0], float + ) + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[Command[int], Command[float]], parsed_instance + ) + + +def test_serialize_missing() -> None: + @dataclasses.dataclass + class TupleGenericVariableMissing(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + x = TupleGenericVariableMissing[int](xyz=(tyro.MISSING, tyro.MISSING)) + assert tyro.extras.to_yaml(x).count("!missing") == 2 + _check_serialization_identity(TupleGenericVariableMissing[int], x) + assert ( + tyro.extras.from_yaml( + TupleGenericVariableMissing[int], tyro.extras.to_yaml(x) + ).xyz[0] + is tyro.MISSING + ) + + +def test_generic_inherited_type_narrowing() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class ActualParentClass(Generic[T]): + x: T # Documentation 1 + + # Documentation 2 + y: T + + z: T = 3 # type: ignore + """Documentation 3""" + + @dataclasses.dataclass + class ChildClass(ActualParentClass[int]): + a: int = 7 + + def main(x: ActualParentClass[int] = ChildClass(5, 5)) -> ActualParentClass: + return x + + assert tyro.cli(main, args="--x.x 3".split(" ")) == ChildClass(3, 5, 3) + + +def test_pculbertson() -> None: + # https://github.com/brentyi/tyro/issues/7 + from typing import Union + + @dataclasses.dataclass(frozen=True) + class TypeA: + data: int + + @dataclasses.dataclass + class TypeB: + data: int + + @dataclasses.dataclass + class Wrapper: + subclass: Union[TypeA, TypeB] = TypeA(1) + + wrapper1 = Wrapper() # Create Wrapper object. + assert wrapper1 == tyro.extras.from_yaml(Wrapper, tyro.extras.to_yaml(wrapper1)) + + +def test_annotated() -> None: + # https://github.com/brentyi/tyro/issues/7 + + @dataclasses.dataclass(frozen=True) + class TypeA: + data: int + + @dataclasses.dataclass + class Wrapper: + subclass: Annotated[TypeA, int] = TypeA(1) + + wrapper1 = Wrapper() # Create Wrapper object. + assert wrapper1 == tyro.extras.from_yaml(Wrapper, tyro.extras.to_yaml(wrapper1)) + + +@dataclasses.dataclass +class TypeA: + data: int + + +@dataclasses.dataclass +class TypeASubclass(TypeA): + pass + + +@dataclasses.dataclass +class Wrapper: + subclass: TypeA + + +def test_superclass() -> None: + # https://github.com/brentyi/tyro/issues/7 + + wrapper1 = Wrapper(TypeASubclass(3)) # Create Wrapper object. + assert wrapper1 == tyro.extras.from_yaml(Wrapper, tyro.extras.to_yaml(wrapper1)) diff --git a/tests/test_helptext.py b/tests/test_helptext.py new file mode 100644 index 000000000..75c79d178 --- /dev/null +++ b/tests/test_helptext.py @@ -0,0 +1,974 @@ +import dataclasses +import enum +import json +import os +import pathlib +from collections.abc import Callable +from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union, cast + +from helptext_utils import get_helptext_with_checks +from torch import nn +from typing_extensions import Annotated, Literal, NotRequired, TypedDict + +import tyro + + +def test_helptext() -> None: + @dataclasses.dataclass + class Helptext: + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: Annotated[int, "ignored"] + + z: int = 3 + """Documentation 3""" + + helptext = get_helptext_with_checks(Helptext) + assert cast(str, helptext) in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_sphinx_autodoc_style() -> None: + @dataclasses.dataclass + class Helptext: + """This docstring should be printed as a description.""" + + x: int #: Documentation 1 + + #:Documentation 2 + y: Annotated[int, "ignored"] + z: int = 3 + + helptext = get_helptext_with_checks(Helptext) + assert cast(str, helptext) in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert ": Documentation 1" not in helptext + assert "Documentation 2 (required)" in helptext + assert ":Documentation 2" not in helptext + + # :Documentation 2 should not be applied to `z`. + assert helptext.count("Documentation 2") == 1 + + +def test_helptext_from_class_docstring() -> None: + @dataclasses.dataclass + class Helptext2: + """This docstring should be printed as a description. + + Attributes: + x: Documentation 1 + y: Documentation 2 + z: Documentation 3 + """ + + x: int + y: Annotated[int, "ignored"] + z: int = 3 + + helptext = get_helptext_with_checks(Helptext2) + assert "This docstring should be printed as a description" in helptext + assert "Attributes" not in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_from_class_docstring_args() -> None: + @dataclasses.dataclass + class Helptext3: + """This docstring should be printed as a description. + + Args: + x: Documentation 1 + y: Documentation 2 + z: Documentation 3 + """ + + x: int + y: Annotated[int, "ignored"] + z: int = 3 + + helptext = get_helptext_with_checks(Helptext3) + assert "This docstring should be printed as a description" in helptext + assert "Args" not in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_inherited() -> None: + class UnrelatedParentClass: + pass + + @dataclasses.dataclass + class ActualParentClass: + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + # fmt: off + + z: int = 3 + def some_method(self) -> None: # noqa + """Coverage stress test.""" + # fmt: on + + @dataclasses.dataclass + class ChildClass(UnrelatedParentClass, ActualParentClass): + pass + + helptext = get_helptext_with_checks(ChildClass) + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + + +def test_helptext_inherited_default_override() -> None: + @dataclasses.dataclass + class ParentClass: + """This docstring should __not__ be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + # fmt: off + + z: int = 3 + def some_method(self) -> None: # noqa + """Coverage stress test.""" + # fmt: on + + @dataclasses.dataclass + class ChildClass(ParentClass): + """This docstring should be printed as a description.""" + + def main(x: ParentClass = ChildClass(x=5, y=5)) -> Any: + """Main function.""" + return x + + helptext = get_helptext_with_checks(main) + assert "Main function." in helptext + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "__not__" not in helptext + assert helptext.count("should be printed") == 1 + + +def test_helptext_nested() -> None: + """For nested classes, we should pull helptext from the outermost docstring if + possible. The class docstring can be used as a fallback.""" + + class Inner: + def __init__(self, a: int): + """Something + + Args: + a (int): Hello world! + """ + + def main_with_docstring(a: Inner) -> None: + """main_with_docstring. + + Args: + a: Documented in function. + """ + + def main_no_docstring(a: Inner) -> None: + """main_no_docstring.""" + + helptext = get_helptext_with_checks(main_with_docstring) + assert "Documented in function" in helptext and str(Inner.__doc__) not in helptext + assert "main_with_docstring." in helptext + assert "Args:" not in helptext + assert "Args:" not in helptext + assert "Hello world!" in helptext + + helptext = get_helptext_with_checks(main_no_docstring) + assert "Something" in helptext + assert "main_no_docstring." in helptext + assert "Args:" not in helptext + assert "Hello world!" in helptext + + +def test_helptext_defaults() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class HelptextWithVariousDefaults: + x: pathlib.Path = pathlib.Path("/some/path/to/a/file") + y: Color = Color.RED + z: str = "%" + + helptext = get_helptext_with_checks(HelptextWithVariousDefaults) + assert "show this help message and exit" in helptext + assert "--x PATH" in helptext + assert "(default: /some/path/to/a/file)" in helptext + assert "--y {RED,GREEN,BLUE}" in helptext + assert "(default: RED)" in helptext + assert "--z STR" in helptext + assert "(default: %)" in helptext + + +def test_multiline_helptext() -> None: + @dataclasses.dataclass + class HelptextMultiline: + x: int # Documentation 1 + + # This comment should be ignored! + + # Documentation 2 + # Next line of documentation 2 + y: int + + z: int = 3 + """Documentation 3 + Next line of documentation 3""" + + helptext = get_helptext_with_checks(HelptextMultiline) + assert "Documentation 1 (required)" in helptext + assert "Documentation 2" in helptext + assert "documentation 2" in helptext + assert "Documentation 3" in helptext + assert "documentation 3" in helptext + + +def test_grouped_helptext() -> None: + @dataclasses.dataclass + class HelptextGrouped: + x: int # Documentation 1 + # Description of both y and z. + y: int + z: int = 3 + + helptext = get_helptext_with_checks(HelptextGrouped) + assert "Documentation 1 (required)" in helptext + assert "Description of both y and z. (required)" in helptext + assert "Description of both y and z. (default: 3)" in helptext + + +def test_none_default_value_helptext() -> None: + @dataclasses.dataclass + class Config: + x: Optional[int] = None + """An optional variable.""" + + helptext = get_helptext_with_checks(Config) + assert "--x {None}|INT" in helptext + assert "An optional variable. (default: None)" in helptext + + +def test_helptext_hard_bool() -> None: + @dataclasses.dataclass + class HelptextHardString: + # fmt: off + x: bool = ( + False + ) + """Helptext. 2% milk.""" + # fmt: on + + # Note that the percent symbol needs some extra handling in argparse. + # https://stackoverflow.com/questions/21168120/python-argparse-errors-with-in-help-string + + helptext = get_helptext_with_checks(HelptextHardString) + assert "--x" in helptext + assert "2% milk." in helptext + + +def test_helptext_with_inheritance() -> None: + @dataclasses.dataclass + class Parent: + # fmt: off + x: str = ( + "This docstring may be tougher to parse!" + ) + """Helptext.""" + # fmt: on + + @dataclasses.dataclass + class Child(Parent): + pass + + helptext = get_helptext_with_checks(Child) + assert "--x STR" in helptext + assert "Helptext." in helptext + assert "(default: 'This docstring" in helptext + + +def test_helptext_with_inheritance_overriden() -> None: + @dataclasses.dataclass + class Parent2: + # fmt: off + x: str = ( + "This docstring may be tougher to parse!" + ) + """Helptext.""" + # fmt: on + + @dataclasses.dataclass + class Child2(Parent2): + # fmt: off + x: str = ( + "This docstring may be tougher to parse?" + ) + """Helptext!""" + # fmt: on + + helptext = get_helptext_with_checks(Child2) + assert "--x STR" in helptext + assert "Helptext! (default: 'This" in helptext + + +def test_tuple_helptext() -> None: + @dataclasses.dataclass + class TupleHelptext: + x: Tuple[int, str, float] + + helptext = get_helptext_with_checks(TupleHelptext) + assert "--x INT STR FLOAT" in helptext + + +def test_tuple_helptext_defaults() -> None: + @dataclasses.dataclass + class TupleHelptextDefaults: + x: Tuple[int, str, str] = (5, "hello world", "hello") + + helptext = get_helptext_with_checks(TupleHelptextDefaults) + assert "--x INT STR STR" in helptext + assert "(default: 5 'hello world' hello)" in helptext + + +def test_generic_helptext() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class GenericTupleHelptext(Generic[T]): + x: T + + helptext = get_helptext_with_checks(GenericTupleHelptext[int]) + assert "--x INT" in helptext + + +def test_generic_tuple_helptext() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class GenericTupleHelptext(Generic[T]): + x: Tuple[T, T, T] + + helptext = get_helptext_with_checks(GenericTupleHelptext[int]) + assert "--x INT INT INT" in helptext + + +def test_generic_list_helptext() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class GenericTupleHelptext(Generic[T]): + x: List[T] + + helptext = get_helptext_with_checks(GenericTupleHelptext[int]) + assert "--x [INT [INT ...]]" in helptext + + +def test_literal_helptext() -> None: + @dataclasses.dataclass + class LiteralHelptext: + x: Literal[1, 2, 3] + """A number.""" + + helptext = get_helptext_with_checks(LiteralHelptext) + assert "--x {1,2,3}" in helptext + assert "A number. (required)" in helptext + + +def test_optional_literal_helptext() -> None: + @dataclasses.dataclass + class OptionalLiteralHelptext: + x: Optional[Literal[1, 2, 3]] = None + """A number.""" + + helptext = get_helptext_with_checks(OptionalLiteralHelptext) + assert "--x {None,1,2,3}" in helptext + assert "A number. (default: None)" in helptext + + +def test_multiple_subparsers_helptext() -> None: + @dataclasses.dataclass + class Subcommand1: + """2% milk.""" # % symbol is prone to bugs in argparse. + + x: int = 0 + + @dataclasses.dataclass + class Subcommand2: + y: int = 1 + + @dataclasses.dataclass + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: Union[Subcommand1, Subcommand2, Subcommand3] + # Field b description. + b: Union[Subcommand1, Subcommand2, Subcommand3] + # Field c description. + c: Union[Subcommand1, Subcommand2, Subcommand3] = dataclasses.field( + default_factory=Subcommand3 + ) + + d: bool = False + + helptext = get_helptext_with_checks(MultipleSubparsers) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + # Not enough args for usage shortening to kick in. + assert "[OPTIONS]" not in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, args=["a:subcommand1", "b:subcommand1", "--help"] + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand3)" in helptext + + # Not enough args for usage shortening to kick in. + assert "[OPTIONS]" not in helptext + assert "[B:SUBCOMMAND1 OPTIONS]" not in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + +def test_multiple_subparsers_helptext_shortened_usage() -> None: + @dataclasses.dataclass + class Subcommand1: + """2% milk.""" # % symbol is prone to bugs in argparse. + + a: int = 0 + b: int = 0 + c: int = 0 + d: int = 0 + e: int = 0 + + @dataclasses.dataclass + class Subcommand2: + a: int = 0 + b: int = 0 + c: int = 0 + d: int = 0 + e: int = 0 + + @dataclasses.dataclass + class Subcommand3: + a: int = 0 + b: int = 0 + c: int = 0 + d: int = 0 + e: int = 0 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: Union[Subcommand1, Subcommand2, Subcommand3] + # Field b description. + b: Union[Subcommand1, Subcommand2, Subcommand3] + # Field c description. + c: Union[Subcommand1, Subcommand2, Subcommand3] = dataclasses.field( + default_factory=Subcommand3 + ) + + d: bool = False + f: bool = False + g: bool = False + h: bool = False + i: bool = False + + helptext = get_helptext_with_checks(MultipleSubparsers) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + assert "[OPTIONS]" in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, args=["a:subcommand1", "b:subcommand1", "--help"] + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand3)" in helptext + + assert "[OPTIONS]" not in helptext + assert "[B:SUBCOMMAND1 OPTIONS]" in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + +def test_optional_helptext() -> None: + @dataclasses.dataclass + class OptionalHelptext: + """This docstring should be printed as a description. 2% milk.""" + + x: Optional[int] # Documentation 1 + + # Documentation 2 + y: List[Optional[int]] + + z: Optional[int] = 3 + """Documentation 3""" + + helptext = get_helptext_with_checks(OptionalHelptext) + assert cast(str, cast(str, OptionalHelptext.__doc__)[:20]) in helptext + assert "2% milk" in helptext + assert "--x {None}|INT" in helptext + assert "--y [{None}|INT [{None}|INT ...]]" in helptext + assert "[--z {None}|INT]" in helptext + + +def test_metavar_0() -> None: + def main(x: Union[Literal[0, 1, 2, 3], Tuple[int, int]]) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3}|{INT INT}" in helptext + + +def test_metavar_1() -> None: + def main( + x: Union[ + Literal[0, 1, 2, 3], + Literal["hey,there", "hello"], + List[int], + ], + ) -> None: + pass + + # The comma formatting is unfortunate, but matches argparse's default behavior. + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3,hey,there,hello}|{[INT [INT ...]]}" in helptext + + +def test_metavar_2() -> None: + def main( + x: Tuple[ + Literal[0, 1, 2, 3], + Union[int, str], + ], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3} INT|STR" in helptext + + +def test_metavar_3() -> None: + def main( + x: Union[ + Literal[0, 1, 2, 3], + Union[Tuple[int, int], Tuple[str]], + ], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3}|{INT INT}|STR" in helptext + + +def test_metavar_4() -> None: + def main( + x: Union[ + Literal[0, 1, 2, 3], + Union[Tuple[int, int], Tuple[str, str, str]], + Literal[True], + ], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3}|{INT INT}|{STR STR STR}|{True}" in helptext + + +def test_metavar_5() -> None: + def main( + x: List[Union[Tuple[int, int], Tuple[str, str]]] = [(1, 1), (2, 2)], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "[--x [{INT INT}|{STR STR} [{INT INT}|{STR STR} ...]]]" in helptext + + +def test_metavar_6() -> None: + def main(x: Dict[Union[Tuple[int, int], Tuple[str, str]], Tuple[int, int]]) -> dict: + return x + + helptext = get_helptext_with_checks(main) + assert ( + "--x [{INT INT}|{STR STR} INT INT [{INT INT}|{STR STR} INT INT ...]]" + in helptext + ) + + +def test_comment_in_subclass_list() -> None: + @dataclasses.dataclass + class Something( + # This text should not show up in the helptext anywhere. + object, + ): + a: int + + # But this text should! + b: int + + helptext = get_helptext_with_checks(Something) + assert "This text should not" not in helptext + assert "But this text should!" in helptext + + +def test_unparsable() -> None: + class Struct: + a: int = 5 + b: str = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x {fixed}" in helptext + + def main2(x: Callable = nn.ReLU): + pass + + helptext = get_helptext_with_checks(main2) + assert "--x {fixed}" in helptext + assert "(fixed to:" in helptext + assert "torch" in helptext + + +def test_pathlike() -> None: + def main(x: os.PathLike) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x PATH " in helptext + + +def test_nested_bool() -> None: + @dataclasses.dataclass + class Child: + x: bool = False + + def main(child: Child) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--child.x | --child.no-x" in helptext + + +def test_multiple_subparsers_helptext_hyphens() -> None: + @dataclasses.dataclass + class SubcommandOne: + """2% milk.""" # % symbol is prone to bugs in argparse. + + arg_x: int = 0 + arg_flag: bool = False + + @dataclasses.dataclass + class SubcommandTwo: + arg_y: int = 1 + + @dataclasses.dataclass + class SubcommandThree: + arg_z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: Union[SubcommandOne, SubcommandTwo, SubcommandThree] + # Field b description. + b: Union[SubcommandOne, SubcommandTwo, SubcommandThree] + # Field c description. + c: Union[SubcommandOne, SubcommandTwo, SubcommandThree] = dataclasses.field( + default_factory=SubcommandThree + ) + + helptext = get_helptext_with_checks(MultipleSubparsers) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, args=["a:subcommand-one", "b:subcommand-one", "--help"] + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand-three)" in helptext + assert "--b.arg-x" in helptext + assert "--b.no-arg-flag" in helptext + assert "--b.arg-flag" in helptext + + +def test_multiple_subparsers_helptext_underscores() -> None: + @dataclasses.dataclass + class SubcommandOne: + """2% milk.""" # % symbol is prone to bugs in argparse. + + arg_x: int = 0 + arg_flag: bool = False + + @dataclasses.dataclass + class SubcommandTwo: + arg_y: int = 1 + + @dataclasses.dataclass + class SubcommandThree: + arg_z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: Union[SubcommandOne, SubcommandTwo, SubcommandThree] + # Field b description. + b: Union[SubcommandOne, SubcommandTwo, SubcommandThree] + # Field c description. + c: Union[SubcommandOne, SubcommandTwo, SubcommandThree] = dataclasses.field( + default_factory=SubcommandThree + ) + + helptext = get_helptext_with_checks(MultipleSubparsers, use_underscores=True) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, + args=["a:subcommand_one", "b:subcommand_one", "--help"], + use_underscores=True, + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand_three)" in helptext + assert "--b.arg_x" in helptext + assert "--b.no_arg_flag" in helptext + assert "--b.arg_flag" in helptext + + +def test_subparsers_wrapping() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CheckoutCompletion: + """Help message.""" + + y: int + + help = get_helptext_with_checks(Union[A, CheckoutCompletion]) # type: ignore + assert help.count("checkout-completion") == 3 + + +def test_subparsers_wrapping1() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CheckoutCompletio: + """Help message.""" + + y: int + + help = get_helptext_with_checks(Union[A, CheckoutCompletio]) # type: ignore + assert help.count("checkout-completio") == 3 + + +def test_subparsers_wrapping2() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CheckoutCompletionn: + """Help message.""" + + y: int + + help = get_helptext_with_checks(Union[A, CheckoutCompletionn]) # type: ignore + assert help.count("checkout-completionn") == 3 + + +def test_subparsers_wrapping3() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CmdCheckout012: + """Help message.""" + + y: int + + help = get_helptext_with_checks(Union[A, CmdCheckout012]) # type: ignore + assert help.count("cmd-checkout012") == 3 + + +def test_tuple_default() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: Tuple[str, str] = ("hello", "world") + + help = get_helptext_with_checks(A) + assert "STR STR" in help + assert "hello world" in help + assert "('hello', 'world')" not in help + + +def test_argconf_constructor() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: Annotated[ + Tuple[str, str], tyro.conf.arg(constructor=lambda x: ("a", "b")) + ] = ("hello", "world") + + help = get_helptext_with_checks(A) + assert "STR STR" not in help + # Unlike case above, should not be converted to 'hello world'. + assert "('hello', 'world')" in help # JSON special case. + + +def test_argconf_constructor_json_special_case() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: Annotated[ + Tuple[str, str], tyro.conf.arg(constructor=json.loads, metavar="JSON") + ] = ("hello", "world") + y: Annotated[ + Tuple[int, int], tyro.conf.arg(constructor=json.loads, metavar="JSON") + ] = (3, 5) + + help = get_helptext_with_checks(A) + assert "STR STR" not in help + assert "JSON" in help + assert '["hello", "world"]' in help # JSON special case. + assert "[3, 5]" in help, help # JSON special case. + + +def test_optional_group() -> None: + def f( + x: Annotated[ + Tuple[str, str], tyro.conf.arg(constructor=json.loads, metavar="JSON") + ] = ("hello", "world"), + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f, default=3) + assert 'default if used: ["hello", "world"]' in help + assert "default if used: 3" in help + + +def test_append_fixed() -> None: + def f( + x: tyro.conf.UseAppendAction[Tuple[str, str]], + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f) + assert "repeatable" not in help, help + + +def test_append_good() -> None: + def f( + x: tyro.conf.UseAppendAction[Tuple[str, ...]], + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f) + assert "repeatable" in help, help + + +def test_append_with_default() -> None: + def f( + x: tyro.conf.UseAppendAction[Tuple[str, ...]] = ("hello world", "hello"), + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f) + assert "repeatable, appends to: 'hello world' hello" in help, help + + +def test_typeddict_exclude() -> None: + class Special(TypedDict): + x: NotRequired[int] + + help = get_helptext_with_checks(Special) + assert "unset by default" in help, help diff --git a/tests/test_initvar_min_py38.py b/tests/test_initvar_min_py38.py new file mode 100644 index 000000000..721462716 --- /dev/null +++ b/tests/test_initvar_min_py38.py @@ -0,0 +1,27 @@ +import dataclasses + +import tyro + + +def test_dataclass_init_var() -> None: + @dataclasses.dataclass + class DataclassWithInitVar: + i: dataclasses.InitVar[int] + x: str + + def __post_init__(self, i: int) -> None: + self.x += str(i) + + # We can directly pass a dataclass to `tyro.cli()`: + assert ( + tyro.cli( + DataclassWithInitVar, + args=[ + "--i", + "5", + "--x", + "5", + ], + ).x + == "55" + ) diff --git a/tests/test_is_nested_type.py b/tests/test_is_nested_type.py new file mode 100644 index 000000000..a586756f4 --- /dev/null +++ b/tests/test_is_nested_type.py @@ -0,0 +1,52 @@ +import dataclasses +import pathlib +from typing import Any, Dict, List, Tuple + +from tyro._fields import MISSING_NONPROP, is_nested_type + + +def test_is_nested_type_simple(): + assert not is_nested_type(int, MISSING_NONPROP) + assert not is_nested_type(bool, MISSING_NONPROP) + assert not is_nested_type(str, MISSING_NONPROP) + assert not is_nested_type(pathlib.Path, MISSING_NONPROP) + + +def test_is_nested_type_containers(): + assert not is_nested_type(List[int], MISSING_NONPROP) + assert not is_nested_type(List[bool], MISSING_NONPROP) + assert not is_nested_type(List[str], MISSING_NONPROP) + assert not is_nested_type(List[pathlib.Path], MISSING_NONPROP) + + +@dataclasses.dataclass +class Color: + r: int + g: int + b: int + + +def test_is_nested_type_actually_nested(): + assert is_nested_type(Color, Color(255, 0, 0)) + + +def test_is_nested_type_actually_nested_narrowing(): + assert is_nested_type(Any, Color(255, 0, 0)) + assert is_nested_type(object, Color(255, 0, 0)) + assert not is_nested_type(int, Color(255, 0, 0)) + + +def test_is_nested_type_actually_nested_in_container(): + assert is_nested_type(Tuple[Color, Color], MISSING_NONPROP) + assert is_nested_type(Tuple[object, ...], (Color(255, 0, 0),)) + assert is_nested_type(Tuple[Any, ...], (Color(255, 0, 0),)) + assert is_nested_type(tuple, (Color(255, 0, 0),)) + assert not is_nested_type(tuple, (1, 2, 3)) + assert is_nested_type(tuple, (1, Color(255, 0, 0), 3)) + assert is_nested_type(List[Any], [Color(255, 0, 0)]) + + +def test_nested_dict(): + assert is_nested_type(Dict[str, int], {"x": 5}) + assert is_nested_type(dict, {"x": 5}) + assert is_nested_type(Any, {"x": 5}) diff --git a/tests/test_missing.py b/tests/test_missing.py new file mode 100644 index 000000000..3f7f9d428 --- /dev/null +++ b/tests/test_missing.py @@ -0,0 +1,82 @@ +"""Tests for tyro.MISSING.""" + +import contextlib +import dataclasses +import io +from typing import Tuple + +import pytest + +import tyro + + +def test_missing() -> None: + def main(a: int = 5, b: int = tyro.MISSING, c: int = 3) -> Tuple[int, int, int]: + return a, b, c + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(main, args=[]) + message = target.getvalue() + assert "Required options" in message + assert "Argument helptext" in message + assert "--a INT" not in message + assert "--b INT" in message + assert "(required)" in message + assert tyro.cli(main, args=["--b", "7"]) == (5, 7, 3) + + +def test_missing_dataclass() -> None: + @dataclasses.dataclass + class Args2: + a: int = 5 + b: int = tyro.MISSING + c: int = 3 + + with pytest.raises(SystemExit): + tyro.cli(Args2, args=[]) + assert tyro.cli(Args2, args=["--b", "7"]) == Args2(5, 7, 3) + + +def test_missing_default() -> None: + @dataclasses.dataclass + class Args2: + a: int + b: int + c: int + + with pytest.raises(SystemExit): + tyro.cli( + Args2, + args=[], + default=Args2(5, tyro.MISSING, 3), + ) + assert tyro.cli( + Args2, + args=["--b", "7"], + default=Args2(5, tyro.MISSING, 3), + ) == Args2(5, 7, 3) + + +def test_missing_nested_default() -> None: + @dataclasses.dataclass + class Child: + a: int = 5 + b: int = 3 + c: int = 7 + + @dataclasses.dataclass + class Parent: + child: Child + + with pytest.raises(SystemExit): + tyro.cli( + Parent, + args=[], + default=Parent(child=tyro.MISSING), + ) + assert tyro.cli( + Parent, + args=["--child.a", "5", "--child.b", "7", "--child.c", "3"], + default=Parent(child=tyro.MISSING), + ) == Parent(Child(5, 7, 3)) diff --git a/tests/test_missing_optional_packages.py b/tests/test_missing_optional_packages.py new file mode 100644 index 000000000..084e6e7f9 --- /dev/null +++ b/tests/test_missing_optional_packages.py @@ -0,0 +1,33 @@ +import builtins +import dataclasses +from typing import Any + +import pytest + + +# https://stackoverflow.com/questions/60227582/making-a-python-test-think-an-installed-package-is-not-available +@pytest.fixture +def hide_optional_packages(monkeypatch: Any) -> None: + import_orig = builtins.__import__ + + def mocked_import(name, *args, **kwargs): + if name in ("attr", "pydantic", "flax", "omegaconf"): + raise ImportError() + return import_orig(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", mocked_import) + + +def test_missing_optional_packages(hide_optional_packages: Any) -> None: + @dataclasses.dataclass + class Extra: + a: int = 3 + b: str = "5" + + def main(x: int, extra: Extra) -> int: + del extra + return x + + import tyro + + assert tyro.cli(main, args=["--x", "5"]) == 5 diff --git a/tests/test_mixed_unions.py b/tests/test_mixed_unions.py new file mode 100644 index 000000000..d6444f9bd --- /dev/null +++ b/tests/test_mixed_unions.py @@ -0,0 +1,97 @@ +"""Tests for unsupported union types. + +Unions like `int | str` or `SomeDataclassA | SomeDataclassB` are OK (note that the latter +will produce a pair of subcommands); when we write things like `int | SomeDataclassA` +handling gets more complicated but should still be supported! +""" + +import dataclasses +from typing import Any, Dict, List, Tuple, Union + +import pytest + +import tyro + + +def test_subparser_strip_non_nested() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + # Note that we add [int, str] to the annotation here... this should be ignored. + bc: Union[int, str, DefaultHTTPServer, DefaultSMTPServer] = dataclasses.field( + default_factory=lambda: DefaultHTTPServer(5) + ) + + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "5"] + ) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-smtp-server", "--bc.z", "3"] + ) == DefaultSubparser(x=1, bc=DefaultSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "8"] + ) + == tyro.cli( + DefaultSubparser, + args=[], + default=DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)), + ) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_strip_nested() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + # Note that we add [int, str] to the annotation here... this should be ignored. + bc: Union[int, str, DefaultHTTPServer, DefaultSMTPServer] = 5 + + assert ( + tyro.cli(DefaultSubparser, args=["--x", "1", "bc:int", "5"]) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=5) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:str", "five"] + ) == DefaultSubparser(x=1, bc="five") + + +def test_with_fancy_types() -> None: + @dataclasses.dataclass + class Args: + y: int + + def main(x: Union[Tuple[int, ...], List[str], Args, Dict[str, int]]) -> Any: + return x + + assert tyro.cli(main, args="x:tuple-int-ellipsis 1 2 3".split(" ")) == (1, 2, 3) + assert tyro.cli(main, args="x:list-str 1 2 3".split(" ")) == ["1", "2", "3"] + assert tyro.cli(main, args="x:args --x.y 5".split(" ")) == Args(5) + assert tyro.cli(main, args="x:dict-str-int 1 2 3 4".split(" ")) == {"1": 2, "3": 4} diff --git a/tests/test_nested.py b/tests/test_nested.py new file mode 100644 index 000000000..0264057df --- /dev/null +++ b/tests/test_nested.py @@ -0,0 +1,1189 @@ +import dataclasses +from typing import Any, Generic, Mapping, NewType, Optional, Tuple, TypeVar, Union + +import pytest +from frozendict import frozendict # type: ignore +from helptext_utils import get_helptext_with_checks +from typing_extensions import Annotated, Final, Literal + +import tyro + + +def test_nested() -> None: + @dataclasses.dataclass + class B: + y: int + + @dataclasses.dataclass + class Nested: + x: int + b: B + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_annotated() -> None: + @dataclasses.dataclass + class B: + y: int + + @dataclasses.dataclass + class Nested: + x: int + b: Annotated[B, "this should be ignored"] + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_final() -> None: + @dataclasses.dataclass + class B: + y: int + + @dataclasses.dataclass + class Nested: + x: int + b: Final[B] # type: ignore + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_accidental_underscores() -> None: + @dataclasses.dataclass + class B: + arg_name: str + + @dataclasses.dataclass + class Nested: + x: int + child_struct: B + + assert ( + tyro.cli(Nested, args=["--x", "1", "--child-struct.arg-name", "three_five"]) + == tyro.cli(Nested, args=["--x", "1", "--child_struct.arg_name", "three_five"]) + == tyro.cli(Nested, args=["--x", "1", "--child_struct.arg-name", "three_five"]) + == tyro.cli(Nested, args=["--x", "1", "--child_struct.arg_name=three_five"]) + == Nested(x=1, child_struct=B(arg_name="three_five")) + ) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_default() -> None: + @dataclasses.dataclass(frozen=True) + class B: + y: int = 1 + + @dataclasses.dataclass + class Nested: + x: int = 2 + b: B = B() + + assert tyro.cli(Nested, args=[], default=Nested(x=1, b=B(y=2))) == Nested( + x=1, b=B(y=2) + ) + + +def test_nested_default_alternate() -> None: + @dataclasses.dataclass + class B: + y: int = 3 + + @dataclasses.dataclass + class Nested: + x: int + b: B + + assert ( + Nested(x=1, b=B(y=3)) + == tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) + == tyro.cli(Nested, args=[], default=Nested(x=1, b=B(y=3))) + ) + assert tyro.cli(Nested, args=["--x", "1"]) == Nested(x=1, b=B(y=3)) + + +def test_default_nested() -> None: + @dataclasses.dataclass(frozen=True) + class B: + y: int = 3 + + @dataclasses.dataclass(frozen=True) + class Nested: + x: int + b: B = B(y=5) + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + assert tyro.cli(Nested, args=["--x", "1"]) == Nested(x=1, b=B(y=5)) + + +def test_double_default_nested() -> None: + @dataclasses.dataclass(frozen=True) + class Child: + y: int + + @dataclasses.dataclass(frozen=True) + class Parent: + c: Child + + @dataclasses.dataclass(frozen=True) + class Grandparent: + x: int + b: Parent = Parent(Child(y=5)) + + assert tyro.cli(Grandparent, args=["--x", "1", "--b.c.y", "3"]) == Grandparent( + x=1, b=Parent(Child(y=3)) + ) + assert tyro.cli(Grandparent, args=["--x", "1"]) == Grandparent( + x=1, b=Parent(Child(y=5)) + ) + + +def test_default_factory_nested() -> None: + @dataclasses.dataclass + class B: + y: int = 3 + + @dataclasses.dataclass + class Nested: + x: int + b: B = dataclasses.field(default_factory=lambda: B(y=5)) + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + assert tyro.cli(Nested, args=["--x", "1"]) == Nested(x=1, b=B(y=5)) + + +def test_optional_nested() -> None: + @dataclasses.dataclass + class OptionalNestedChild: + y: int + z: int + + @dataclasses.dataclass + class OptionalNested: + x: int + b: Optional[OptionalNestedChild] = None + + assert tyro.cli(OptionalNested, args=["--x", "1"]) == OptionalNested(x=1, b=None) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, args=["--x", "1", "b:optional-nested-child", "--b.y", "3"] + ) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, args=["--x", "1", "b:optional-nested-child", "--b.z", "3"] + ) + + assert tyro.cli( + OptionalNested, + args=["--x", "1", "b:optional-nested-child", "--b.y", "2", "--b.z", "3"], + ) == OptionalNested(x=1, b=OptionalNestedChild(y=2, z=3)) + + +def test_optional_nested_newtype() -> None: + @dataclasses.dataclass + class OptionalNestedChild: + y: int + z: int + + SpecialOptionalNestedChild = NewType( + "SpecialOptionalNestedChild", OptionalNestedChild + ) + + @dataclasses.dataclass + class OptionalNested: + x: int + b: Optional[SpecialOptionalNestedChild] = None + + assert tyro.cli(OptionalNested, args=["--x", "1"]) == OptionalNested(x=1, b=None) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, + args=["--x", "1", "b:special-optional-nested-child", "--b.y", "3"], + ) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, + args=["--x", "1", "b:special-optional-nested-child", "--b.z", "3"], + ) + + assert tyro.cli( + OptionalNested, + args=[ + "--x", + "1", + "b:special-optional-nested-child", + "--b.y", + "2", + "--b.z", + "3", + ], + ) == OptionalNested( + x=1, b=SpecialOptionalNestedChild(OptionalNestedChild(y=2, z=3)) + ) + + +def test_optional_nested_multiple() -> None: + """Adapted from: https://github.com/brentyi/tyro/issues/60""" + + @dataclasses.dataclass(frozen=True) + class OutputHeadSettings: + number_of_outputs: int = 1 + + @dataclasses.dataclass(frozen=True) + class OptimizerSettings: + name: str = "adam" + + @dataclasses.dataclass(frozen=True) + class ModelSettings: + output_head_settings: Optional[OutputHeadSettings] = None + optimizer_settings: Optional[OptimizerSettings] = None + + assert tyro.cli( + ModelSettings, + args="output-head-settings:None optimizer-settings:None".split(" "), + ) == ModelSettings(None, None) + + with pytest.raises(SystemExit): + # Order cannot be flipped, unfortunately. + tyro.cli( + ModelSettings, + args="optimizer-settings:None output-head-settings:None".split(" "), + ) + + assert tyro.cli( + ModelSettings, + args="output-head-settings:output-head-settings optimizer-settings:None".split( + " " + ), + ) == ModelSettings(OutputHeadSettings(1), None) + + assert tyro.cli( + ModelSettings, + args=( + "output-head-settings:output-head-settings" + " --output-head-settings.number-of-outputs 5 optimizer-settings:None".split( + " " + ) + ), + ) == ModelSettings(OutputHeadSettings(5), None) + + assert tyro.cli( + tyro.conf.OmitSubcommandPrefixes[ + tyro.conf.ConsolidateSubcommandArgs[ModelSettings] + ], + args=("output-head-settings None" " --number-of-outputs 5".split(" ")), + ) == ModelSettings(OutputHeadSettings(5), None) + + assert tyro.cli( + tyro.conf.OmitSubcommandPrefixes[ + tyro.conf.ConsolidateSubcommandArgs[ModelSettings] + ], + args=( + "output-head-settings" + " optimizer-settings --name sgd --number-of-outputs 5".split(" ") + ), + ) == ModelSettings(OutputHeadSettings(5), OptimizerSettings("sgd")) + + +def test_subparser() -> None: + @dataclasses.dataclass + class HTTPServer: + y: int + + @dataclasses.dataclass + class SMTPServer: + z: int + + @dataclasses.dataclass + class Subparser: + x: int + bc: Union[HTTPServer, SMTPServer] + + assert tyro.cli( + Subparser, args=["--x", "1", "bc:http-server", "--bc.y", "3"] + ) == Subparser(x=1, bc=HTTPServer(y=3)) + assert tyro.cli( + Subparser, args=["--x", "1", "bc:smtp-server", "--bc.z", "3"] + ) == Subparser(x=1, bc=SMTPServer(z=3)) + + with pytest.raises(SystemExit): + # Missing subcommand. + tyro.cli(Subparser, args=["--x", "1"]) + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli(Subparser, args=["--x", "1", "bc:http-server", "--bc.z", "3"]) + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli(Subparser, args=["--x", "1", "bc:smtp-server", "--bc.y", "3"]) + + +def test_subparser_root() -> None: + @dataclasses.dataclass + class HTTPServer: + y: int + + @dataclasses.dataclass + class SMTPServer: + z: int + + @dataclasses.dataclass + class Subparser: + x: int + bc: Union[HTTPServer, SMTPServer] + + assert tyro.cli( + Union[HTTPServer, SMTPServer], + args=["http-server", "--y", "3"], # type: ignore + ) == HTTPServer(y=3) + + +def test_subparser_with_default() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + bc: Union[DefaultHTTPServer, DefaultSMTPServer] = dataclasses.field( + default_factory=lambda: DefaultHTTPServer(5) + ) + + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "5"] + ) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-smtp-server", "--bc.z", "3"] + ) == DefaultSubparser(x=1, bc=DefaultSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "8"] + ) + == tyro.cli( + DefaultSubparser, + args=[], + default=DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)), + ) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_with_default_and_newtype() -> None: + @dataclasses.dataclass + class DefaultHTTPServer_: + y: int + + DefaultHTTPServer__ = NewType("DefaultHTTPServer__", DefaultHTTPServer_) + DefaultHTTPServer = NewType("DefaultHTTPServer", DefaultHTTPServer__) # type: ignore + # ^nesting NewType is not technically allowed and pyright will complain, + # but we should try to be robust to it anyways. + + def make_http_server(y: int) -> DefaultHTTPServer: + return DefaultHTTPServer(DefaultHTTPServer__(DefaultHTTPServer_(y))) + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + bc: Union[DefaultHTTPServer, DefaultSMTPServer] = dataclasses.field( + default_factory=lambda: make_http_server(5) + ) + + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "5"] + ) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=make_http_server(y=5)) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-smtp-server", "--bc.z", "3"] + ) == DefaultSubparser(x=1, bc=DefaultSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "8"] + ) + == tyro.cli( + DefaultSubparser, + args=[], + default=DefaultSubparser(x=1, bc=make_http_server(y=8)), + ) + == DefaultSubparser(x=1, bc=make_http_server(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_with_default_alternate() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + bc: Union[DefaultInstanceHTTPServer, DefaultInstanceSMTPServer] + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "5"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=[], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["bc:default-instance-http-server"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultInstanceSubparser, + args=["bc:default-instance-smtp-server", "--bc.z", "3"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "8"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=[], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultInstanceSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultInstanceSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_with_default_bad() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + bc: Union[DefaultHTTPServer, DefaultSMTPServer] = dataclasses.field( + default_factory=lambda: 5 # type: ignore + ) + + # Tolerate bad static types: https://github.com/brentyi/tyro/issues/20 + # Should give us a bunch of warnings! + with pytest.warns(UserWarning): + assert tyro.cli(DefaultSubparser, args=["--x", "1"]) == DefaultSubparser( + 1, + 5, # type: ignore + ) + + +def test_subparser_with_default_bad_alt() -> None: + @dataclasses.dataclass + class A: + a: int + + @tyro.conf.configure(tyro.conf.subcommand(name="bbbb")) + @dataclasses.dataclass + class B: + b: int + + @dataclasses.dataclass + class C: + c: int + + with pytest.warns(UserWarning): + assert tyro.cli( + Union[A, Annotated[B, None]], # type: ignore + default=C(3), + args=["c", "--c", "2"], + ) == C(2) + + +def test_optional_subparser() -> None: + @dataclasses.dataclass + class OptionalHTTPServer: + y: int + + @dataclasses.dataclass + class OptionalSMTPServer: + z: int + + @dataclasses.dataclass + class OptionalSubparser: + x: int + bc: Optional[Union[OptionalHTTPServer, OptionalSMTPServer]] + + assert tyro.cli( + OptionalSubparser, args=["--x", "1", "bc:optional-http-server", "--bc.y", "3"] + ) == OptionalSubparser(x=1, bc=OptionalHTTPServer(y=3)) + assert tyro.cli( + OptionalSubparser, args=["--x", "1", "bc:optional-smtp-server", "--bc.z", "3"] + ) == OptionalSubparser(x=1, bc=OptionalSMTPServer(z=3)) + assert tyro.cli( + OptionalSubparser, args=["--x", "1", "bc:None"] + ) == OptionalSubparser(x=1, bc=None) + + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli( + OptionalSubparser, + args=["--x", "1", "bc:optional-http-server", "--bc.z", "3"], + ) + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli( + OptionalSubparser, + args=["--x", "1", "bc:optional-smtp-server", "--bc.y", "3"], + ) + + +def test_post_init_default() -> None: + @dataclasses.dataclass + class DataclassWithDynamicDefault: + x: int = 3 + y: Optional[int] = None + + def __post_init__(self): + # If unspecified, set y = x. + if self.y is None: + self.y = self.x + + @dataclasses.dataclass + class NoDefaultPostInitArgs: + inner: DataclassWithDynamicDefault + + @dataclasses.dataclass + class DefaultFactoryPostInitArgs: + inner: DataclassWithDynamicDefault = dataclasses.field( + default_factory=DataclassWithDynamicDefault + ) + + assert ( + tyro.cli(NoDefaultPostInitArgs, args=["--inner.x", "5"]).inner + == tyro.cli(DefaultFactoryPostInitArgs, args=["--inner.x", "5"]).inner + == DataclassWithDynamicDefault(x=5, y=5) + ) + + +def test_multiple_subparsers() -> None: + @dataclasses.dataclass + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass + class Subcommand2: + y: int = 1 + + @dataclasses.dataclass + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + a: Union[Subcommand1, Subcommand2, Subcommand3] + b: Union[Subcommand1, Subcommand2, Subcommand3] + c: Union[Subcommand1, Subcommand2, Subcommand3] + + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=[]) + + assert tyro.cli( + MultipleSubparsers, args="a:subcommand1 b:subcommand2 c:subcommand3".split(" ") + ) == MultipleSubparsers(Subcommand1(), Subcommand2(), Subcommand3()) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand1 --a.x 5 b:subcommand2 --b.y 7 c:subcommand3 --c.z 3".split( + " " + ), + ) == MultipleSubparsers(Subcommand1(x=5), Subcommand2(y=7), Subcommand3(z=3)) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand2 --a.y 5 b:subcommand1 --b.x 7 c:subcommand3 --c.z 3".split( + " " + ), + ) == MultipleSubparsers(Subcommand2(y=5), Subcommand1(x=7), Subcommand3(z=3)) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand3 --a.z 5 b:subcommand1 --b.x 7 c:subcommand3 --c.z 3".split( + " " + ), + ) == MultipleSubparsers(Subcommand3(z=5), Subcommand1(x=7), Subcommand3(z=3)) + + +def test_multiple_subparsers_with_default() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: int = 1 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + a: Union[Subcommand1, Subcommand2, Subcommand3] = Subcommand1(tyro.MISSING) + b: Union[Subcommand1, Subcommand2, Subcommand3] = Subcommand2(7) + c: Union[Subcommand1, Subcommand2, Subcommand3] = Subcommand3(3) + + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=[], + ) + + assert tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "--a.x", "5"], + ) == MultipleSubparsers(Subcommand1(x=5), Subcommand2(y=7), Subcommand3(z=3)) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand1 --a.x 3".split(" "), + ) == MultipleSubparsers(Subcommand1(x=3), Subcommand2(y=7), Subcommand3(z=3)) + + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=[], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=[ + "a:subcommand1", + ], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2", "c:subcommand3"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + assert tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2", "c:subcommand3", "--c.z", "3"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) == MultipleSubparsers(Subcommand1(x=0), Subcommand2(y=1), Subcommand3(z=3)) + assert tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2", "c:subcommand2"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) == MultipleSubparsers(Subcommand1(x=0), Subcommand2(y=1), Subcommand2(y=1)) + + +def test_nested_subparsers_with_default() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: Union[Subcommand1, Subcommand3] + + @dataclasses.dataclass(frozen=True) + class MultipleSubparsers: + a: Union[Subcommand1, Subcommand2] = Subcommand2(Subcommand1(tyro.MISSING)) + + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=[]) + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=["a:subcommand2"]) + + assert tyro.cli( + MultipleSubparsers, args="a:subcommand1 --a.x 3".split(" ") + ) == MultipleSubparsers(Subcommand1(3)) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand2 a.y:subcommand3 --a.y.z 2".split(" ") + ) == MultipleSubparsers(Subcommand2(Subcommand3())) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand2 a.y:subcommand3 --a.y.z 7".split(" ") + ) == MultipleSubparsers(Subcommand2(Subcommand3(7))) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand2 a.y:subcommand1 --a.y.x 7".split(" ") + ) == MultipleSubparsers(Subcommand2(Subcommand1(7))) + + +def test_nested_subparsers_multiple() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: Union[Subcommand1, Subcommand3] + + @dataclasses.dataclass(frozen=True) + class MultipleSubparsers: + a: Union[Subcommand1, Subcommand2] + b: Union[Subcommand1, Subcommand2] + + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=[]) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand1 b:subcommand1".split(" ") + ) == MultipleSubparsers(Subcommand1(), Subcommand1()) + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand1 b:subcommand2 b.y:subcommand1".split(" "), + ) == MultipleSubparsers(Subcommand1(), Subcommand2(Subcommand1())) + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand2 a.y:subcommand1 b:subcommand2 b.y:subcommand1".split(" "), + ) == MultipleSubparsers(Subcommand2(Subcommand1()), Subcommand2(Subcommand1())) + assert tyro.cli( + MultipleSubparsers, + args=( + "a:subcommand2 a.y:subcommand1 --a.y.x 3 b:subcommand2 b.y:subcommand1" + " --b.y.x 7".split(" ") + ), + ) == MultipleSubparsers(Subcommand2(Subcommand1(3)), Subcommand2(Subcommand1(7))) + + +def test_nested_subparsers_multiple_in_container() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: Union[Subcommand1, Subcommand3] + + @dataclasses.dataclass(frozen=True) + class MultipleSubparsers: + a: Union[Subcommand1, Subcommand2] + b: Union[Subcommand1, Subcommand2] + + @dataclasses.dataclass(frozen=True) + class Root: + inner: MultipleSubparsers + + with pytest.raises(SystemExit): + tyro.cli(Root, args=[]) + assert tyro.cli( + Root, args="inner.a:subcommand1 inner.b:subcommand1".split(" ") + ) == Root(MultipleSubparsers(Subcommand1(), Subcommand1())) + assert tyro.cli( + Root, + args="inner.a:subcommand1 inner.b:subcommand2 inner.b.y:subcommand1".split(" "), + ) == Root(MultipleSubparsers(Subcommand1(), Subcommand2(Subcommand1()))) + assert tyro.cli( + Root, + args=( + "inner.a:subcommand2 inner.a.y:subcommand1 inner.b:subcommand2" + " inner.b.y:subcommand1".split(" ") + ), + ) == Root( + MultipleSubparsers(Subcommand2(Subcommand1()), Subcommand2(Subcommand1())) + ) + assert tyro.cli( + Root, + args=( + "inner.a:subcommand2 inner.a.y:subcommand1 --inner.a.y.x 3" + " inner.b:subcommand2 inner.b.y:subcommand1 --inner.b.y.x 7".split(" ") + ), + ) == Root( + MultipleSubparsers(Subcommand2(Subcommand1(3)), Subcommand2(Subcommand1(7))) + ) + + +def test_tuple_nesting() -> None: + @dataclasses.dataclass(frozen=True) + class Color: + r: int + g: int + b: int + + @dataclasses.dataclass(frozen=True) + class Location: + x: float + y: float + z: float + + def main(x: Tuple[Tuple[Color], Location, float]): + return x + + assert tyro.cli( + main, + args=( + "--x.0.0.r 255 --x.0.0.g 0 --x.0.0.b 0 --x.1.x 5.0 --x.1.y 0.0" + " --x.1.z 2.0 --x.2 4.0".split(" ") + ), + ) == ((Color(255, 0, 0),), Location(5.0, 0.0, 2.0), 4.0) + + +def test_tuple_nesting_union() -> None: + @dataclasses.dataclass(frozen=True) + class Color: + r: int + g: int + b: int + + @dataclasses.dataclass(frozen=True) + class Location: + x: float + y: float + z: float + + def main(x: Union[Tuple[Tuple[Color], Location, float], Tuple[Color, Color]]): + return x + + assert tyro.cli( + main, + args=( + "x:tuple-tuple-color-location-float --x.0.0.r 255 --x.0.0.g 0 --x.0.0.b 0" + " --x.1.x 5.0 --x.1.y 0.0 --x.1.z 2.0 --x.2 4.0".split(" ") + ), + ) == ((Color(255, 0, 0),), Location(5.0, 0.0, 2.0), 4.0) + + +def test_generic_subparsers() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class A(Generic[T]): + x: T + + def main(x: Union[A[int], A[float]]) -> Any: + return x + + assert tyro.cli(main, args="x:a-float --x.x 3.2".split(" ")) == A(3.2) + assert tyro.cli(main, args="x:a-int --x.x 3".split(" ")) == A(3) + + def main_with_default(x: Union[A[int], A[float]] = A(5)) -> Any: + return x + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main_with_default, args=[]) + + +def test_generic_inherited() -> None: + class UnrelatedParentClass: + pass + + T = TypeVar("T") + + @dataclasses.dataclass + class ActualParentClass(Generic[T]): + x: T # Documentation 1 + + # Documentation 2 + y: T + + z: T = 3 # type: ignore + """Documentation 3""" + + @dataclasses.dataclass + class ChildClass(UnrelatedParentClass, ActualParentClass[int]): + pass + + assert tyro.cli( + ChildClass, args=["--x", "1", "--y", "2", "--z", "3"] + ) == ChildClass(x=1, y=2, z=3) + + +def test_subparser_in_nested() -> None: + @dataclasses.dataclass + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + + @dataclasses.dataclass + class Nested2: + subcommand: Union[A, B] + + @dataclasses.dataclass + class Nested1: + nested2: Nested2 + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:a --nested1.nested2.subcommand.a 3".split(" "), + ) == Parent(Nested1(Nested2(A(3)))) + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:b --nested1.nested2.subcommand.b 7".split(" "), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_frozen_dict() -> None: + def main( + x: Mapping[str, float] = frozendict( # type: ignore + { + "num_epochs": 20, + "batch_size": 64, + } + ), + ): + return x + + assert hash(tyro.cli(main, args="--x.num-epochs 10".split(" "))) == hash( + frozendict({"num_epochs": 10, "batch_size": 64}) # type: ignore + ) + + +def test_nested_in_subparser() -> None: + # https://github.com/brentyi/tyro/issues/9 + @dataclasses.dataclass(frozen=True) + class Subtype: + data: int = 1 + + @dataclasses.dataclass(frozen=True) + class TypeA: + subtype: Subtype = Subtype(1) + + @dataclasses.dataclass(frozen=True) + class TypeB: + subtype: Subtype = Subtype(2) + + @dataclasses.dataclass(frozen=True) + class Wrapper: + supertype: Union[TypeA, TypeB] = TypeA() + + assert tyro.cli(Wrapper, args=[]) == Wrapper() + assert ( + tyro.cli(Wrapper, args="supertype:type-a --supertype.subtype.data 1".split(" ")) + == Wrapper() + ) + + +def test_nested_in_subparser_override_with_default() -> None: + @dataclasses.dataclass(frozen=True) + class Mnist: + binary: bool = False + """Set to load binary version of MNIST dataset.""" + + @dataclasses.dataclass + class ImageNet: + subset: Literal[50, 100, 1000] + """Choose between ImageNet-50, ImageNet-100, ImageNet-1000, etc.""" + + # Possible optimizer configurations. + + Selector = tyro.extras.subcommand_type_from_defaults( + { + "m": Mnist(), + "i": ImageNet(50), + } + ) + + @dataclasses.dataclass(frozen=True) + class DatasetContainer: + dataset: Selector = Mnist() # type: ignore + + @dataclasses.dataclass + class Adam: + learning_rate: float = 1e-3 + betas: Tuple[float, float] = (0.9, 0.999) + container: DatasetContainer = DatasetContainer() + + @dataclasses.dataclass + class Sgd: + learning_rate: float = 3e-4 + container: DatasetContainer = DatasetContainer() + + # Train script. + + Optimizers = tyro.extras.subcommand_type_from_defaults( + {"adam": Adam(container=DatasetContainer(ImageNet(50))), "sgd": Sgd()}, + prefix_names=False, + ) + + @tyro.conf.configure(tyro.conf.OmitSubcommandPrefixes) + def train( + optimizer: Optimizers = Adam(container=DatasetContainer(ImageNet(50))), # type: ignore + ) -> Union[Adam, Sgd]: + return optimizer + + assert tyro.cli(train, args=[]) == Adam(container=DatasetContainer(ImageNet(50))) + assert tyro.cli(train, args=["adam"]) == Adam( + container=DatasetContainer(ImageNet(50)) + ) + assert tyro.cli(train, args=["sgd"]) == Sgd(container=DatasetContainer(Mnist())) + + +def test_underscore_prefix() -> None: + """https://github.com/brentyi/tyro/issues/77""" + + @dataclasses.dataclass + class PrivateConfig: + pass + + @dataclasses.dataclass + class BaseConfig: + _private: PrivateConfig = dataclasses.field( + default_factory=lambda: PrivateConfig() + ) + + @dataclasses.dataclass + class Level1(BaseConfig): + pass + + @dataclasses.dataclass + class Level2(BaseConfig): + child: Level1 = dataclasses.field(default_factory=lambda: Level1()) + + @dataclasses.dataclass + class Level3(BaseConfig): + child: Level2 = dataclasses.field(default_factory=lambda: Level2()) + + tyro.cli(Level3, args=[]) + + +def test_subcommand_dict_helper() -> None: + def checkout(branch: str) -> str: + """Check out a branch.""" + return branch + + def commit(message: str, all: bool = False) -> Tuple[str, bool]: + """Make a commit.""" + return (message, all) + + # If we only get one, we unfortunately can't form subcommands. This is + # because unions in Python require at least 2 types. + with pytest.raises(AssertionError): + tyro.extras.subcommand_cli_from_dict( + { + "commit": commit, + }, + args="--message hello --all".split(" "), + ) + + assert ( + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + }, + args="checkout --branch main".split(" "), + ) + == "main" + ) + assert tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + }, + args="commit --message hello".split(" "), + ) == ("hello", False) + assert tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + }, + args="commit --message hello --all".split(" "), + ) == ("hello", True) + + +def test_subcommand_by_type_tree() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: Union[int, str] + + @dataclasses.dataclass + class Args: + inner: Union[ + Annotated[A, tyro.conf.subcommand(name="alt", default=A(5))], A + ] = A("hello") + + assert tyro.cli(Args, args=[]) == Args(A("hello")) + assert "default: inner:alt" in get_helptext_with_checks(Args) diff --git a/tests/test_nested_in_containers.py b/tests/test_nested_in_containers.py new file mode 100644 index 000000000..819998176 --- /dev/null +++ b/tests/test_nested_in_containers.py @@ -0,0 +1,379 @@ +import dataclasses +import enum +from typing import Any, Dict, Generic, List, Set, Tuple, TypeVar + +import pytest + +import tyro + + +@dataclasses.dataclass(frozen=True) +class Color: + r: int + g: int + b: int + + +def test_nested_tuple_fixed_single() -> None: + def main(x: Tuple[Color]) -> Any: + return x + + assert tyro.cli(main, args="--x.0.r 255 --x.0.g 127 --x.0.b 5".split(" ")) == ( + Color(255, 127, 5), + ) + + +def test_nested_tuple_fixed_two() -> None: + def main(x: Tuple[Color, Color]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 255 --x.0.g 127 --x.0.b 5 --x.1.r 255 --x.1.g 127 --x.1.b 0" + ).split(" "), + ) == ( + Color(255, 127, 5), + Color(255, 127, 0), + ) + + +def test_nested_tuple_fixed_three() -> None: + def main(x: Tuple[Color, int, Color]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 255 --x.0.g 127 --x.0.b 5 --x.1 94709 --x.2.r 255 --x.2.g 127" + " --x.2.b 0" + ).split(" "), + ) == ( + Color(255, 127, 5), + 94709, + Color(255, 127, 0), + ) + + +def test_nested_tuple_recursive() -> None: + def main(x: Tuple[Color, Tuple[Color, Color]]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 255 --x.0.g 127 --x.0.b 5 --x.1.0.r 255 --x.1.0.g 127 --x.1.0.b 0" + " --x.1.1.r 255 --x.1.1.g 127 --x.1.1.b 0" + ).split(" "), + ) == ( + Color(255, 127, 5), + ( + Color(255, 127, 0), + Color(255, 127, 0), + ), + ) + + +def test_tuple_bad() -> None: + # Unable to infer input length. + def main(x: Tuple[Color, ...]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_set_bad() -> None: + # Unable to infer input length. + def main(x: Set[Color]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_set_ok() -> None: + def main(x: Set[Color] = {Color(255, 0, 0)}) -> Any: + return x + + assert tyro.cli(main, args=[]) == {Color(255, 0, 0)} + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == {Color(127, 0, 0)} + + +def test_list_bad() -> None: + # Unable to infer input length. + def main(x: List[Color]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_list_ok() -> None: + def main(x: List[Color] = [Color(255, 0, 0)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [Color(255, 0, 0)] + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == [Color(127, 0, 0)] + + +def test_list_object() -> None: + def main(x: List[object] = [Color(255, 0, 0)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [Color(255, 0, 0)] + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == [Color(127, 0, 0)] + + +def test_list_any() -> None: + def main(x: List[Any] = [Color(255, 0, 0)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [Color(255, 0, 0)] + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == [Color(127, 0, 0)] + + +def test_tuple_in_list() -> None: + def main(x: List[Tuple[Color]] = [(Color(255, 0, 0),)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [(Color(255, 0, 0),)] + assert tyro.cli(main, args="--x.0.0.r 127".split(" ")) == [(Color(127, 0, 0),)] + + +def test_tuple_variable() -> None: + def main(x: Tuple[Color, ...] = (Color(255, 0, 0), Color(255, 0, 127))) -> Any: + return x + + assert tyro.cli(main, args=[]) == (Color(255, 0, 0), Color(255, 0, 127)) + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == ( + Color(127, 0, 0), + Color(255, 0, 127), + ) + + +def test_dict_bad() -> None: + def main(x: Dict[str, Color]) -> Any: + return x + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_dict_ok() -> None: + def main( + x: Dict[str, Color] = { + "red": Color(255, 0, 0), + "green": Color(0, 255, 0), + "blue": Color(0, 0, 255), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])["green"] == Color(0, 255, 0) + assert tyro.cli(main, args="--x.green.g 127".split(" "))["green"] == Color( + 0, 127, 0 + ) + + +def test_dict_key_int() -> None: + def main( + x: Dict[int, Color] = { + 0: Color(255, 0, 0), + 1: Color(0, 255, 0), + 2: Color(0, 0, 255), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])[1] == Color(0, 255, 0) + assert tyro.cli(main, args="--x.1.g 127".split(" "))[1] == Color(0, 127, 0) + + +def test_dict_key_enum() -> None: + class ColorType(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + def main( + x: Dict[ColorType, Color] = { + ColorType.RED: Color(255, 0, 0), + ColorType.GREEN: Color(0, 255, 0), + ColorType.BLUE: Color(0, 0, 255), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])[ColorType.GREEN] == Color(0, 255, 0) + assert tyro.cli(main, args="--x.GREEN.g 127".split(" "))[ColorType.GREEN] == Color( + 0, 127, 0 + ) + + +def test_dict_nested() -> None: + def main( + x: Dict[str, Tuple[Color, int]] = { + # For each color: RGB and xterm color code. + "red": (Color(255, 0, 0), 9), + "green": (Color(0, 255, 0), 10), + "blue": (Color(0, 0, 255), 12), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])["green"] == (Color(0, 255, 0), 10) + assert tyro.cli(main, args="--x.green.0.g 127 --x.green.1 2".split(" "))[ + "green" + ] == (Color(0, 127, 0), 2) + + +def test_generic_in_tuple() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main(x: Tuple[GenericColor[float], GenericColor[int]]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 0.5 --x.0.g 0.2 --x.0.b 0.3 --x.1.r 25 --x.1.g 2 --x.1.b 3".split( + " " + ) + ), + ) == (GenericColor(0.5, 0.2, 0.3), GenericColor(25, 2, 3)) + + +def test_generic_in_tuple_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Tuple[GenericColor[float], GenericColor[int]] = ( + GenericColor(0.5, 0.2, 0.3), + GenericColor[int](25, 2, 3), # The subscript should be optional. + ), + ) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 0.5 --x.0.g 0.2 --x.0.b 0.3 --x.1.r 25 --x.1.g 2 --x.1.b 3".split( + " " + ) + ), + ) == (GenericColor(0.5, 0.2, 0.3), GenericColor(25, 2, 3)) + + +def test_generic_in_variable_tuple_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Tuple[GenericColor, ...] = ( + GenericColor[float](0.5, 0.2, 0.3), + GenericColor[int](25, 2, 3), + ), + ) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 0.5 --x.0.g 0.9 --x.0.b 0.3 --x.1.r 25 --x.1.g 2 --x.1.b 3".split( + " " + ) + ), + ) == (GenericColor(0.5, 0.9, 0.3), GenericColor(25, 2, 3)) + + +def test_generic_in_dict_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Dict[str, GenericColor] = { + "float": GenericColor(0.5, 0.2, 0.3), + "int": GenericColor[int](25, 2, 3), + }, + ) -> Any: + return x + + assert tyro.cli( + main, + args="--x.float.g 0.1".split(" "), + )["float"] == GenericColor(0.5, 0.1, 0.3) + assert tyro.cli( + main, + args="--x.int.g 0".split(" "), + ) == {"float": GenericColor(0.5, 0.2, 0.3), "int": GenericColor(25, 0, 3)} + + +def test_generic_in_double_nested_dict_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Dict[str, Dict[str, GenericColor]] = { + "hello": { + "float": GenericColor(0.5, 0.2, 0.3), + "int": GenericColor[int](25, 2, 3), + } + }, + ) -> Any: + return x + + assert tyro.cli( + main, + args="--x.hello.float.g 0.1".split(" "), + )["hello"]["float"] == GenericColor(0.5, 0.1, 0.3) + assert tyro.cli( + main, + args="--x.hello.int.g 0".split(" "), + ) == { + "hello": {"float": GenericColor(0.5, 0.2, 0.3), "int": GenericColor(25, 0, 3)} + } + + +def test_double_nested_dict_with_inferred_type() -> None: + def main( + x: Dict[str, Any] = { + "hello": { + "a": Color(5, 2, 3), + "b": Color(25, 2, 3), + } + }, + ) -> Any: + return x + + assert tyro.cli( + main, + args="--x.hello.a.g 1".split(" "), + )["hello"]["a"] == Color(5, 1, 3) diff --git a/tests/test_new_style_annotations_generics.py b/tests/test_new_style_annotations_generics.py new file mode 100644 index 000000000..bf365f0f7 --- /dev/null +++ b/tests/test_new_style_annotations_generics.py @@ -0,0 +1,55 @@ +from __future__ import annotations # Should enable support for all versions of Python. + +from typing import Any, Optional, Union + +import pytest +from typing_extensions import Literal + +import tyro + + +def test_list() -> None: + def main(x: list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_tuple() -> None: + def main(x: tuple[bool, str]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == (True, "False") + + +def test_tuple_variable() -> None: + def main(x: tuple[Union[bool, str], ...]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "Wrong"]) == (True, "Wrong") + + +def test_super_nested() -> None: + def main( + x: Optional[ + list[ + tuple[ + Optional[int], + Literal[3, 4], + Union[tuple[int, int], tuple[str, str]], + ] + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) diff --git a/tests/test_new_style_annotations_min_py310.py b/tests/test_new_style_annotations_min_py310.py new file mode 100644 index 000000000..dc17e3db4 --- /dev/null +++ b/tests/test_new_style_annotations_min_py310.py @@ -0,0 +1,64 @@ +from typing import Any, Literal + +import pytest + +import tyro + + +def test_union_direct(): + assert tyro.cli(int | str, args=["5"]) == 5 + assert tyro.cli(int | str, args=["five"]) == "five" + + +def test_union_basic(): + def main(x: int | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_union_with_list(): + def main(x: int | str | list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + assert tyro.cli(main, args=["--x", "True"]) == "True" + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_union_literal(): + def main(x: Literal[1, 2] | Literal[3, 4, 5] | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == "6" + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_super_nested(): + def main( + x: None + | list[ + tuple[ + None | int, + Literal[3, 4], + tuple[int, int] | tuple[str, str], + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) diff --git a/tests/test_new_style_annotations_min_py312.py b/tests/test_new_style_annotations_min_py312.py new file mode 100644 index 000000000..d04c73f3b --- /dev/null +++ b/tests/test_new_style_annotations_min_py312.py @@ -0,0 +1,54 @@ +# mypy: ignore-errors +# +# PEP 695 isn't yet supported in mypy. (April 4, 2024) +from dataclasses import dataclass + +import tyro + + +def test_simple_generic(): + @dataclass(frozen=True) + class Container[T]: + a: T + b: T + + assert tyro.cli(Container[int], args="--a 1 --b 2".split(" ")) == Container(1, 2) + + +type X = int +type Y = list[int] +type Z = Inner[int] + + +@dataclass(frozen=True) +class Inner[T]: + a: T + b: T + + +def test_generic_with_type_statement_0(): + @dataclass(frozen=True) + class Container[T]: + a: T + b: T + + assert tyro.cli(Container[X], args="--a 1 --b 2".split(" ")) == Container(1, 2) + + +def test_generic_with_type_statement_1(): + @dataclass(frozen=True) + class Container[T]: + a: tuple[X, ...] + b: T + + assert tyro.cli(Container[Y], args="--a 1 --b 2".split(" ")) == Container((1,), [2]) + + +def test_generic_with_type_statement_2(): + @dataclass(frozen=True) + class Container[T]: + a: Z + + assert tyro.cli(Container[Y], args="--a.a 1 --a.b 2".split(" ")) == Container( + Inner(1, 2) + ) diff --git a/tests/test_new_style_annotations_min_py39.py b/tests/test_new_style_annotations_min_py39.py new file mode 100644 index 000000000..9b8908a35 --- /dev/null +++ b/tests/test_new_style_annotations_min_py39.py @@ -0,0 +1,69 @@ +import dataclasses +from typing import Any, Literal, Optional, Union + +import pytest + +import tyro + + +def test_list() -> None: + def main(x: list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_tuple() -> None: + def main(x: tuple[bool, str]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == (True, "False") + + +def test_tuple_nested() -> None: + @dataclasses.dataclass + class Args: + a: int + + def main(x: tuple[Args, Args]) -> Any: + return x + + assert tyro.cli(main, args=["--x.0.a", "3", "--x.1.a", "4"]) == (Args(3), Args(4)) + + +def test_tuple_variable() -> None: + def main(x: tuple[Union[bool, str], ...]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "Wrong"]) == (True, "Wrong") + + +def test_super_nested() -> None: + def main( + x: Optional[ + list[ + tuple[ + Optional[int], + Literal[3, 4], + Union[tuple[int, int], tuple[str, str]], + ] + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_tuple_direct() -> None: + assert tyro.cli(tuple[int, ...], args="1 2".split(" ")) == (1, 2) # type: ignore + assert tyro.cli(tuple[int, int], args="1 2".split(" ")) == (1, 2) # type: ignore diff --git a/tests/test_new_style_annotations_unions.py b/tests/test_new_style_annotations_unions.py new file mode 100644 index 000000000..22b38ab72 --- /dev/null +++ b/tests/test_new_style_annotations_unions.py @@ -0,0 +1,62 @@ +from __future__ import annotations # Should enable support for all versions of Python. + +from typing import Any + +import pytest +from typing_extensions import Literal + +import tyro + + +def test_union_basic(): + def main(x: int | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_union_with_list(): + def main(x: int | str | list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + assert tyro.cli(main, args=["--x", "True"]) == "True" + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_union_literal(): + def main(x: Literal[1, 2] | Literal[3, 4, 5] | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == "6" + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_super_nested(): + def main( + x: None + | list[ + tuple[ + None | int, + Literal[3, 4], + tuple[int, int] | tuple[str, str], + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) diff --git a/tests/test_optional_annotated.py b/tests/test_optional_annotated.py new file mode 100644 index 000000000..640c30695 --- /dev/null +++ b/tests/test_optional_annotated.py @@ -0,0 +1,57 @@ +"""Test adapted from @KolinGuo: https://github.com/brentyi/tyro/issues/156""" + +import dataclasses +from typing import Optional, Tuple + +from typing_extensions import Annotated + +import tyro + + +def detect_usb_device( + vendor_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--vid"])] = None, + product_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--pid"])] = "", +) -> Tuple[Optional[str], Optional[str]]: + """ + Detect connected USB device by vendor_id + product_id. + + :param vendor_id: vendor_id of a specific USB device. + :param product_id: product_id of a specific USB device. + """ + return vendor_id, product_id + + +def test_detect_usb_device() -> None: + assert tyro.cli(detect_usb_device, args=["--vid", "0x1234", "--pid", "0x5678"]) == ( + "0x1234", + "0x5678", + ) + assert tyro.cli( + detect_usb_device, args=["--vendor-id", "0x1234", "--product-id", "0x5678"] + ) == ( + "0x1234", + "0x5678", + ) + assert tyro.cli(detect_usb_device, args=[]) == (None, "") + + +@dataclasses.dataclass +class DeviceStruct: + vendor_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--vid"])] = None + product_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--pid"])] = "" + + +def test_detect_usb_device_dataclass() -> None: + assert tyro.cli( + DeviceStruct, args=["--vid", "0x1234", "--pid", "0x5678"] + ) == DeviceStruct( + "0x1234", + "0x5678", + ) + assert tyro.cli( + DeviceStruct, args=["--vendor-id", "0x1234", "--product-id", "0x5678"] + ) == DeviceStruct( + "0x1234", + "0x5678", + ) + assert tyro.cli(DeviceStruct, args=[]) == DeviceStruct(None, "") diff --git a/tests/test_partial.py b/tests/test_partial.py new file mode 100644 index 000000000..89fe78552 --- /dev/null +++ b/tests/test_partial.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import dataclasses +import functools + +from helptext_utils import get_helptext_with_checks + +import tyro + + +def test_partial_func() -> None: + def main(a: int, b: str) -> str: + return b * a + + assert tyro.cli(functools.partial(main, a=3), args=["--b", "hi"]) == "hihihi" + + +def test_partial_class() -> None: + class Main: + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + assert tyro.cli(functools.partial(Main, a=3), args=["--b", "hi"]).inner == "hihihi" + + +def test_partial_helptext_func() -> None: + def main(a: int, b: str) -> str: + """Hello!""" + return b * a + + helptext = get_helptext_with_checks(functools.partial(main, b="3")) + assert "partial" not in helptext + assert "Hello!" in helptext + + +def test_partial_helptext_class() -> None: + class Main: + """Hello!""" + + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + helptext = get_helptext_with_checks(functools.partial(Main, b="3")) + assert "partial" not in helptext + assert "Hello!" in helptext + + +def test_partial_helptext_dataclass() -> None: + @dataclasses.dataclass + class Config: + a: int + b: int + """Hello!""" + c: int + + ConfigWithDefaults = functools.partial(functools.partial(Config, a=3), c=5) + assert tyro.cli(ConfigWithDefaults, args=["--b", "4"]) == Config(3, 4, 5) + + helptext = get_helptext_with_checks(ConfigWithDefaults) + assert "partial" not in helptext + assert "Hello!" in helptext diff --git a/tests/test_positional_min_py38.py b/tests/test_positional_min_py38.py new file mode 100644 index 000000000..15f8f7911 --- /dev/null +++ b/tests/test_positional_min_py38.py @@ -0,0 +1,118 @@ +from typing import List, Optional, Tuple + +import pytest + +import tyro + + +def test_positional(): + def main( + x: int, + y: int, + /, + # Note: it's generally a bad idea to have a mutable object (like a list) as a + # default value. But it should still work. + z: List[int] = [1, 2, 3], + ) -> Tuple[int, int, int]: + """main. + + Args: + x: x + y: y + z: z + + Returns: + Tuple[int, int, int]: Output. + """ + return (x, y, z[0]) + + assert tyro.cli(main, args="1 2 --z 3".split(" ")) == (1, 2, 3) + with pytest.raises(SystemExit): + assert tyro.cli(main, args="--x 1 --y 2 --z 3".split(" ")) == (1, 2, 3) + + +def test_nested_positional(): + class A: + def __init__(self, a: int, hello_world: int, /, c: int): + self.hello_world = hello_world + + def nest1(a: int, b: int, thing: A, /, c: int) -> A: + return thing + + assert isinstance(tyro.cli(nest1, args="0 1 2 3 4 --c 4".split(" ")), A) + assert tyro.cli(nest1, args="0 1 2 3 4 --c 4".split(" ")).hello_world == 3 + with pytest.raises(SystemExit): + tyro.cli(nest1, args="0 1 2 3 4 4 --c 4".split(" ")) + + +def test_nested_positional_alt(): + class B: + def __init__(self, a: int, b: int, /, c: int): + pass + + def nest2(a: int, b: int, /, thing: B, c: int): + return thing + + assert isinstance(tyro.cli(nest2, args="0 1 2 3 --thing.c 4 --c 4".split(" ")), B) + with pytest.raises(SystemExit): + tyro.cli(nest2, args="0 1 2 3 4 --thing.c 4 --c 4".split(" ")) + + +def test_positional_with_underscores(): + """Hyphen replacement works a bit different for positional arguments.""" + + def main(a_multi_word_input: int, /) -> int: + return a_multi_word_input + + assert tyro.cli(main, args=["5"]) == 5 + + +def test_positional_booleans(): + """Make sure that flag behavior is disabled for positional booleans.""" + + def main( + flag1: bool, + flag2: bool = True, + flag3: bool = False, + /, + ) -> Tuple[bool, bool, bool]: + return flag1, flag2, flag3 + + assert tyro.cli(main, args=["True"]) == (True, True, False) + assert tyro.cli(main, args=["True", "False"]) == (True, False, False) + assert tyro.cli(main, args=["False", "False", "True"]) == (False, False, True) + + with pytest.raises(SystemExit): + tyro.cli(main, args=["hmm"]) + with pytest.raises(SystemExit): + tyro.cli(main, args=["true"]) + with pytest.raises(SystemExit): + tyro.cli(main, args=["True", "false"]) + + +def test_optional_list(): + def main(a: Optional[List[int]], /) -> Optional[List[int]]: + return a + + assert tyro.cli(main, args=["None"]) is None + assert tyro.cli(main, args=["1", "2"]) == [1, 2] + assert tyro.cli(main, args=[]) == [] + with pytest.raises(SystemExit): + tyro.cli(main, args=["hm"]) + + +def test_optional_list_with_default(): + def main(a: Optional[List[int]] = None, /) -> Optional[List[int]]: + return a + + assert tyro.cli(main, args=["None"]) is None + assert tyro.cli(main, args=["5", "5"]) == [5, 5] + with pytest.raises(SystemExit): + tyro.cli(main, args=["None", "5"]) + + +def test_positional_tuple(): + def main(x: Tuple[int, int], y: Tuple[str, str], /): + return x, y + + assert tyro.cli(main, args="1 2 3 4".split(" ")) == ((1, 2), ("3", "4")) diff --git a/tests/test_py311_generated/README.md b/tests/test_py311_generated/README.md new file mode 100644 index 000000000..50b518af5 --- /dev/null +++ b/tests/test_py311_generated/README.md @@ -0,0 +1,3 @@ +This folder contains autogenerated tests, which replaces `typing_extension` +imports with `typing` imports and old-style union syntax (`Union[A, B]`) with +new ones (`A | B`). diff --git a/tests/test_py311_generated/_generate.py b/tests/test_py311_generated/_generate.py new file mode 100644 index 000000000..6ff72a438 --- /dev/null +++ b/tests/test_py311_generated/_generate.py @@ -0,0 +1,62 @@ +"""Generate a Python 3.11 version of tests. This will use imports from `typing` instead +of `typing_extensions`, and replace Union[A, B] types with A | B.""" + +import pathlib +import subprocess +from concurrent.futures import ThreadPoolExecutor + + +def generate_from_path(test_path: pathlib.Path) -> None: + content = test_path.read_text() + content = content.replace("typing_extensions", "typing") + + if "ReadOnly" in content: + content = content.replace("ReadOnly,", "", 1) + content = content.replace( + "\nfrom typing import ", + "\nfrom typing_extensions import ReadOnly\nfrom typing import ", + ) + + while "Union[" in content: + new_content, _, b = content.partition("Union[") + + if b.strip()[0] == '"': + break # Don't bother with forward references! + + new_content_parts = [new_content] + + bracket_count = 0 + for i, char in enumerate(b): + if char == "[": + bracket_count += 1 + elif char == "]": + bracket_count -= 1 + + if char == "," and bracket_count == 0: + new_content_parts.append("|") + elif bracket_count == -1: + while new_content_parts[-1] in (" ", "|"): + new_content_parts.pop(-1) + new_content_parts.append(b[i + 1 :]) + break + elif char != "\n": + new_content_parts.append(char) + + content = "".join(new_content_parts) + + out_path = pathlib.Path(__file__).absolute().parent / ( + test_path.stem + "_generated.py" + ) + out_path.write_text(content) + + subprocess.run(["ruff", "format", str(out_path)], check=True) + subprocess.run(["ruff", "check", "--fix", str(out_path)], check=True) + + +with ThreadPoolExecutor(max_workers=8) as executor: + list( + executor.map( + generate_from_path, + pathlib.Path(__file__).absolute().parent.parent.glob("test_*.py"), + ) + ) diff --git a/tests/test_py311_generated/test_attrs_generated.py b/tests/test_py311_generated/test_attrs_generated.py new file mode 100644 index 000000000..6f44068d3 --- /dev/null +++ b/tests/test_py311_generated/test_attrs_generated.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import contextlib +import io +import pathlib +from typing import cast + +import attr +import pytest +from attrs import define, field + +import tyro +import tyro._strings + + +def test_attrs_basic() -> None: + @attr.s + class ManyTypesA: + i: tyro.conf.Positional[int] = attr.ib() + s: str = attr.ib() + f: float = attr.ib() + p: pathlib.Path = attr.ib() + ignored: int = attr.ib(default=3, init=False) + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesA, + args=[ + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypesA(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + +def test_attrs_defaults() -> None: + @attr.s + class ManyTypesB: + i: int = attr.ib() + s: str = attr.ib() + f: float = attr.ib(default=1.0) + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesB, + args=[ + "--i", + "5", + "--s", + "5", + ], + ) == ManyTypesB(i=5, s="5", f=1.0) + + +def test_attrs_helptext() -> None: + @attr.s + class Helptext: + """This docstring should be printed as a description.""" + + x: int = attr.ib() # Documentation 1 + + # Documentation 2 + y: int = attr.ib() + + z: int = attr.ib(default=3) + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + assert tyro._strings.strip_ansi_sequences(cast(str, Helptext.__doc__)) in helptext + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_attrs_next_gen_and_factory() -> None: + @define + class Helptext: + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int = field(factory=lambda: 3) + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + assert tyro._strings.strip_ansi_sequences(cast(str, Helptext.__doc__)) in helptext + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_attrs_default_instance() -> None: + @attr.s + class ManyTypesB: + i: int = attr.ib() + s: str = attr.ib() + f: float = attr.ib(default=1.0) + k: float = attr.ib(default=1.0) + + assert tyro.cli( + ManyTypesB, + args=[ + "--i", + "5", + "--s", + "5", + ], + default=ManyTypesB(i=5, s="5", f=2.0), + ) == ManyTypesB(i=5, s="5", f=2.0) + assert tyro.cli( + ManyTypesB, + args=["--i", "5"], + default=ManyTypesB(i=5, s="5", f=2.0), + ) == ManyTypesB(i=5, s="5", f=2.0) diff --git a/tests/test_py311_generated/test_base_configs_nested_generated.py b/tests/test_py311_generated/test_base_configs_nested_generated.py new file mode 100644 index 000000000..9ce369220 --- /dev/null +++ b/tests/test_py311_generated/test_base_configs_nested_generated.py @@ -0,0 +1,231 @@ +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Literal, Tuple + +from torch import nn + +import tyro + + +@dataclass(frozen=True) +class AdamOptimizer: + learning_rate: float = 1e-3 + betas: Tuple[float, float] = (0.9, 0.999) + + +@dataclass(frozen=True) +class SGDOptimizer: + learning_rate: float = 1e-3 + momentum: float = 0.9 + + +@dataclass(frozen=True) +class DataConfig: + test: int = 0 + + +@dataclass(frozen=True) +class ExperimentConfig: + # Dataset to run experiment on. + dataset: Literal["mnist", "imagenet-50"] + + # Optimizer parameters. + optimizer: AdamOptimizer + + # Model size. + num_layers: int + units: int + + # Batch size. + batch_size: int + + # Total number of training steps. + train_steps: int + + # Random seed. This is helpful for making sure that our experiments are all + # reproducible! + seed: int + + # Activation to use. Not specifiable via the commandline. + activation: Callable[[], nn.Module] + + +DataConfigDataParserUnion = tyro.extras.subcommand_type_from_defaults( + { + "small-data": DataConfig( + test=2221, + ), + "big-data": DataConfig( + test=2, + ), + }, + prefix_names=False, # Omit prefixes in subcommands themselves. +) + + +# Note that we could also define this library using separate YAML files (similar to +# `config_path`/`config_name` in Hydra), but staying in Python enables seamless type +# checking + IDE support. +ExperimentConfigDataParserUnion = tyro.extras.subcommand_type_from_defaults( + { + "small": ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + "big": ExperimentConfig( + dataset="imagenet-50", + optimizer=AdamOptimizer(), + batch_size=32, + num_layers=8, + units=256, + train_steps=100_000, + seed=0, + activation=nn.GELU, + ), + }, + prefix_names=False, # Omit prefixes in subcommands themselves. +) + + +if TYPE_CHECKING: + AnnotatedDataParserUnion = DataConfig + AnnotatedExperimentParserUnion = ExperimentConfig +else: + AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ + DataConfigDataParserUnion + ] # Omit prefixes of flags in subcommands. + AnnotatedExperimentParserUnion = tyro.conf.OmitSubcommandPrefixes[ + ExperimentConfigDataParserUnion + ] # Omit prefixes of flags in subcommands. + + +@dataclass +class BaseConfig: + # The name of the experiment to run. + experiment: str + + # The path to the output directory. + output_dir: str + + # The experiment configuration. + experiment_config: AnnotatedExperimentParserUnion + + # The experiment configuration. + data_config: AnnotatedDataParserUnion = DataConfig() + + +def test_base_configs_nested() -> None: + def main(cfg: BaseConfig) -> BaseConfig: + return cfg + + assert tyro.cli( + main, + args="--cfg.experiment test --cfg.output-dir test small".split(" "), + ) == BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(0), + ) + assert tyro.cli( + main, + args="--cfg.experiment test --cfg.output-dir test small small-data".split(" "), + ) == BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(2221), + ) + assert tyro.cli( + main, + args="--cfg.experiment test --cfg.output-dir test small big-data".split(" "), + ) == BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(2), + ) + + +def test_pernicious_override(): + """From: https://github.com/nerfstudio-project/nerfstudio/issues/2789 + + Situation where we: + - have a default value in the config class + - override that default value with a subcommand annotation + - override it again with a default instance + """ + assert ( + tyro.cli( + BaseConfig, + default=BaseConfig( + "test", + "test", + ExperimentConfig( + dataset="mnist", + optimizer=AdamOptimizer(), + batch_size=2048, + num_layers=4, + units=64, + train_steps=30_000, + seed=0, + activation=nn.ReLU, + ), + DataConfig(0), + ), + args="small small-data".split(" "), + ).data_config.test + == 0 + ) + + +def test_overridable_config_helper(): + assert tyro.extras.overridable_config_cli( + { + "small-data": ( + "Small data", + DataConfig( + test=2221, + ), + ), + "big-data": ( + "Big data", + DataConfig( + test=2, + ), + ), + }, + args=["small-data", "--test", "100"], + ) == DataConfig(100) diff --git a/tests/test_py311_generated/test_boolean_optional_generated.py b/tests/test_py311_generated/test_boolean_optional_generated.py new file mode 100644 index 000000000..ba48bc9f3 --- /dev/null +++ b/tests/test_py311_generated/test_boolean_optional_generated.py @@ -0,0 +1,71 @@ +import dataclasses + +import tyro + + +def test_flag_default_false() -> None: + """Test for argparse.BooleanOptionalAction-style usage.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(False), + ) == A(True) + + assert tyro.cli( + A, + args=["--no-x"], + default=A(False), + ) == A(False) + + assert tyro.cli( + A, + args=[], + default=A(False), + ) == A(False) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + + +def test_flag_default_true() -> None: + """Test for argparse.BooleanOptionalAction-style usage.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(True), + ) == A(True) + + assert tyro.cli( + A, + args=["--no-x"], + default=A(True), + ) == A(False) + + assert tyro.cli( + A, + args=[], + default=A(True), + ) == A(True) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) diff --git a/tests/test_py311_generated/test_collections_generated.py b/tests/test_py311_generated/test_collections_generated.py new file mode 100644 index 000000000..5b6e3d811 --- /dev/null +++ b/tests/test_py311_generated/test_collections_generated.py @@ -0,0 +1,630 @@ +import collections +import collections.abc +import contextlib +import dataclasses +import enum +import io +import sys +from typing import ( + Any, + Deque, + Dict, + FrozenSet, + List, + Literal, + Optional, + Sequence, + Set, + Tuple, + Type, +) + +import pytest + +import tyro + + +def test_tuples_fixed() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, int, int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_fixed_mixed() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, str, int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, "2", 3)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_with_default() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, int, int] = (0, 1, 2) + + assert tyro.cli(A, args=[]) == A(x=(0, 1, 2)) + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + + +def test_tuple_with_literal_and_default() -> None: + @dataclasses.dataclass + class A: + x: Tuple[Literal[1, 2, 3], ...] = (1, 2) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=(1, 2)) + assert tyro.cli(A, args=["--x"]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3", "4"]) + + +def test_positional_tuple_with_literal_and_default() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.Positional[Tuple[Literal[1, 2, 3], ...]] = (1, 2) + + assert tyro.cli(A, args=["1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=(1, 2)) + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(A, args=["1", "2", "3", "4"]) + assert "invalid choice" in target.getvalue() + + +def test_tuples_fixed_multitype() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, str, float] + + assert tyro.cli(A, args=["--x", "1", "2", "3.5"]) == A(x=(1, "2", 3.5)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_fixed_bool() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, bool, bool] + + assert tyro.cli(A, args=["--x", "True", "True", "False"]) == A( + x=(True, True, False) + ) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_variable() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, ...] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=["--x"]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_variable_bool() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, ...] + + assert tyro.cli(A, args=["--x", "True", "True", "False"]) == A( + x=(True, True, False) + ) + assert tyro.cli(A, args=["--x"]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_tuples_variable_optional() -> None: + @dataclasses.dataclass + class A: + x: Optional[Tuple[int, ...]] = None + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=["--x"]) == A(x=()) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_sequences() -> None: + @dataclasses.dataclass + class A: + x: Sequence[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_sequences_narrow() -> None: + @dataclasses.dataclass + class A: + x: Sequence = dataclasses.field(default_factory=lambda: [0]) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +def test_sequences_narrow_any() -> None: + @dataclasses.dataclass + class A: + x: Sequence[Any] = dataclasses.field(default_factory=lambda: [0]) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +if sys.version_info >= (3, 9): + + def test_abc_sequences() -> None: + @dataclasses.dataclass + class A: + x: collections.abc.Sequence[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_abc_sequences_narrow() -> None: + @dataclasses.dataclass + class A: + x: collections.abc.Sequence = dataclasses.field(default_factory=lambda: [0]) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +if sys.version_info >= (3, 9): + + def test_abc_sequences_narrow_any() -> None: + @dataclasses.dataclass + class A: + x: collections.abc.Sequence[Any] = dataclasses.field( + default_factory=lambda: [0] + ) + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[0]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + + +def test_lists() -> None: + @dataclasses.dataclass + class A: + x: List[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_list_with_literal() -> None: + @dataclasses.dataclass + class A: + x: List[Literal[1, 2, 3]] + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3", "4"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_list_with_enums() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class A: + x: List[Color] + + assert tyro.cli(A, args=["--x", "RED", "RED", "BLUE"]) == A( + x=[Color.RED, Color.RED, Color.BLUE] + ) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "RED", "RED", "YELLOW"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_lists_with_default() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class A: + x: List[Color] = dataclasses.field( + default_factory=[Color.RED, Color.GREEN].copy + ) + + assert tyro.cli(A, args=[]) == A(x=[Color.RED, Color.GREEN]) + assert tyro.cli(A, args=["--x", "RED", "GREEN", "BLUE"]) == A( + x=[Color.RED, Color.GREEN, Color.BLUE] + ) + + +def test_lists_bool() -> None: + @dataclasses.dataclass + class A: + x: List[bool] + + assert tyro.cli(A, args=["--x", "True", "False", "True"]) == A( + x=[True, False, True] + ) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_sets() -> None: + @dataclasses.dataclass + class A: + x: Set[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A(x={1, 2, 3}) + assert tyro.cli(A, args=["--x"]) == A(set()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_frozen_sets() -> None: + @dataclasses.dataclass + class A: + x: FrozenSet[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A(x=frozenset({1, 2, 3})) + assert tyro.cli(A, args=["--x"]) == A(x=frozenset()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_deque() -> None: + @dataclasses.dataclass + class A: + x: Deque[int] + + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A( + x=collections.deque([1, 2, 3, 3]) + ) + assert tyro.cli(A, args=["--x"]) == A(collections.deque()) + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_sets_with_default() -> None: + @dataclasses.dataclass + class A: + x: Set[int] = dataclasses.field(default_factory={0, 1, 2}.copy) + + assert tyro.cli(A, args=[]) == A(x={0, 1, 2}) + assert tyro.cli(A, args=["--x", "1", "2", "3", "3"]) == A(x={1, 2, 3}) + assert tyro.cli(A, args=["--x"]) == A(x=set()) + + +def test_optional_sequences() -> None: + @dataclasses.dataclass + class A: + x: Optional[Sequence[int]] = None + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + assert tyro.cli(A, args=["--x", "None"]) == A(x=None) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_optional_lists() -> None: + @dataclasses.dataclass + class A: + x: Optional[List[int]] = None + + assert tyro.cli(A, args=["--x", "1", "2", "3"]) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=["--x"]) == A(x=[]) + assert tyro.cli(A, args=["--x", "None"]) == A(x=None) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_nested_optional_types() -> None: + """We support "None" as a special-case keyword. (note: this is a bit weird because + Optional[str] might interpret "None" as either a string or an actual `None` + value)""" + + @dataclasses.dataclass + class A: + x: Tuple[Optional[int], ...] + + assert tyro.cli(A, args=["--x", "0", "1"]) == A((0, 1)) + assert tyro.cli(A, args=["--x", "0", "None", "1"]) == A((0, None, 1)) + + +def test_union_over_collections() -> None: + def main(a: Tuple[float, ...] | Tuple[int, ...]) -> Any: + return a + + assert tyro.cli(main, args="--a 3.3 3.3 7.0".split(" ")) == (3.3, 3.3, 7.0) + assert tyro.cli(main, args="--a 3 3 7".split(" ")) == (3, 3, 7) + + +def test_union_over_collections_2() -> None: + def main(a: Tuple[str, float, str] | Tuple[str, str, float]) -> Any: + return a + + assert tyro.cli(main, args="--a 3.3 hey 7.0".split(" ")) == ("3.3", "hey", 7.0) + assert tyro.cli(main, args="--a 3 3 hey".split(" ")) == ("3", 3.0, "hey") + + +def test_union_over_collections_3() -> None: + def main(a: Tuple[int, int] | Tuple[int, int, int]) -> Tuple[int, ...]: + return a + + assert tyro.cli(main, args=["--a", "5", "5"]) == (5, 5) + assert tyro.cli(main, args=["--a", "1", "2", "3"]) == (1, 2, 3) + + with pytest.raises(SystemExit): + tyro.cli(main, args=["--a", "5", "5", "2", "1"]) + + with pytest.raises(SystemExit): + tyro.cli(main, args=["--a"]) + with pytest.raises(SystemExit): + tyro.cli(main, args=[]) + + +def test_choices_in_tuples_0() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, bool] + + assert tyro.cli(A, args=["--x", "True", "False"]) == A((True, False)) + + +def test_choices_in_tuples_1() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, Literal["True", "False"]] + + assert tyro.cli(A, args=["--x", "True", "False"]) == A((True, "False")) + + +def test_choices_in_tuples_2() -> None: + @dataclasses.dataclass + class A: + x: Tuple[bool, Literal["True", "False", "None"]] + + assert tyro.cli(A, args=["--x", "True", "False"]).x == (True, "False") + assert tyro.cli(A, args=["--x", "False", "None"]).x == (False, "None") + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "None", "False"]) + + +def test_nested_tuple_types() -> None: + @dataclasses.dataclass + class A: + x: Tuple[Tuple[int, int], Tuple[str, str]] + + assert tyro.cli(A, args="--x 5 5 5 5".split(" ")).x == ((5, 5), ("5", "5")) + + +def test_variable_nested_tuple() -> None: + def main(x: Tuple[Tuple[int, str], ...]) -> tuple: + return x + + assert tyro.cli(main, args="--x 1 1 2 2".split(" ")) == ((1, "1"), (2, "2")) + with pytest.raises(SystemExit): + tyro.cli(main, args="--x 1 1 2".split(" ")) + + +def test_super_nested() -> None: + def main( + x: Optional[ + List[ + Tuple[ + Optional[int], + Literal[3, 4], + Tuple[int, int] | Tuple[str, str], + ] + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_dict_no_annotation() -> None: + def main(x: Dict[str, Any] = {"int": 5, "str": "5"}): + return x + + assert tyro.cli(main, args=[]) == {"int": 5, "str": "5"} + assert tyro.cli(main, args="--x.int 3 --x.str 7".split(" ")) == { + "int": 3, + "str": "7", + } + + +def test_double_dict_no_annotation() -> None: + def main( + x: Dict[str, Any] = { + "wow": {"int": 5, "str": "5"}, + }, + ): + return x + + assert tyro.cli(main, args=[]) == {"wow": {"int": 5, "str": "5"}} + assert tyro.cli(main, args="--x.wow.int 3 --x.wow.str 7".split(" ")) == { + "wow": { + "int": 3, + "str": "7", + } + } + + +def test_list_narrowing() -> None: + def main(x: list = [0, 1, 2, "hello"]) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", 5] + + +def test_list_narrowing_any() -> None: + def main(x: List[Any] = [0, 1, 2, "hello"]) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", 5] + + +def test_list_narrowing_empty() -> None: + def main(x: list = []) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", "5"] + + +def test_list_narrowing_empty_any() -> None: + def main(x: List[Any] = []) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == ["hi", "there", "5"] + + +def test_set_narrowing() -> None: + def main(x: set = {0, 1, 2, "hello"}) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", 5} + + +def test_set_narrowing_any() -> None: + def main(x: Set[Any] = {0, 1, 2, "hello"}) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", 5} + + +def test_set_narrowing_empty() -> None: + def main(x: set = set()) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", "5"} + + +def test_set_narrowing_any_empty() -> None: + def main(x: Set[Any] = set()) -> Any: + return x + + assert tyro.cli(main, args="--x hi there 5".split(" ")) == {"hi", "there", "5"} + + +def test_tuple_narrowing() -> None: + def main(x: tuple = (0, 1, 2, "hello")) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == (0, 1, 2, "3") + + +def test_tuple_narrowing_any() -> None: + def main(x: Tuple[Any, ...] = (0, 1, 2, "hello")) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == (0, 1, 2, "3") + + +def test_tuple_narrowing_empty() -> None: + def main(x: tuple = ()) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == ("0", "1", "2", "3") + + +def test_tuple_narrowing_empty_any() -> None: + def main(x: Tuple[Any, ...] = ()) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == ("0", "1", "2", "3") + + +def test_tuple_narrowing_empty_default() -> None: + def main(x: tuple = ()) -> Any: + return x + + assert tyro.cli(main, args="--x 0 1 2 3".split(" ")) == ("0", "1", "2", "3") + + +def test_narrowing_edge_case() -> None: + """https://github.com/brentyi/tyro/issues/136""" + + @dataclasses.dataclass + class Config: + _target: Type = dataclasses.field(default_factory=lambda: MyClass) + + class MyClass: + def __len__(self): + return 0 + + assert tyro.cli(Config, args=[]) == Config() + + +def test_no_type_collections(): + assert tyro.cli(dict, args="a b c d".split(" ")) == {"a": "b", "c": "d"} + assert tyro.cli(list, args="a b c d".split(" ")) == ["a", "b", "c", "d"] + assert tyro.cli(tuple, args="a b c d".split(" ")) == ("a", "b", "c", "d") + assert tyro.cli(set, args="a b c d".split(" ")) == {"a", "b", "c", "d"} + + +def test_list_narrowing_alt() -> None: + def main(x: list = [1, "1"]) -> list: + return x + + assert tyro.cli(main, args="--x 3 four 5".split(" ")) == [3, "four", 5] + + +def test_list_narrowing_direct() -> None: + assert tyro.cli(list, default=[1, "2"], args="3 four 5".split(" ")) == [ + 3, + "four", + 5, + ] + + +def test_tuple_direct() -> None: + assert tyro.cli(Tuple[int, ...], args="1 2".split(" ")) == (1, 2) # type: ignore + assert tyro.cli(Tuple[int, int], args="1 2".split(" ")) == (1, 2) # type: ignore diff --git a/tests/test_py311_generated/test_completion_generated.py b/tests/test_py311_generated/test_completion_generated.py new file mode 100644 index 000000000..91f0e7fa9 --- /dev/null +++ b/tests/test_py311_generated/test_completion_generated.py @@ -0,0 +1,74 @@ +import contextlib +import dataclasses +import io +from typing import Annotated, Literal, Optional + +import pytest + +import tyro + + +# https://github.com/brentyi/tyro/issues/9 +@dataclasses.dataclass(frozen=True) +class Subtype: + data: int = 1 + + +@dataclasses.dataclass(frozen=True) +class TypeA: + subtype: Subtype = Subtype(1) + + +@dataclasses.dataclass(frozen=True) +class TypeB: + subtype: Subtype = Subtype(2) + + +@dataclasses.dataclass(frozen=True) +class Wrapper: + supertype: TypeA | TypeB = TypeA() + + +def test_bash(): + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(Wrapper, args=["--tyro-print-completion", "bash"]) + assert "# AUTOMATICALLY GENERATED by `shtab`" in target.getvalue() + + +def test_zsh(): + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(Wrapper, args=["--tyro-print-completion", "zsh"]) + assert "# AUTOMATICALLY GENERATED by `shtab`" in target.getvalue() + + +def test_completion_zsh(): + """https://github.com/brentyi/tyro/issues/158""" + + def start_device( + preset: Annotated[ + Optional[Literal["rgb", "depth", "ir"]], tyro.conf.arg(aliases=["-p"]) + ] = None, + frame: Annotated[ + Literal["world", "base"], tyro.conf.arg(aliases=["-f"]) + ] = "world", + ) -> None: + """ + Start device with the given preset. + + :param preset: device preset to use. + :param frame: coordinate frame to use. + """ + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + tyro.cli(start_device, args=["--tyro-print-completion", "bash"]) + + completion_script = target.getvalue() + print(completion_script) + assert "# AUTOMATICALLY GENERATED by `shtab`" in completion_script + assert "preset_choices=(" in completion_script + assert "p_choices=(" in completion_script + assert "frame_choices=(" in completion_script + assert "f_choices=(" in completion_script diff --git a/tests/test_py311_generated/test_conf_generated.py b/tests/test_py311_generated/test_conf_generated.py new file mode 100644 index 000000000..980fd79b6 --- /dev/null +++ b/tests/test_py311_generated/test_conf_generated.py @@ -0,0 +1,1468 @@ +import argparse +import contextlib +import dataclasses +import io +import json as json_ +import shlex +import sys +from typing import Annotated, Any, Dict, Generic, List, Tuple, Type, TypeVar + +import pytest +from helptext_utils import get_helptext_with_checks + +import tyro + + +def test_suppress_subcommand() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: DefaultInstanceHTTPServer| DefaultInstanceSMTPServer + bc: tyro.conf.Suppress[ + DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + ] = dataclasses.field(default_factory=DefaultInstanceHTTPServer) + + assert "bc" not in get_helptext_with_checks(DefaultInstanceSubparser) + + +def test_omit_subcommand_prefix() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: DefaultInstanceHTTPServer| DefaultInstanceSMTPServer + bc: tyro.conf.OmitSubcommandPrefixes[ + DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + ] + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=[ + "--x", + "1", + "default-instance-http-server", + "--y", + "5", + "--no-flag", + ], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "default-instance-http-server", "--y", "5"], + default=DefaultInstanceSubparser( + x=1, bc=DefaultInstanceHTTPServer(y=3, flag=False) + ), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5, flag=False)) + ) + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "default-instance-http-server", "--y", "8"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "default-instance-http-server", "--y", "8"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + +def test_avoid_subparser_with_default() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + bc: tyro.conf.AvoidSubcommands[ + DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + ] + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "5"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "--bc.y", "5"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=3)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)) + ) + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "8"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["--bc.y", "8"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + +def test_avoid_subparser_with_default_recursive() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + bc: DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "5"], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.AvoidSubcommands[DefaultInstanceSubparser], + args=["--x", "1", "--bc.y", "5"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=3)), # type: ignore + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultInstanceSubparser, + args=["bc:default-instance-smtp-server", "--bc.z", "3"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceSMTPServer(z=3)) + assert ( + tyro.cli( + tyro.conf.AvoidSubcommands[DefaultInstanceSubparser], + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "8"], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.AvoidSubcommands[DefaultInstanceSubparser], + args=["--bc.y", "8"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), # type: ignore + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + +def test_subparser_in_nested_with_metadata() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + @dataclasses.dataclass + class Nested2: + subcommand: ( + Annotated[A, tyro.conf.subcommand("command-a", default=A(7))] + | Annotated[B, tyro.conf.subcommand("command-b", default=B(9))] + ) + + @dataclasses.dataclass + class Nested1: + nested2: Nested2 + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_suppressed() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: tyro.conf.Suppress[int] + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + @dataclasses.dataclass + class Nested2: + subcommand: ( + Annotated[A, tyro.conf.subcommand("command-a", default=A(7))] + | Annotated[B, tyro.conf.subcommand("command-b", default=B(9))] + ) + + @dataclasses.dataclass + class Nested1: + nested2: Nested2 + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_generic() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + T = TypeVar("T") + + @dataclasses.dataclass + class Nested2(Generic[T]): + subcommand: T + + @dataclasses.dataclass + class Nested1: + nested2: Nested2[ + Annotated[A, tyro.conf.subcommand("command-a", default=A(7))] + | Annotated[B, tyro.conf.subcommand("command-b", default=B(9))] + ] + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_generic_alt() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + T = TypeVar("T") + + @dataclasses.dataclass + class Nested2(Generic[T]): + subcommand: ( + Annotated[T, tyro.conf.subcommand("command-a", default=A(7))] + | Annotated[B, tyro.conf.subcommand("command-b", default=B(9))] + ) + + @dataclasses.dataclass + class Nested1: + nested2: Nested2[A] + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-a".split(" "), + ) == Parent(Nested1(Nested2(A(7)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-a --nested1.nested2.subcommand.a 3".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(A(3)))) + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:command-b".split(" "), + ) == Parent(Nested1(Nested2(B(9)))) + assert tyro.cli( + Parent, + args=( + "nested1.nested2.subcommand:command-b --nested1.nested2.subcommand.b 7".split( + " " + ) + ), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_subparser_in_nested_with_metadata_default_matching() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + a: A = A(5) + + default_one = B(3) + default_three = B(9) + + @dataclasses.dataclass + class Nested: + subcommand: ( + Annotated[B, tyro.conf.subcommand("one", default=default_one)] + | Annotated[B, tyro.conf.subcommand("two")] + | Annotated[B, tyro.conf.subcommand("three", default=default_three)] + ) + + # Match by hash. + def main_one(x: Nested = Nested(default_one)) -> None: + pass + + assert "default: x.subcommand:one" in get_helptext_with_checks(main_one) + + # Match by value. + def main_two(x: Nested = Nested(B(9))) -> None: + pass + + assert "default: x.subcommand:three" in get_helptext_with_checks(main_two) + + # Match by type. + def main_three(x: Nested = Nested(B(15))) -> None: + pass + + assert "default: x.subcommand:one" in get_helptext_with_checks(main_three) + + +def test_flag() -> None: + """When boolean flags have no default value, they must be explicitly specified.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(False), + ) == A(True) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + assert tyro.cli( + A, + args=["--x", "True"], + default=A(False), + config=(tyro.conf.FlagConversionOff,), + ) == A(True) + + +def test_fixed() -> None: + """When an argument is fixed, we shouldn't be able to override it from the CLI.""" + + @dataclasses.dataclass + class A: + x: tyro.conf.Fixed[bool] + + assert tyro.cli( + A, + args=[], + default=A(True), + ) == A(True) + + with pytest.raises(SystemExit): + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + assert tyro.cli( + tyro.conf.FlagConversionOff[A], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + + with pytest.raises(SystemExit): + assert tyro.cli( + A, + args=["--x", "True"], + default=A(False), # type: ignore + config=(tyro.conf.FlagConversionOff,), + ) == A(True) + + +def test_fixed_recursive() -> None: + """When an argument is fixed, we shouldn't be able to override it from the CLI.""" + + @dataclasses.dataclass + class A: + x: bool + + assert tyro.cli( + A, + args=["--x"], + default=A(False), + ) == A(True) + + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + with pytest.raises(SystemExit): + assert tyro.cli( + tyro.conf.Fixed[tyro.conf.FlagConversionOff[A]], + args=["--x", "True"], + default=A(False), # type: ignore + ) == A(True) + + +def test_suppressed_group() -> None: + """Reproduction of https://github.com/nerfstudio-project/nerfstudio/issues/882.""" + + @dataclasses.dataclass + class Inner: + a: int + b: int + + def main( + value: int, + inner: tyro.conf.Suppress[Inner] = Inner(1, 2), + ) -> int: + return value + inner.a + inner.b + + assert tyro.cli(main, args=["--value", "5"]) == 8 + + +def test_fixed_group() -> None: + """Inspired by https://github.com/nerfstudio-project/nerfstudio/issues/882.""" + + @dataclasses.dataclass + class Inner: + a: int + b: int + + def main( + value: int, + inner: tyro.conf.Fixed[Inner] = Inner(1, 2), + ) -> int: + return value + inner.a + inner.b + + assert tyro.cli(main, args=["--value", "5"]) == 8 + + +def test_fixed_suppressed_group() -> None: + """Reproduction of https://github.com/nerfstudio-project/nerfstudio/issues/882.""" + + @dataclasses.dataclass + class Inner: + a: int + b: int + + def main( + value: int, + inner: tyro.conf.Fixed[Inner] = Inner(1, 2), + ) -> int: + return value + inner.a + inner.b + + assert tyro.cli(main, args=["--value", "5"]) == 8 + + +def test_suppressed() -> None: + @dataclasses.dataclass + class Struct: + a: int = 5 + b: tyro.conf.Suppress[str] = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" in helptext + assert "--x.b" not in helptext + + +def test_suppress_manual_fixed() -> None: + @dataclasses.dataclass + class Struct: + a: int = 5 + b: tyro.conf.SuppressFixed[tyro.conf.Fixed[str]] = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" in helptext + assert "--x.b" not in helptext + + +def test_suppress_manual_fixed_one_arg_only() -> None: + @dataclasses.dataclass + class Struct: + b: tyro.conf.SuppressFixed[tyro.conf.Fixed[str]] = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" not in helptext + assert "--x.b" not in helptext + + +def test_suppress_auto_fixed() -> None: + @dataclasses.dataclass + class Struct: + a: int = 5 + + def b(self, x): + return 5 + + def main(x: tyro.conf.SuppressFixed[Any] = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x.a" in helptext + assert "--x.b" not in helptext + + +def test_argconf_help() -> None: + @dataclasses.dataclass + class Struct: + a: Annotated[ + int, tyro.conf.arg(name="nice", help="Hello world", metavar="NUMBER") + ] = 5 + b: tyro.conf.Suppress[str] = "7" + + def main(x: Any = Struct()) -> int: + return x.a + + helptext = get_helptext_with_checks(main) + assert "Hello world" in helptext + assert "INT" not in helptext + assert "NUMBER" in helptext + assert "--x.a" not in helptext + assert "--x.nice" in helptext + assert "--x.b" not in helptext + + assert tyro.cli(main, args=[]) == 5 + assert tyro.cli(main, args=["--x.nice", "3"]) == 3 + + +def test_argconf_no_prefix_help() -> None: + @dataclasses.dataclass + class Struct: + a: Annotated[ + int, + tyro.conf.arg( + name="nice", help="Hello world", metavar="NUMBER", prefix_name=False + ), + ] = 5 + b: tyro.conf.Suppress[str] = "7" + + def main(x: Any = Struct()) -> int: + return x.a + + helptext = get_helptext_with_checks(main) + assert "Hello world" in helptext + assert "INT" not in helptext + assert "NUMBER" in helptext + assert "--x.a" not in helptext + assert "--x.nice" not in helptext + assert "--nice" in helptext + assert "--x.b" not in helptext + + assert tyro.cli(main, args=[]) == 5 + with pytest.raises(SystemExit): + assert tyro.cli(main, args=["--x.nice", "3"]) == 3 + assert tyro.cli(main, args=["--nice", "3"]) == 3 + + +def test_positional() -> None: + def main(x: tyro.conf.Positional[int], y: int) -> int: + return x + y + + assert tyro.cli(main, args="5 --y 3".split(" ")) == 8 + assert tyro.cli(main, args="--y 3 5".split(" ")) == 8 + + +def test_positional_required_args() -> None: + @dataclasses.dataclass + class Args: + x: int + y: int = 3 + + assert tyro.cli( + tyro.conf.PositionalRequiredArgs[Args], args="5 --y 3".split(" ") + ) == Args(5, 3) + assert tyro.cli( + tyro.conf.PositionalRequiredArgs[Args], args="--y 3 5".split(" ") + ) == Args(5, 3) + + +def test_positional_order_swap() -> None: + def main(x: int, y: tyro.conf.Positional[int]) -> int: + return x + y + + assert tyro.cli(main, args="5 --x 3".split(" ")) == 8 + assert tyro.cli(main, args="--x 3 5".split(" ")) == 8 + + +def test_omit_subcommand_prefix_and_consolidate_subcommand_args() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: DefaultInstanceHTTPServer| DefaultInstanceSMTPServer + bc: tyro.conf.OmitSubcommandPrefixes[ + DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + ] = dataclasses.field(default_factory=DefaultInstanceHTTPServer) + + assert ( + tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "5", + "--no-flag", + ], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "5", + ], + default=DefaultInstanceSubparser( # type: ignore + x=1, bc=DefaultInstanceHTTPServer(y=3, flag=False) + ), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5, flag=False)) + ) + assert ( + tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "8", + ], + ) + # Type ignore can be removed once TypeForm lands. + # https://discuss.python.org/t/typeform-spelling-for-a-type-annotation-object-at-runtime/51435 + == tyro.cli( + tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], + args=[ + "default-instance-http-server", + "--x", + "1", + "--y", + "8", + ], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=7)), # type: ignore + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + # Despite all defaults being set, a subcommand should be required. + with pytest.raises(SystemExit): + tyro.cli(tyro.conf.ConsolidateSubcommandArgs[DefaultInstanceSubparser], args=[]) + + +def test_omit_subcommand_prefix_and_consolidate_subcommand_args_in_function() -> None: + @tyro.conf.configure(tyro.conf.subcommand(name="http-server")) + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + flag: bool = True + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + # bc: DefaultInstanceHTTPServer| DefaultInstanceSMTPServer + bc: DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + + @tyro.conf.configure( + tyro.conf.OmitSubcommandPrefixes, + tyro.conf.ConsolidateSubcommandArgs, + ) + def func(parent: DefaultInstanceSubparser) -> DefaultInstanceSubparser: + return parent + + assert tyro.cli( + func, + args=[ + "http-server", + "--parent.x", + "1", + # --y and --no-flag are in a subcommand with prefix omission. + "--y", + "5", + "--no-flag", + ], + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5, flag=False)) + assert tyro.cli( + func, + args=[ + "http-server", + "--parent.x", + "1", + # --y is in a subcommand with prefix omission. + "--y", + "8", + ], + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + + +def test_append_lists() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[List[int]] + + assert tyro.cli(A, args=[]) == A(x=[]) + assert tyro.cli(A, args="--x 1 --x 2 --x 3".split(" ")) == A(x=[1, 2, 3]) + assert tyro.cli(A, args=[]) == A(x=[]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_tuple() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[int, ...]] + + assert tyro.cli(A, args=[]) == A(x=()) + assert tyro.cli(A, args="--x 1 --x 2 --x 3".split(" ")) == A(x=(1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_tuple_with_default() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[int, ...]] = (1, 2) + + assert tyro.cli(A, args="--x 1 --x 2 --x 3".split(" ")) == A(x=(1, 2, 1, 2, 3)) + assert tyro.cli(A, args=[]) == A(x=(1, 2)) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_nested_tuple_fixed_length() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, int], ...]] + + assert tyro.cli(A, args="--x 1 1 --x 2 2 --x 3 3".split(" ")) == A( + x=(("1", 1), ("2", 2), ("3", 3)) + ) + assert tyro.cli(A, args=[]) == A(x=()) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_nested_tuple_with_default_fixed_length() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, int], ...]] = (("1", 1), ("2", 2)) + + assert tyro.cli(A, args="--x 1 1 --x 2 2 --x 3 3".split(" ")) == A( + x=(("1", 1), ("2", 2), ("1", 1), ("2", 2), ("3", 3)) + ) + assert tyro.cli(A, args=[]) == A(x=(("1", 1), ("2", 2))) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_dict() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Dict[str, int]] + + assert tyro.cli(A, args="--x 1 1 --x 2 2 --x 3 3".split(" ")) == A( + x={"1": 1, "2": 2, "3": 3} + ) + assert tyro.cli(A, args=[]) == A(x={}) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1", "2", "3"]) + + +def test_append_dict_with_default() -> None: + """Append has no impact when a dictionary has a default value.""" + + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Dict[str, int]] = dataclasses.field( + default_factory=lambda: {"1": 1} + ) + + assert tyro.cli(A, args=[]) == A(x={"1": 1}) + assert tyro.cli(A, args=["--x.1", "2"]) == A(x={"1": 2}) + with pytest.raises(SystemExit): + assert tyro.cli(A, args="--x 2 2 --x 3 3".split(" ")) == A( + x={"1": 1, "2": 2, "3": 3} + ) + + +def test_append_nested_tuple() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, ...], ...]] + + assert tyro.cli(A, args="--x 1 2 3 --x 4 5".split(" ")) == A( + x=(("1", "2", "3"), ("4", "5")) + ) + assert tyro.cli(A, args=[]) == A(x=()) + + +def test_append_nested_tuple_with_default() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Tuple[Tuple[str, ...], ...]] = (("1", "2"),) + + assert tyro.cli(A, args="--x 1 2 3 --x 4 5".split(" ")) == A( + x=(("1", "2"), ("1", "2", "3"), ("4", "5")) + ) + assert tyro.cli(A, args=[]) == A(x=(("1", "2"),)) + + +def test_append_nested_list() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[List[List[int]]] + + assert tyro.cli(A, args="--x 1 2 3 --x 4 5".split(" ")) == A(x=[[1, 2, 3], [4, 5]]) + assert tyro.cli(A, args=[]) == A(x=[]) + + +def test_append_nested_dict() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[List[Dict[str, int]]] + + assert tyro.cli(A, args="--x 1 2 3 4 --x 4 5".split(" ")) == A( + x=[{"1": 2, "3": 4}, {"4": 5}] + ) + assert tyro.cli(A, args=[]) == A(x=[]) + + +def test_append_nested_dict_double() -> None: + @dataclasses.dataclass + class A: + x: tyro.conf.UseAppendAction[Dict[str, Dict[str, int]]] + + assert tyro.cli(A, args="--x 0 1 2 3 4 --x 4 5 6".split(" ")) == A( + x={"0": {"1": 2, "3": 4}, "4": {"5": 6}} + ) + assert tyro.cli(A, args=[]) == A(x={}) + + +def test_duplicated_arg() -> None: + # Loosely inspired by: https://github.com/brentyi/tyro/issues/49 + @dataclasses.dataclass + class ModelConfig: + num_slots: Annotated[int, tyro.conf.arg(prefix_name=False)] + + @dataclasses.dataclass + class TrainConfig: + num_slots: int + model: ModelConfig + + with pytest.raises(argparse.ArgumentError): + tyro.cli(TrainConfig, args="--num-slots 3".split(" ")) + + +def test_omit_arg_prefixes() -> None: + # Loosely inspired by: https://github.com/brentyi/tyro/issues/49 + @dataclasses.dataclass + class ModelConfig: + num_slots: int + + @dataclasses.dataclass + class TrainConfig: + model: ModelConfig + + assert tyro.cli( + tyro.conf.OmitSubcommandPrefixes[TrainConfig], + args="--model.num-slots 3".split(" "), + ) == TrainConfig(ModelConfig(num_slots=3)) + + assert tyro.cli( + tyro.conf.OmitArgPrefixes[TrainConfig], args="--num-slots 3".split(" ") + ) == TrainConfig(ModelConfig(num_slots=3)) + + +def test_custom_constructor_0() -> None: + def times_two(n: str) -> int: + return int(n) * 2 + + @dataclasses.dataclass + class Config: + x: Annotated[int, tyro.conf.arg(name="x-renamed", constructor=times_two)] + + assert tyro.cli(Config, args="--x-renamed.n 5".split(" ")) == Config(x=10) + + +def test_custom_constructor_1() -> None: + def times_two(n: int) -> int: + return int(n) * 2 + + @dataclasses.dataclass + class Config: + x: Annotated[int, tyro.conf.arg(constructor=times_two)] + + assert tyro.cli(Config, args="--x.n 5".split(" ")) == Config(x=10) + + +def test_custom_constructor_2() -> None: + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=int)] + + assert tyro.cli(Config, args="--x 5".split(" ")) == Config(x=5) + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x 5.23".split(" ")) + + +def test_custom_constructor_3() -> None: + def dict_from_json(json: str) -> dict: + out = json_.loads(json) + if not isinstance(out, dict): + raise ValueError(f"{json} is not a dict!") + return out + + @dataclasses.dataclass + class Config: + x: Annotated[ + dict, + tyro.conf.arg( + metavar="JSON", + constructor=dict_from_json, + ), + ] + + assert tyro.cli( + Config, args=shlex.split('--x.json \'{"hello": "world"}\'') + ) == Config(x={"hello": "world"}) + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.json 5".split(" ")) + + error = target.getvalue() + assert "Error parsing x: 5 is not a dict!" in error + + +def test_custom_constructor_4() -> None: + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=int)] = 3.23 + + assert tyro.cli(Config, args="--x 5".split(" ")) == Config(x=5) + assert tyro.cli(Config, args=[]) == Config(x=3.23) + + +def test_custom_constructor_5() -> None: + def make_float(a: float, b: float, c: float = 3) -> float: + return a * b * c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli(Config, args="--x.a 5 --x.b 2 --x.c 3".split(" ")) == Config(x=30) + assert tyro.cli(Config, args="--x.a 5 --x.b 2".split(" ")) == Config(x=30) + + # --x.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.a 5".split(" ")) + + # --x.a and --x.b are required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.c 5".split(" ")) + + +def test_custom_constructor_6() -> None: + def make_float( + a: tyro.conf.Positional[float], + b2: Annotated[float, tyro.conf.arg(name="b")], + c: float = 3, + ) -> float: + return a * b2 * c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli(Config, args="--x.b 2 --x.c 3 5".split(" ")) == Config(x=30) + assert tyro.cli(Config, args="--x.b 2 5".split(" ")) == Config(x=30) + + # --x.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="5".split(" ")) + + # --x.a and --x.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.c 5".split(" ")) + error = target.getvalue() + assert "either all arguments must be provided" in error + assert "or none of them" in error + assert "We're missing arguments" in error + + +def test_custom_constructor_7() -> None: + @dataclasses.dataclass + class Struct: + a: int + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli( + Config, args="--x.struct.a 5 --x.struct.b 2 --x.struct.c 3".split(" ") + ) == Config(x=30) + assert tyro.cli(Config, args="--x.struct.a 5 --x.struct.b 2".split(" ")) == Config( + x=30 + ) + + # --x.struct.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.struct.a 5".split(" ")) + + # --x.struct.a and --x.struct.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.struct.c 5".split(" ")) + error = target.getvalue() + assert "We're missing arguments" in error + assert "'--x.struct.b'" in error + assert "'--x.struct.a'" in error # The 5 is parsed into `a`. + + +def test_custom_constructor_8() -> None: + @dataclasses.dataclass + class Struct: + a: tyro.conf.Positional[int] + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 5".split(" ") + ) == Config(x=30) + assert tyro.cli(Config, args="--x.struct.b 2 5".split(" ")) == Config(x=30) + + # --x.struct.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="5".split(" ")) + + # --x.struct.a and --x.struct.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.struct.b 5".split(" ")) + error = target.getvalue() + assert "We're missing arguments" in error + assert "'x.struct.a'" in error + assert "'--x.struct.b'" not in error + + +def test_custom_constructor_9() -> None: + def commit(branch: str) -> int: + """Commit""" + print(f"commit branch={branch}") + return 3 + + assert ( + tyro.cli( + Annotated[Any, tyro.conf.arg(constructor=commit)], # type: ignore + args="--branch 5".split(" "), + ) + == 3 + ) + + +def test_custom_constructor_10() -> None: + def commit(branch: str) -> int: + """Commit""" + print(f"commit branch={branch}") + return 3 + + def inner(x: Annotated[Any, tyro.conf.arg(constructor=commit)]) -> None: + return x + + def inner_no_prefix( + x: Annotated[Any, tyro.conf.arg(constructor=commit, prefix_name=False)], + ) -> None: + return x + + def outer(x: Annotated[Any, tyro.conf.arg(constructor=inner)]) -> None: + return x + + def outer_no_prefix( + x: Annotated[Any, tyro.conf.arg(constructor=inner_no_prefix)], + ) -> None: + return x + + assert ( + tyro.cli( + outer, + args="--x.x.branch 5".split(" "), + ) + == 3 + ) + assert ( + tyro.cli( + outer_no_prefix, + args="--x.branch 5".split(" "), + ) + == 3 + ) + + +def test_alias() -> None: + """Arguments with aliases.""" + + @dataclasses.dataclass + class Struct: + a: Annotated[int, tyro.conf.arg(aliases=["--all", "-d"])] + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + assert tyro.cli(Config, args=[]) == Config(x=3.23) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 --x.struct.a 5".split(" ") + ) == Config(x=30) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 -d 5".split(" ") + ) == Config(x=30) + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 --all 5".split(" ") + ) == Config(x=30) + assert tyro.cli(Config, args="--x.struct.b 2 --x.struct.a 5".split(" ")) == Config( + x=30 + ) + + # --x.struct.b is required! + with pytest.raises(SystemExit): + tyro.cli(Config, args="--x.struct.a 5".split(" ")) + + # --x.struct.a and --x.struct.b are required! + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Config, args="--x.struct.b 5".split(" ")) + error = target.getvalue() + assert "We're missing arguments" in error + assert "'--x.struct.a'" in error + assert "'--x.struct.b'" not in error + + assert "--x.struct.a INT, --all INT, -d INT" in get_helptext_with_checks(Config) + + +def test_positional_alias() -> None: + """Positional arguments with aliases (which will be ignored).""" + + @dataclasses.dataclass + class Struct: + a: Annotated[tyro.conf.Positional[int], tyro.conf.arg(aliases=["--all", "-d"])] + b: int + c: int = 3 + + def make_float(struct: Struct) -> float: + return struct.a * struct.b * struct.c + + @dataclasses.dataclass + class Config: + x: Annotated[float, tyro.conf.arg(constructor=make_float)] = 3.23 + + with pytest.warns(UserWarning): + assert tyro.cli(Config, args=[]) == Config(x=3.23) + with pytest.warns(UserWarning): + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 5".split(" ") + ) == Config(x=30) + + with pytest.raises(SystemExit), pytest.warns(UserWarning): + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 -d 5".split(" ") + ) == Config(x=30) + with pytest.raises(SystemExit), pytest.warns(UserWarning): + assert tyro.cli( + Config, args="--x.struct.b 2 --x.struct.c 3 --all 5".split(" ") + ) == Config(x=30) + + +def test_flag_alias() -> None: + @dataclasses.dataclass + class Struct: + flag: Annotated[bool, tyro.conf.arg(aliases=["-f", "--flg"])] = False + + assert tyro.cli(Struct, args=[]).flag is False + assert tyro.cli(Struct, args="--flag".split(" ")).flag is True + assert tyro.cli(Struct, args="--no-flag".split(" ")).flag is False + assert tyro.cli(Struct, args="--flg".split(" ")).flag is True + assert tyro.cli(Struct, args="--no-flg".split(" ")).flag is False + assert tyro.cli(Struct, args="-f".split(" ")).flag is True + + # BooleanOptionalAction will ignore arguments that aren't prefixed with --. + with pytest.raises(SystemExit): + tyro.cli(Struct, args="-no-f".split(" ")) + + +def test_subcommand_constructor_mix() -> None: + """https://github.com/brentyi/tyro/issues/89""" + + def checkout(branch: str) -> str: + """Check out a branch.""" + return branch + + def commit(message: str, all: bool = False) -> str: + """Make a commit.""" + return f"{message} {all}" + + @dataclasses.dataclass + class Arg: + foo: int = 1 + + t: Any = Annotated[ + Annotated[ + Any, + tyro.conf.subcommand(name="checkout-renamed", constructor=checkout), + ] + | Annotated[ + Any, + tyro.conf.subcommand(name="commit", constructor=commit), + tyro.conf.FlagConversionOff, + ] + | Arg, + tyro.conf.OmitArgPrefixes, # Should do nothing. + ] + + assert tyro.cli(t, args=["arg"]) == Arg() + assert tyro.cli(t, args=["checkout-renamed", "--branch", "main"]) == "main" + assert tyro.cli(t, args=["commit", "--message", "hi", "--all", "True"]) == "hi True" + + +def test_merge() -> None: + """Test effect of tyro.conf.arg() on nested structures by approximating an + HfArgumentParser-style API.""" + + T = TypeVar("T") + + # We could add a lot overloads here if we were doing this for real. :) + def instantiate_dataclasses( + classes: Tuple[Type[T], ...], args: List[str] + ) -> Tuple[T, ...]: + return tyro.cli( + tyro.conf.OmitArgPrefixes[ # type: ignore + # Convert (type1, type2) into Tuple[type1, type2] + Tuple.__getitem__( # type: ignore + tuple(Annotated[c, tyro.conf.arg(name=c.__name__)] for c in classes) + ) + ], + args=args, + ) + + @dataclasses.dataclass(frozen=True) + class OptimizerConfig: + lr: float = 1e-4 + weight: int = 10 + + @dataclasses.dataclass(frozen=True) + class DatasetConfig: + batch_size: int = 1 + shuffle: bool = False + + assert instantiate_dataclasses( + (OptimizerConfig, DatasetConfig), args=["--lr", "1e-3"] + ) == ( + OptimizerConfig(1e-3), + DatasetConfig(), + ) + assert instantiate_dataclasses( + (OptimizerConfig, DatasetConfig), args=["--lr", "1e-3", "--shuffle"] + ) == ( + OptimizerConfig(1e-3), + DatasetConfig(shuffle=True), + ) + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stdout(target): + instantiate_dataclasses((OptimizerConfig, DatasetConfig), args=["--help"]) + helptext = target.getvalue() + assert "OptimizerConfig options" in helptext + assert "DatasetConfig options" in helptext + + +def test_counter_action() -> None: + def main( + verbosity: tyro.conf.UseCounterAction[int], + aliased_verbosity: Annotated[ + tyro.conf.UseCounterAction[int], tyro.conf.arg(aliases=["-v"]) + ], + ) -> Tuple[int, int]: + """Example showing how to use counter actions. + Args: + verbosity: Verbosity level. + aliased_verbosity: Same as above, but can also be specified with -v, -vv, -vvv, etc. + """ + return verbosity, aliased_verbosity + + assert tyro.cli(main, args=[]) == (0, 0) + assert tyro.cli(main, args="--verbosity --verbosity".split(" ")) == (2, 0) + assert tyro.cli(main, args="--verbosity --verbosity -v".split(" ")) == (2, 1) + if sys.version_info >= (3, 8): + # Doesn't work in Python 3.7 because of argparse limitations. + assert tyro.cli(main, args="--verbosity --verbosity -vv".split(" ")) == (2, 2) + assert tyro.cli(main, args="--verbosity --verbosity -vvv".split(" ")) == (2, 3) diff --git a/tests/test_py311_generated/test_dcargs_generated.py b/tests/test_py311_generated/test_dcargs_generated.py new file mode 100644 index 000000000..b8836572b --- /dev/null +++ b/tests/test_py311_generated/test_dcargs_generated.py @@ -0,0 +1,879 @@ +import argparse +import copy +import dataclasses +import datetime +import enum +import os +import pathlib +from typing import ( + Annotated, + Any, + AnyStr, + Callable, + ClassVar, + Dict, + Final, + List, + Literal, + Optional, + Tuple, + TypeAlias, + TypeVar, +) + +import pytest +import torch + +import tyro + + +def test_no_args() -> None: + def main() -> int: + return 5 + + assert tyro.cli(main, args=[]) == 5 + with pytest.raises(SystemExit): + tyro.cli(main, args=["3"]) + + +def test_basic() -> None: + @dataclasses.dataclass + class ManyTypes: + i: int + s: str + f: float + p: pathlib.Path + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypes, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + # We can directly pass a function to `tyro.cli()`: + def function(i: int, s: str, f: float, p: pathlib.Path) -> ManyTypes: + return ManyTypes(i=i, s=s, f=f, p=p) + + assert tyro.cli( + function, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + # We can directly pass a generic class to `tyro.cli()`: + class Wrapper: + def __init__(self, i: int, s: str, f: float, p: pathlib.Path): + self.inner = ManyTypes(i=i, s=s, f=f, p=p) + + assert tyro.cli( + Wrapper, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ).inner == ManyTypes(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + +def test_init_false() -> None: + @dataclasses.dataclass + class InitFalseDataclass: + i: int + s: str + f: float + dir: pathlib.Path + ignored: str = dataclasses.field(default="hello", init=False) + + assert tyro.cli( + InitFalseDataclass, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--dir", + "~", + ], + ) == InitFalseDataclass(i=5, s="5", f=5.0, dir=pathlib.Path("~")) + + with pytest.raises(SystemExit): + tyro.cli( + InitFalseDataclass, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--dir", + "~", + "--ignored", + "blah", + ], + ) + + +def test_required() -> None: + @dataclasses.dataclass + class A: + x: int + + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + +def test_flag() -> None: + """When boolean flags have no default value, they must be explicitly specified.""" + + @dataclasses.dataclass + class A: + x: bool + + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "true"]) + assert tyro.cli(A, args=["--x", "True"]) == A(True) + + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "0"]) + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "false"]) + assert tyro.cli(A, args=["--x", "False"]) == A(False) + + +def test_flag_default_false() -> None: + """When boolean flags default to False, a --flag-name flag must be passed in to flip it to True.""" + + @dataclasses.dataclass + class A: + x: bool = False + + assert tyro.cli(A, args=[]) == A(False) + assert tyro.cli(A, args=["--x"]) == A(True) + + +def test_flag_default_true() -> None: + """When boolean flags default to True, a --no-flag-name flag must be passed in to flip it to False.""" + + @dataclasses.dataclass + class A: + x: bool = True + + assert tyro.cli(A, args=[]) == A(True) + assert tyro.cli(A, args=["--no-x"]) == A(False) + + +def test_flag_default_true_nested() -> None: + """When boolean flags default to True, a --no-flag-name flag must be passed in to flip it to False.""" + + @dataclasses.dataclass + class NestedDefaultTrue: + x: bool = True + + @dataclasses.dataclass + class A: + x: NestedDefaultTrue + + assert tyro.cli(A, args=[]) == A(NestedDefaultTrue(True)) + assert tyro.cli(A, args=["--x.no-x"]) == A(NestedDefaultTrue(False)) + + +def test_default() -> None: + @dataclasses.dataclass + class A: + x: int = 5 + + assert tyro.cli(A, args=[]) == A() + + +def test_default_factory() -> None: + @dataclasses.dataclass + class A: + x: int = dataclasses.field(default_factory=lambda: 5) + + assert tyro.cli(A, args=[]) == A() + + +def test_optional() -> None: + @dataclasses.dataclass + class A: + x: Optional[int] = None + + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_union_basic() -> None: + def main(x: int | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_union_with_list() -> None: + def main(x: int | str | List[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + assert tyro.cli(main, args=["--x", "True"]) == "True" + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_union_literal() -> None: + def main(x: Literal[1, 2] | Literal[3, 4, 5] | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == "6" + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_func_typevar() -> None: + T = TypeVar("T", int, str) + + def main(x: T) -> T: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_func_typevar_bound() -> None: + T = TypeVar("T", bound=int) + + def main(x: T) -> T: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "five"]) + + +def test_enum() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class EnumClassA: + color: Color + + @dataclasses.dataclass + class EnumClassB: + color: Color = Color.GREEN + + assert tyro.cli(EnumClassA, args=["--color", "RED"]) == EnumClassA(color=Color.RED) + assert tyro.cli(EnumClassB, args=[]) == EnumClassB() + + +def test_enum_alias() -> None: + class Color(enum.Enum): + RED = 1 + ROUGE = 1 + + @dataclasses.dataclass + class A: + color: Color + + assert tyro.cli(A, args=["--color", "RED"]) == A(color=Color.RED) + assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.ROUGE) + assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.RED) + + +def test_literal() -> None: + @dataclasses.dataclass + class A: + x: Literal[0, 1, 2] + + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) + with pytest.raises(SystemExit): + assert tyro.cli(A, args=["--x", "3"]) + + +def test_literal_none() -> None: + @dataclasses.dataclass + class A: + x: Literal[0, 1, None, 2] + + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) + assert tyro.cli(A, args=["--x", "None"]) == A(x=None) + with pytest.raises(SystemExit): + assert tyro.cli(A, args=["--x", "3"]) + + +Choices = int +Choices = tyro.extras.literal_type_from_choices([0, 1, 2]) # type: ignore + + +def test_dynamic_literal() -> None: + @dataclasses.dataclass + class A: + x: Choices # type: ignore + + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) + with pytest.raises(SystemExit): + assert tyro.cli(A, args=["--x", "3"]) + + +def test_literal_bool() -> None: + def main(x: Literal[True]) -> bool: + return x + + assert tyro.cli(main, args=["--x", "True"]) is True + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "False"]) + + def main2(x: Literal[True, False]) -> bool: + return x + + assert tyro.cli(main2, args=["--x", "True"]) is True + assert tyro.cli(main2, args=["--x", "False"]) is False + with pytest.raises(SystemExit): + tyro.cli(main2, args=["--x", "Tru"]) + + +def test_literal_enum() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class A: + x: Literal[Color.RED, Color.GREEN] + + assert tyro.cli(A, args=["--x", "RED"]) == A(x=Color.RED) + assert tyro.cli(A, args=["--x", "GREEN"]) == A(x=Color.GREEN) + with pytest.raises(SystemExit): + assert tyro.cli(A, args=["--x", "BLUE"]) + + +def test_optional_literal() -> None: + @dataclasses.dataclass + class A: + x: Optional[Literal[0, 1, 2]] = None + + assert tyro.cli(A, args=["--x", "1"]) == A(x=1) + with pytest.raises(SystemExit): + assert tyro.cli(A, args=["--x", "3"]) + assert tyro.cli(A, args=[]) == A(x=None) + + +def test_multitype_literal() -> None: + def main(x: Literal[0, "5"]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "0"]) == 0 + assert tyro.cli(main, args=["--x", "5"]) == "5" + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "6"]) + + +def test_annotated() -> None: + """Annotated[] is a no-op.""" + + @dataclasses.dataclass + class A: + x: Annotated[int, "some label"] = 3 + + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_annotated_optional() -> None: + """Annotated[] is a no-op.""" + + @dataclasses.dataclass + class A: + x: Annotated[Optional[int], "some label"] = 3 + + assert tyro.cli(A, args=[]) == A(x=3) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_optional_annotated() -> None: + """Annotated[] is a no-op.""" + + @dataclasses.dataclass + class A: + x: Optional[Annotated[int, "some label"]] = 3 + + assert tyro.cli(A, args=[]) == A(x=3) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_final() -> None: + """Final[] is a no-op.""" + + @dataclasses.dataclass + class A: + x: Final[int] = 3 + + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_final_optional() -> None: + @dataclasses.dataclass + class A: + x: Final[Optional[int]] = 3 + + assert tyro.cli(A, args=[]) == A(x=3) + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_classvar() -> None: + """ClassVar[] types should be skipped.""" + + @dataclasses.dataclass + class A: + x: ClassVar[int] = 5 + + with pytest.raises(SystemExit): + tyro.cli(A, args=["--x", "1"]) + assert tyro.cli(A, args=[]) == A() + + +def test_parse_empty_description() -> None: + """If the file has no dosctring, it should be treated as an empty string.""" + + @dataclasses.dataclass + class A: + x: int = 0 + + assert tyro.cli(A, description=None, args=[]) == A(x=0) + + +SomeTypeAlias: TypeAlias = int + + +def test_type_alias() -> None: + def add(a: SomeTypeAlias, b: SomeTypeAlias) -> SomeTypeAlias: + return a + b + + assert tyro.cli(add, args=["--a", "5", "--b", "7"]) == 12 + + +def test_any() -> None: + def main(x: Any = 5) -> Any: + return x + + assert tyro.cli(main, args=[]) == 5 + + +def test_bytes() -> None: + def main(x: bytes) -> bytes: + return x + + assert tyro.cli(main, args=["--x", "hello"]) == b"hello" + + +def test_any_str() -> None: + def main(x: AnyStr) -> AnyStr: + return x + + # Use bytes when provided ascii-compatible inputs. + assert tyro.cli(main, args=["--x", "hello"]) == b"hello" + assert tyro.cli(main, args=["--x", "hello„"]) == "hello„" + + +def test_fixed() -> None: + def main(x: Callable[[int], int] = lambda x: x * 2) -> Callable[[int], int]: + return x + + assert tyro.cli(main, args=[])(3) == 6 + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "something"]) + + +def test_fixed_dataclass_type() -> None: + def dummy(): + return 5 # noqa + + def main(x: Callable = dummy) -> Callable: + return x + + assert tyro.cli(main, args=[]) is dummy + with pytest.raises(SystemExit): + tyro.cli(main, args=["--x", "something"]) + + +def test_missing_singleton() -> None: + assert tyro.MISSING is copy.deepcopy(tyro.MISSING) + + +def test_torch_device() -> None: + def main(device: torch.device) -> torch.device: + return device + + assert tyro.cli(main, args=["--device", "cpu"]) == torch.device("cpu") + + +def test_supports_inference_mode_decorator() -> None: + @torch.inference_mode() + def main(x: int, device: str) -> Tuple[int, str]: + return x, device + + assert tyro.cli(main, args="--x 3 --device cuda".split(" ")) == (3, "cuda") + + +def test_torch_device_2() -> None: + assert tyro.cli(torch.device, args=["cpu"]) == torch.device("cpu") + + +def test_just_int() -> None: + assert tyro.cli(int, args=["123"]) == 123 + + +def test_just_dict() -> None: + assert tyro.cli(Dict[str, str], args="key value key2 value2".split(" ")) == { + "key": "value", + "key2": "value2", + } + + +def test_just_list() -> None: + assert tyro.cli(List[int], args="1 2 3 4".split(" ")) == [1, 2, 3, 4] + + +def test_just_tuple() -> None: + # Need a type: ignore for mypy. Seems like a mypy bug. + assert tyro.cli(Tuple[int, int, int, int], args="1 2 3 4".split(" ")) == ( # type: ignore + 1, + 2, + 3, + 4, + ) + + +def test_return_parser() -> None: + def main() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + return parser + + assert isinstance(tyro.cli(main, args=[]), argparse.ArgumentParser) + + +def test_pathlike_custom_class() -> None: + class CustomPath(pathlib.PurePosixPath): + def __new__(cls, *args: str | os.PathLike): + return super().__new__(cls, *args) + + def main(a: CustomPath) -> CustomPath: + return a + + assert tyro.cli(main, args=["--a", "/dev/null"]) == CustomPath("/dev/null") + + +def test_class_with_new_and_no_init() -> None: + class A(object): + def __new__(cls, x: int = 5): + return cls._custom_initializer(x) + + @classmethod + def _custom_initializer(cls, x: int = 5): + self = object.__new__(cls) + self.x = x # type: ignore + return self + + def __eq__(self, other) -> bool: + return self.x == other.x # type: ignore + + assert tyro.cli(A, args=["--x", "5"]) == A(x=5) + + +def test_return_unknown_args() -> None: + @dataclasses.dataclass + class A: + x: int = 0 + + a, unknown_args = tyro.cli( + A, args=["positional", "--x", "5", "--y", "7"], return_unknown_args=True + ) + assert a == A(x=5) + assert unknown_args == ["positional", "--y", "7"] + + +def test_unknown_args_with_arg_fixing() -> None: + @dataclasses.dataclass + class A: + x: int = 0 + + a, unknown_args = tyro.cli( + A, + args=["--x", "5", "--a_b", "--a-c"], + return_unknown_args=True, + ) + assert a == A(x=5) + # Should return the unfixed arguments + assert unknown_args == ["--a_b", "--a-c"] + + +def test_allow_ambiguous_args_when_not_returning_unknown_args() -> None: + @dataclasses.dataclass + class A: + a_b: List[int] = dataclasses.field(default_factory=list) + + a = tyro.cli( + A, + args=["--a_b", "5", "--a-b", "7"], + ) + assert a == A(a_b=[7]) + + +def test_disallow_ambiguous_args_when_returning_unknown_args() -> None: + @dataclasses.dataclass + class A: + x: int = 0 + + # If there's an argument that's ambiguous then we should raise an error when we're + # returning unknown args. + with pytest.raises(RuntimeError, match="Ambiguous .* --a_b and --a-b"): + tyro.cli( + A, + args=["--x", "5", "--a_b", "--a-b"], + return_unknown_args=True, + ) + + +def test_unknown_args_with_consistent_duplicates() -> None: + @dataclasses.dataclass + class A: + a_b: List[int] = dataclasses.field(default_factory=list) + c_d: List[int] = dataclasses.field(default_factory=list) + + # Tests logic for consistent duplicate arguments when performing argument fixing. + # i.e., we can fix arguments if the separator is consistent (all _'s or all -'s). + a, unknown_args = tyro.cli( + A, + args=[ + "--a-b", + "5", + "--a-b", + "7", + "--c_d", + "5", + "--c_d", + "7", + "--e-f", + "--e-f", + "--g_h", + "--g_h", + ], + return_unknown_args=True, + ) + assert a == A(a_b=[7], c_d=[7]) + assert unknown_args == ["--e-f", "--e-f", "--g_h", "--g_h"] + + +def test_pathlike(): + def main(x: os.PathLike) -> os.PathLike: + return x + + assert tyro.cli(main, args=["--x", "/dev/null"]) == pathlib.Path("/dev/null") + + +def test_unpack() -> None: + def main(*args: int, **kwargs: float) -> Tuple[Tuple[int, ...], Dict[str, float]]: + return args, kwargs + + assert tyro.cli(main, args=[]) == ((), {}) + assert tyro.cli(main, args="--args --kwargs".split(" ")) == ((), {}) + assert tyro.cli( + main, args="--args 1 2 3 --kwargs learning_rate 1e-4 beta1 0.99".split(" ") + ) == ((1, 2, 3), {"learning_rate": 1e-4, "beta1": 0.99}) + + +def test_unpack_with_other_args() -> None: + def main( + other1: str, other2: str = "3", *args: int, **kwargs: float + ) -> Tuple[str, str, Tuple[int, ...], Dict[str, float]]: + return other1, other2, args, kwargs + + assert tyro.cli(main, args=["--other1", "hi"]) == ("hi", "3", (), {}) + assert tyro.cli(main, args="--other1 ok --args --kwargs".split(" ")) == ( + "ok", + "3", + (), + {}, + ) + assert tyro.cli( + main, + args="--args 1 2 3 --kwargs learning_rate 1e-4 beta1 0.99 --other1 hello".split( + " " + ), + ) == ("hello", "3", (1, 2, 3), {"learning_rate": 1e-4, "beta1": 0.99}) + + +def test_empty_container() -> None: + @dataclasses.dataclass + class A: + x: Tuple[int, ...] = (1, 2, 3) + y: int | str | List[bool] = dataclasses.field( + default_factory=lambda: [False, False, True] + ) + + assert tyro.cli(A, args="--x".split(" ")).x == () + assert tyro.cli(A, args="--y".split(" ")).y == [] + + +def test_unknown_args_with_consistent_duplicates_use_underscores() -> None: + @dataclasses.dataclass + class A: + a_b: List[int] = dataclasses.field(default_factory=list) + c_d: List[int] = dataclasses.field(default_factory=list) + + # Tests logic for consistent duplicate arguments when performing argument fixing. + # i.e., we can fix arguments if the separator is consistent (all _'s or all -'s). + a, unknown_args = tyro.cli( + A, + args=[ + "--a-b", + "5", + "--a-b", + "7", + "--c_d", + "5", + "--c_d", + "7", + "--e-f", + "--e-f", + "--g_h", + "--g_h", + ], + return_unknown_args=True, + use_underscores=True, + ) + assert a == A(a_b=[7], c_d=[7]) + assert unknown_args == ["--e-f", "--e-f", "--g_h", "--g_h"] + + +def test_datetime_parsing() -> None: + assert tyro.cli( + datetime.datetime, args=["2024-08-25T14:30:00"] + ) == datetime.datetime(2024, 8, 25, 14, 30, 0) + assert tyro.cli( + datetime.datetime, args=["2024-08-25T14:30:00.123456"] + ) == datetime.datetime(2024, 8, 25, 14, 30, 0, 123456) + assert tyro.cli( + datetime.datetime, args=["2024-08-25T14:30:00+05:30"] + ) == datetime.datetime( + 2024, 8, 25, 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_date_parsing() -> None: + assert tyro.cli(datetime.date, args=["2024-08-25"]) == datetime.date(2024, 8, 25) + assert tyro.cli(datetime.date, args=["1999-12-31"]) == datetime.date(1999, 12, 31) + + +def test_time_parsing() -> None: + assert tyro.cli(datetime.time, args=["14:30:00"]) == datetime.time(14, 30, 0) + assert tyro.cli(datetime.time, args=["14:30:00.123456"]) == datetime.time( + 14, 30, 0, 123456 + ) + assert tyro.cli(datetime.time, args=["14:30:00+05:30"]) == datetime.time( + 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_datetime_parsing_with_main(): + def main(dt: datetime.datetime) -> datetime.datetime: + return dt + + assert tyro.cli(main, args=["--dt", "2024-08-25T14:30:00"]) == datetime.datetime( + 2024, 8, 25, 14, 30, 0 + ) + assert tyro.cli( + main, args=["--dt", "2024-08-25T14:30:00.123456"] + ) == datetime.datetime(2024, 8, 25, 14, 30, 0, 123456) + assert tyro.cli( + main, args=["--dt", "2024-08-25T14:30:00+05:30"] + ) == datetime.datetime( + 2024, 8, 25, 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_date_parsing_with_main(): + def main(dt: datetime.date) -> datetime.date: + return dt + + assert tyro.cli(main, args=["--dt", "2024-08-25"]) == datetime.date(2024, 8, 25) + assert tyro.cli(main, args=["--dt", "1999-12-31"]) == datetime.date(1999, 12, 31) + + +def test_time_parsing_with_main(): + def main(dt: datetime.time) -> datetime.time: + return dt + + assert tyro.cli(main, args=["--dt", "14:30:00"]) == datetime.time(14, 30, 0) + assert tyro.cli(main, args=["--dt", "14:30:00.123456"]) == datetime.time( + 14, 30, 0, 123456 + ) + assert tyro.cli(main, args=["--dt", "14:30:00+05:30"]) == datetime.time( + 14, 30, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800)) + ) + + +def test_date_parsing_harder_format(): + def main(dt: datetime.date) -> datetime.date: + return dt + + # Wrong separator '/'. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "2024/08/25"]) + + # Day-Month-Year format instead of Year-Month-Day. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "25-08-2024"]) + + +def test_time_parsing_harder_format(): + def main(dt: datetime.time) -> datetime.time: + return dt + + # Wrong separator '/'. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "14/30"]) + + # Missing seconds. + assert tyro.cli(main, args=["--dt", "14:30"]) == datetime.time(14, 30, 0) + + # Invalid minute value. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "14:60:00"]) + + # Invalid hour value. + with pytest.raises(SystemExit): + tyro.cli(main, args=["--dt", "25:00:00"]) diff --git a/tests/test_py311_generated/test_dict_namedtuple_generated.py b/tests/test_py311_generated/test_dict_namedtuple_generated.py new file mode 100644 index 000000000..638a2a12c --- /dev/null +++ b/tests/test_py311_generated/test_dict_namedtuple_generated.py @@ -0,0 +1,646 @@ +import contextlib +import copy +import dataclasses +import io +import pathlib +from typing import ( + Any, + Dict, + Literal, + Mapping, + NamedTuple, + NotRequired, + Required, + Tuple, + TypedDict, + cast, +) + +import pytest + +import tyro +import tyro._strings + + +def test_basic_dict() -> None: + def main(params: Dict[str, int]) -> Dict[str, int]: + return params + + assert tyro.cli(main, args="--params hey 5 hello 2".split(" ")) == { + "hey": 5, + "hello": 2, + } + assert tyro.cli(main, args="--params hey 5 hello 2".split(" ")) == { + "hey": 5, + "hello": 2, + } + assert tyro.cli(main, args="--params".split(" ")) == {} + with pytest.raises(SystemExit): + tyro.cli(main, args="--params hey 5 hello hey".split(" ")) + with pytest.raises(SystemExit): + tyro.cli(main, args="--params hey 5 hello".split(" ")) + + +def test_dict_with_default() -> None: + def main(params: Mapping[Literal[1, 3, 5, 7], bool] = {5: False, 1: True}) -> Any: + return params + + assert tyro.cli(main, args=[]) == {5: False, 1: True} + assert tyro.cli(main, args="--params.5 --params.no-1".split(" ")) == { + 5: True, + 1: False, + } + with pytest.raises(SystemExit): + tyro.cli(main, args="--params".split(" ")) + + +def test_tuple_in_dict() -> None: + def main(x: Dict[Tuple[int, int] | Tuple[str, str], Tuple[int, int]]) -> dict: + return x + + assert tyro.cli(main, args="--x 1 1 2 2 3 3 4 4".split(" ")) == { + (1, 1): (2, 2), + (3, 3): (4, 4), + } + + +def test_basic_typeddict() -> None: + class ManyTypesTypedDict(TypedDict): + i: int + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) + + +def test_positional_in_typeddict() -> None: + class ManyTypesTypedDict(TypedDict): + i: tyro.conf.Positional[int] + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="5".split(" ")) + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) + + +def test_total_false_typeddict() -> None: + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + assert tyro.cli(ManyTypesTypedDict, args=[]) == dict() + assert tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) == dict(i=5) + assert tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) == dict(s="5") + + +def test_total_false_required_typeddict() -> None: + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: Required[str] + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + # `s` is Required[]. + assert tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) == dict(i=5) + assert tyro.cli( + ManyTypesTypedDict, args="--i 5".split(" "), default={"s": "5"} + ) == dict(i=5, s="5") + assert tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) == dict(s="5") + + +def test_total_true_not_required_typeddict() -> None: + class ManyTypesTypedDict(TypedDict, total=True): + i: NotRequired[int] + s: str + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + # `s` is Required[]. + assert tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) == dict(i=5) + assert tyro.cli( + ManyTypesTypedDict, args="--i 5".split(" "), default={"s": "5"} + ) == dict(i=5, s="5") + assert tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) == dict(s="5") + + +def test_total_false_nested_typeddict() -> None: + class ChildTypedDict(TypedDict, total=False): + i: int + s: str + + class ParentTypedDict(TypedDict, total=False): + child: ChildTypedDict + + assert tyro.cli( + ParentTypedDict, + args="--child.i 5 --child.s 5".split(" "), + ) == {"child": {"i": 5, "s": "5"}} + + # total=False is ~ignored on the parent. + assert tyro.cli( + ParentTypedDict, + args=[], + ) == {"child": {}} + + +def test_total_false_typeddict_with_nested() -> None: + @dataclasses.dataclass + class Inner: + j: float + + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: Inner + + # --s.j is (unfortunately) still required. + with pytest.raises(SystemExit): + tyro.cli( + ManyTypesTypedDict, + args="".split(" "), + ) + + assert tyro.cli( + ManyTypesTypedDict, + args="--s.j 5".split(" "), + ) == {"s": Inner(5.0)} + + +def test_total_false_typeddict_with_tuple() -> None: + class ManyTypesTypedDict(TypedDict, total=False): + i: int + s: Tuple[str, str] + + assert ( + tyro.cli( + ManyTypesTypedDict, + args=[], + ) + == dict() + ) + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5 5".split(" "), + ) == dict(i=5, s=("5", "5")) + + +def test_nested_typeddict() -> None: + class ChildTypedDict(TypedDict): + y: int + + class NestedTypedDict(TypedDict): + x: int + b: ChildTypedDict + + assert tyro.cli(NestedTypedDict, args=["--x", "1", "--b.y", "3"]) == dict( + x=1, b=dict(y=3) + ) + with pytest.raises(SystemExit): + tyro.cli(NestedTypedDict, args=["--x", "1"]) + + +def test_helptext_and_default_typeddict() -> None: + class HelptextTypedDict(TypedDict): + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(HelptextTypedDict, default={"z": 3}, args=["--help"]) + helptext = tyro._strings.strip_ansi_sequences(f.getvalue()) + assert cast(str, HelptextTypedDict.__doc__) in helptext + assert "--x INT" in helptext + assert "--y INT" in helptext + assert "--z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_basic_namedtuple() -> None: + class ManyTypesNamedTuple(NamedTuple): + i: int + s: str + f: float + p: pathlib.Path + + assert tyro.cli( + ManyTypesNamedTuple, + args=[ + "--i", + "5", + "--s", + "5", + "--f", + "5", + "--p", + "~", + ], + ) == ManyTypesNamedTuple(i=5, s="5", f=5.0, p=pathlib.Path("~")) + + +def test_nested_namedtuple() -> None: + class ChildNamedTuple(NamedTuple): + y: int + + class NestedNamedTuple(NamedTuple): + x: int + b: ChildNamedTuple + + assert tyro.cli( + NestedNamedTuple, args=["--x", "1", "--b.y", "3"] + ) == NestedNamedTuple(x=1, b=ChildNamedTuple(y=3)) + with pytest.raises(SystemExit): + tyro.cli(NestedNamedTuple, args=["--x", "1"]) + + +def test_helptext_and_default_namedtuple() -> None: + class HelptextNamedTupleDefault(NamedTuple): + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int = 3 + """Documentation 3""" + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(HelptextNamedTupleDefault, args=["--help"]) + helptext = tyro._strings.strip_ansi_sequences(f.getvalue()) + assert cast(str, HelptextNamedTupleDefault.__doc__) in helptext + assert "--x INT" in helptext + assert "--y INT" in helptext + assert "--z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_and_default_namedtuple_alternate() -> None: + class HelptextNamedTuple(NamedTuple): + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + z: int + """Documentation 3""" + + with pytest.raises(SystemExit): + tyro.cli( + HelptextNamedTuple, + default=tyro.MISSING, + args=[], + ) + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli( + HelptextNamedTuple, + default=HelptextNamedTuple( + x=tyro.MISSING, + y=tyro.MISSING, + z=3, + ), + args=["--help"], + ) + helptext = tyro._strings.strip_ansi_sequences(f.getvalue()) + assert cast(str, HelptextNamedTuple.__doc__) in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3" in helptext + assert "(default: 3)" in helptext + + +def test_nested_dict() -> None: + loaded_config = { + "batch_size": 32, + "optimizer": { + "learning_rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule_type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule_type", + "exponential", + ], + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch_size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule_type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_use_underscores() -> None: + loaded_config = { + "batch_size": 32, + "optimizer": { + "learning_rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule_type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule-type", + "exponential", + ], + use_underscores=True, + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch_size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule_type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_hyphen() -> None: + # We do a lot of underscore <=> conversion in the code; this is just to make sure it + # doesn't break anything! + loaded_config = { + "batch-size": 32, + "optimizer": { + "learning-rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule-type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch-size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_hyphen_use_underscores() -> None: + # We do a lot of underscore <=> conversion in the code; this is just to make sure it + # doesn't break anything! + loaded_config = { + "batch-size": 32, + "optimizer": { + "learning-rate": 1e-4, + "epsilon": 1e-8, + "scheduler": {"schedule-type": "constant"}, + }, + } + backup_config = copy.deepcopy(loaded_config) + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch-size", + "16", + "--optimizer.scheduler.schedule-type", + "exponential", + ], + use_underscores=True, + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch-size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--batch_size", + "16", + "--optimizer.scheduler.schedule_type", + "exponential", + ], + use_underscores=True, + ) + + # Overridden config should be different from loaded config. + assert overrided_config != loaded_config + assert overrided_config["batch-size"] == 16 + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + + # Original loaded config should not be mutated. + assert loaded_config == backup_config + + +def test_nested_dict_annotations() -> None: + loaded_config = { + "optimizer": { + "scheduler": {"schedule-type": "constant"}, + }, + } + + overrided_config = tyro.cli( + dict, + default=loaded_config, + args=[ + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + del overrided_config + + overrided_config = tyro.cli( + Dict[str, Dict], + default=loaded_config, + args=[ + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + del overrided_config + + overrided_config = tyro.cli( + Dict[str, Dict[str, Dict]], + default=loaded_config, + args=[ + "--optimizer.scheduler.schedule-type", + "exponential", + ], + ) + assert overrided_config["optimizer"]["scheduler"]["schedule-type"] == "exponential" + del overrided_config + + +def test_functional_typeddict(): + """Source: https://github.com/brentyi/tyro/issues/87""" + NerfMLPHiddenLayers_0 = TypedDict( + "NerfMLPHiddenLayers_0", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + ) + NerfMLPHiddenLayers_1 = TypedDict( + "NerfMLPHiddenLayers_1", + { + "hidden_layers.0": NotRequired[int], + "hidden_layers.1": NotRequired[int], + "hidden_layers.2": NotRequired[int], + "hidden_layers.3": NotRequired[int], + "hidden_layers.4": NotRequired[int], + "hidden_layers.5": NotRequired[int], + "hidden_layers.6": NotRequired[int], + "hidden_layers.7": NotRequired[int], + }, + ) + NerfMLPHiddenLayers_2 = TypedDict( + "NerfMLPHiddenLayers_2", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + total=False, + ) + with pytest.raises(SystemExit): + tyro.cli(NerfMLPHiddenLayers_0, args=["--hidden_layers.0", "3"]) + assert tyro.cli(NerfMLPHiddenLayers_1, args=["--hidden_layers.0", "3"]) == { + "hidden_layers.0": 3 + } + assert tyro.cli(NerfMLPHiddenLayers_2, args=["--hidden_layers.0", "3"]) == { + "hidden_layers.0": 3 + } + + +def test_functional_typeddict_with_default(): + """Source: https://github.com/brentyi/tyro/issues/87""" + NerfMLPHiddenLayers_0 = TypedDict( + "NerfMLPHiddenLayers_0", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + ) + NerfMLPHiddenLayers_1 = TypedDict( + "NerfMLPHiddenLayers_1", + { + "hidden_layers.0": NotRequired[int], + "hidden_layers.1": NotRequired[int], + "hidden_layers.2": NotRequired[int], + "hidden_layers.3": NotRequired[int], + "hidden_layers.4": NotRequired[int], + "hidden_layers.5": NotRequired[int], + "hidden_layers.6": NotRequired[int], + "hidden_layers.7": NotRequired[int], + }, + ) + NerfMLPHiddenLayers_2 = TypedDict( + "NerfMLPHiddenLayers_2", + { + "hidden_layers.0": int, + "hidden_layers.1": int, + "hidden_layers.2": int, + "hidden_layers.3": int, + "hidden_layers.4": int, + "hidden_layers.5": int, + "hidden_layers.6": int, + "hidden_layers.7": int, + }, + total=False, + ) + with pytest.raises(SystemExit): + tyro.cli(NerfMLPHiddenLayers_0, args=["--hidden_layers.0", "3"], default={}) + assert tyro.cli( + NerfMLPHiddenLayers_1, args=["--hidden_layers.0", "3"], default={} + ) == {"hidden_layers.0": 3} + assert tyro.cli( + NerfMLPHiddenLayers_2, args=["--hidden_layers.0", "3"], default={} + ) == {"hidden_layers.0": 3} diff --git a/tests/test_py311_generated/test_dynamic_dataclasses_generated.py b/tests/test_py311_generated/test_dynamic_dataclasses_generated.py new file mode 100644 index 000000000..b0f30091a --- /dev/null +++ b/tests/test_py311_generated/test_dynamic_dataclasses_generated.py @@ -0,0 +1,14 @@ +from dataclasses import field, make_dataclass + +import pytest + +import tyro + + +def test_dynamic(): + B = make_dataclass("B", [("c", int, field())]) + A = make_dataclass("A", [("b", B, field())]) + + with pytest.raises(SystemExit): + tyro.cli(A, args=[]) + assert tyro.cli(A, args=["--b.c", "5"]) == A(b=B(c=5)) diff --git a/tests/test_py311_generated/test_errors_generated.py b/tests/test_py311_generated/test_errors_generated.py new file mode 100644 index 000000000..b7a9f5cfc --- /dev/null +++ b/tests/test_py311_generated/test_errors_generated.py @@ -0,0 +1,598 @@ +import contextlib +import dataclasses +import io +from typing import Dict, List, Literal, Tuple, TypeVar + +import pytest + +import tyro + + +def test_ambiguous_collection_0() -> None: + @dataclasses.dataclass + class A: + x: Tuple[Tuple[int, ...], ...] + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(A, args=["--x", "0", "1"]) + + +def test_ambiguous_collection_1() -> None: + @dataclasses.dataclass + class A: + x: List[List[int]] + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(A, args=["--x", "0", "1"]) + + +def test_ambiguous_collection_2() -> None: + def main(x: Tuple[List[str], List[str]]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError) as e: + tyro.cli(main, args=["--help"]) + assert "Unsupported type annotation for the field x" in e.value.args[0] + + +def test_ambiguous_collection_3() -> None: + def main(x: List[Tuple[int, int] | Tuple[int, int, int]]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError) as e: + tyro.cli(main, args=["--help"]) + assert "Unsupported type annotation for the field x" in e.value.args[0] + + +def test_ambiguous_collection_4() -> None: + X = List[Tuple[int, int] | Tuple[int, int, int]] + with pytest.raises(tyro.UnsupportedTypeAnnotationError) as e: + tyro.cli(X, args=["--help"]) + assert "Unsupported type annotation for the field" not in e.value.args[0] + + +# Must be global. +@dataclasses.dataclass +class _CycleDataclass: + x: "_CycleDataclass" + + +def test_cycle() -> None: + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(_CycleDataclass, args=[]) + + +def test_uncallable_annotation() -> None: + def main(arg: 5) -> None: # type: ignore + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_nested_annotation() -> None: + @dataclasses.dataclass + class OneIntArg: + x: int + + def main(arg: List[OneIntArg]) -> List[OneIntArg]: # type: ignore + return arg + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + @dataclasses.dataclass + class OneStringArg: + x: str + + def main(arg: List[OneStringArg]) -> List[OneStringArg]: # type: ignore + return arg + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--arg", "0", "1", "2"]) + + @dataclasses.dataclass + class TwoStringArg: + x: str + y: str + + def main2(arg: List[TwoStringArg]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main2, args=[]) + + +def test_missing_annotation_1() -> None: + def main(a, b) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_missing_annotation_2() -> None: + def main(*, a) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_tuple_needs_default() -> None: + def main(arg: tuple) -> None: # type: ignore + pass + + # This formerly raised an error, but now defaults to Tuple[str, ...]. + # with pytest.raises(tyro.UnsupportedTypeAnnotationError): + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_unbound_typevar() -> None: + T = TypeVar("T") + + def main(arg: T) -> None: # type: ignore + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_missing_default_fixed() -> None: + def main(value: tyro.conf.SuppressFixed[tyro.conf.Fixed[int]]) -> int: + return value + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_missing_default_suppressed() -> None: + def main(value: tyro.conf.Suppress[int]) -> int: + return value + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=["--help"]) + + +def test_ambiguous_sequence() -> None: + def main(value: list) -> None: + return None + + # This formerly raised an error, but now defaults to List[str]. + # with pytest.raises(tyro.UnsupportedTypeAnnotationError): + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_similar_arguments_basic() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class Class: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Class, args="--reward.trac".split(" ")) + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + + assert error.count("--reward.track") == 1 + assert error.count("--help") == 1 + + +def test_suppress_console_outputs() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class Class: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(Class, args="--reward.trac".split(" "), console_outputs=False) + + error = target.getvalue() + assert error == "" + + +def test_similar_arguments_subcommands() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(ClassA | ClassB, args="--reward.trac".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track") == 1 + assert error.count("--help") == 3 + + +def test_similar_arguments_subcommands_multiple() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + trace: int + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(ClassA | ClassB, args="--fjdkslaj --reward.trac".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Arguments similar to --reward.trac" in error + assert error.count("--reward.track {True,False}") == 1 + assert error.count("--reward.trace INT") == 1 + assert error.count("--help") == 5 + + +def test_similar_arguments_subcommands_multiple_contains_match() -> None: + @dataclasses.dataclass + class RewardConfigA: + track: bool + trace: int + + @dataclasses.dataclass + class RewardConfigB: ... + + @dataclasses.dataclass + class ClassA: + reward: RewardConfigA + + @dataclasses.dataclass + class ClassB: + reward: RewardConfigB + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( + ClassA | ClassB, + args="class-b --reward.track True --reward.trace 7".split(" "), + ) # type: ignore + + error = target.getvalue() + assert "Unrecognized or misplaced" in error + assert ( + "(applied to " in error and "class-b)" in error + ) # (applied to {root prog} class-b) + assert "Arguments similar to" in error + assert error.count("--reward.track {True,False}") == 2 + assert error.count("--reward.trace INT") == 2 + assert error.count("--help") == 5 # 2 subcommands * 2 arguments + usage hint. + + +def test_similar_arguments_subcommands_multiple_contains_match_alt() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + trace: int + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(ClassA | ClassB, args="--track".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track {True,False}") == 1 + assert ( + error.count("--help") == 3 + ) # Should show two possible subcommands + usage hint. + + +def test_similar_arguments_subcommands_overflow_different() -> None: + @dataclasses.dataclass + class RewardConfig: + track0: bool + track1: bool + track2: bool + track3: bool + track4: bool + track5: bool + track6: bool + track7: bool + track8: bool + track9: bool + track10: bool + track11: bool + track12: bool + track13: bool + track14: bool + track15: bool + track16: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(ClassA | ClassB, args="--track".split(" ")) # type: ignore + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track") == 10 + assert "[...]" not in error + assert error.count("--help") == 21 + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + # --tracked is intentionally between 0.8 ~ 0.9 similarity to track{i} for test + # coverage. + tyro.cli(RewardConfig, args="--tracked".split(" ")) # type: ignore + + # Usage print should be clipped. + error = target.getvalue() + assert "For full helptext, run" in error + + +def test_similar_arguments_subcommands_overflow_same() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + @dataclasses.dataclass + class ClassC: + reward: RewardConfig + + @dataclasses.dataclass + class ClassD: + reward: RewardConfig + + @dataclasses.dataclass + class ClassE: + reward: RewardConfig + + @dataclasses.dataclass + class ClassF: + reward: RewardConfig + + @dataclasses.dataclass + class ClassG: + reward: RewardConfig + + @dataclasses.dataclass + class ClassH: + reward: RewardConfig + + @dataclasses.dataclass + class ClassI: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( # type: ignore + ClassA + | ClassB + | ClassC + | ClassD + | ClassE + | ClassF + | ClassG + | ClassH + | ClassI, + args="--track".split(" "), + ) + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Perhaps you meant:" in error + assert error.count("--reward.track") == 1 + assert "[...]" in error + assert error.count("--help") == 5 + + +def test_similar_arguments_subcommands_overflow_same_startswith_multiple() -> None: + @dataclasses.dataclass + class RewardConfig: + track: bool + + @dataclasses.dataclass + class ClassA: + reward: RewardConfig + + @dataclasses.dataclass + class ClassB: + reward: RewardConfig + + @dataclasses.dataclass + class ClassC: + reward: RewardConfig + + @dataclasses.dataclass + class ClassD: + reward: RewardConfig + + @dataclasses.dataclass + class ClassE: + reward: RewardConfig + rewarde: tyro.conf.Suppress[int] = 10 + + @dataclasses.dataclass + class ClassF: + reward: RewardConfig + + @dataclasses.dataclass + class ClassG: + reward: RewardConfig + + @dataclasses.dataclass + class ClassH: + reward: RewardConfig + + @dataclasses.dataclass + class ClassI: + reward: RewardConfig + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( # type: ignore + ClassA + | ClassB + | ClassC + | ClassD + | ClassE + | ClassF + | ClassG + | ClassH + | ClassI, + args="--track --ffff".split(" "), + ) + + error = target.getvalue() + assert "Unrecognized option" in error + assert "Arguments similar to --track" in error + assert error.count("--rewar") == 1 + assert "rewarde" not in error + assert "[...]" in error + assert error.count("--help") == 5 + + +def test_similar_flag() -> None: + @dataclasses.dataclass + class Args: + flag: bool = False + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli( + Args, + args="--lag".split(" "), + ) + + error = target.getvalue() + + # We don't print usage text anymore. + assert error.count("--flag | --no-flag") == 0 + + # Printed in the similar argument list. + assert error.count("--flag, --no-flag") == 1 + + +def test_value_error() -> None: + """https://github.com/brentyi/tyro/issues/86""" + + def main() -> None: + raise ValueError("This shouldn't be caught by tyro") + + with pytest.raises(ValueError): + tyro.cli(main, args=[]) + + +def test_value_error_subcommand() -> None: + """https://github.com/brentyi/tyro/issues/86""" + + def main1() -> None: + raise ValueError("This shouldn't be caught by tyro") + + def main2() -> None: + raise ValueError("This shouldn't be caught by tyro") + + with pytest.raises(ValueError): + tyro.extras.subcommand_cli_from_dict( + {"main1": main1, "main2": main2}, args=["main1"] + ) + + +def test_wrong_annotation() -> None: + @dataclasses.dataclass + class Args: + x: dict = None # type: ignore + + with pytest.warns(UserWarning): + assert tyro.cli(Args, args=[]).x is None + with pytest.warns(UserWarning): + assert tyro.cli(Args, args="--x key value k v".split(" ")).x == { + "key": "value", + "k": "v", + } + + +def test_bad_dict_control() -> None: + assert tyro.cli( + Dict[Literal["a", "b"], Literal["c", "d"]], args="a c b d".split(" ") + ) == {"a": "c", "b": "d"} + + +def test_bad_dict_key() -> None: + with pytest.raises(SystemExit): + tyro.cli(Dict[Literal["a", "b"], Literal["c", "d"]], args="c c d d".split(" ")) + + +def test_bad_dict_val() -> None: + assert tyro.cli( + Dict[Literal["a", "b"], Literal["c", "d"]], args="a c b d".split(" ") + ) == {"a": "c", "b": "d"} + with pytest.raises(SystemExit): + tyro.cli(Dict[Literal["a", "b"], Literal["c", "d"]], args="a a b b".split(" ")) + + +def test_invalid_subcommand_default() -> None: + @dataclasses.dataclass + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + + @dataclasses.dataclass + class C: + c: int + + assert tyro.cli(A | B, args="a --a 5".split(" ")) == A(5) # type: ignore + assert tyro.cli(A | B, args="b --b 5".split(" ")) == B(5) # type: ignore + + with pytest.warns(UserWarning): + assert tyro.cli(A | B, default=C(3), args=[]) == C(3) # type: ignore + + with pytest.warns(UserWarning): + assert tyro.cli(A | B, default=C(3), args="b --b 5".split(" ")) == B(5) # type: ignore diff --git a/tests/test_py311_generated/test_errors_new_annotations_generated.py b/tests/test_py311_generated/test_errors_new_annotations_generated.py new file mode 100644 index 000000000..da4aee733 --- /dev/null +++ b/tests/test_py311_generated/test_errors_new_annotations_generated.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import sys + +import pytest + +import tyro + + +@pytest.mark.skipif( + sys.version_info >= (3, 10), reason="No error for newer versions of Python." +) +def test_new_union_error() -> None: + """PEP 604 allows `|` to be used as a type annotation in Python >=3.10.""" + + def main(x: int | str) -> None: ... + + with pytest.raises(TypeError) as e: + tyro.cli(main) + assert "You may be using a Union in the form of `X | Y`" in e.value.args[0] + + +@pytest.mark.skipif( + sys.version_info >= (3, 9), reason="No error for newer versions of Python." +) +def test_new_collection_error() -> None: + """PEP 585 allows standard collections to be used as generics in Python >=3.9.""" + + def main(x: list[int]) -> None: ... + + with pytest.raises(TypeError) as e: + tyro.cli(main) + assert "You may be using a standard collection as a generic" in e.value.args[0] diff --git a/tests/test_py311_generated/test_flax_min_py38_generated.py b/tests/test_py311_generated/test_flax_min_py38_generated.py new file mode 100644 index 000000000..ab450c77a --- /dev/null +++ b/tests/test_py311_generated/test_flax_min_py38_generated.py @@ -0,0 +1,68 @@ +"""Tests initializing flax modules directly via tyro.""" + +from typing import cast + +import jax +import pytest +from flax import linen as nn +from helptext_utils import get_helptext_with_checks +from jax import numpy as jnp + +import tyro + + +class Classifier(nn.Module): + layers: int + units: int + output_dim: int + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # type: ignore + for i in range(self.layers - 1): + x = nn.Dense( + self.units, + kernel_init=nn.initializers.kaiming_normal(), + )(x) + x = nn.relu(x) + + x = nn.Dense( + self.output_dim, + kernel_init=nn.initializers.xavier_normal(), + )(x) + x = nn.sigmoid(x) + return x + + +def test_ok(): + network = tyro.cli( + Classifier, + args=[ + "--layers", + "3", + "--units", + "8", + "--output-dim", + "3", + ], + ) + + x = jnp.zeros((10, 4)) + params = network.init(jax.random.PRNGKey(0), x) + assert cast(jax.Array, network.apply(params, x)).shape == (10, 3) + + helptext = get_helptext_with_checks(Classifier) + assert "parent" not in helptext + assert "name" not in helptext + + +def test_missing(): + with pytest.raises(SystemExit): + tyro.cli( + Classifier, + args=[ + "--layers", + "3", + "--units", + "8", + ], + ) diff --git a/tests/test_py311_generated/test_forward_ref_generated.py b/tests/test_py311_generated/test_forward_ref_generated.py new file mode 100644 index 000000000..e54438bd0 --- /dev/null +++ b/tests/test_py311_generated/test_forward_ref_generated.py @@ -0,0 +1,48 @@ +import dataclasses +from typing import Union + +import pytest + +import tyro + + +@dataclasses.dataclass +class A1: + x: int + bc: "B| C" + + +@dataclasses.dataclass +class A2: + x: int + bc: Union["B", "C"] + + +@dataclasses.dataclass +class B: + y: "int" + + +@dataclasses.dataclass +class C: + z: int + + +def test_forward_ref_1(): + assert tyro.cli(A1, args=["--x", "1", "bc:b", "--bc.y", "3"]) == A1(x=1, bc=B(y=3)) + assert tyro.cli(A1, args=["--x", "1", "bc:c", "--bc.z", "3"]) == A1(x=1, bc=C(z=3)) + + with pytest.raises(SystemExit): + tyro.cli(A1, args=["--x", "1", "bc:b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(A1, args=["--x", "1", "bc:c", "--bc.y", "3"]) + + +def test_forward_ref_2(): + assert tyro.cli(A2, args=["--x", "1", "bc:b", "--bc.y", "3"]) == A2(x=1, bc=B(y=3)) + assert tyro.cli(A2, args=["--x", "1", "bc:c", "--bc.z", "3"]) == A2(x=1, bc=C(z=3)) + + with pytest.raises(SystemExit): + tyro.cli(A2, args=["--x", "1", "bc:b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(A2, args=["--x", "1", "bc:c", "--bc.y", "3"]) diff --git a/tests/test_py311_generated/test_functools_generated.py b/tests/test_py311_generated/test_functools_generated.py new file mode 100644 index 000000000..f1a0aaf06 --- /dev/null +++ b/tests/test_py311_generated/test_functools_generated.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import dataclasses +import functools + +import pytest +from helptext_utils import get_helptext_with_checks + +import tyro + + +def test_partial_func() -> None: + def main(a: int, b: str) -> str: + return b * a + + assert tyro.cli(functools.partial(main, a=3), args=["--b", "hi"]) == "hihihi" + + +def test_partial_class() -> None: + class Main: + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + assert tyro.cli(functools.partial(Main, a=3), args=["--b", "hi"]).inner == "hihihi" + + +def test_partial_helptext_func() -> None: + def main(a: int, b: str) -> str: + """Hello!""" + return b * a + + helptext = get_helptext_with_checks(functools.partial(main, b="hello world")) + assert "partial" not in helptext + assert "Hello!" in helptext + assert "hello world" in helptext + + +def test_partial_helptext_class() -> None: + class Main: + """Hello!""" + + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + helptext = get_helptext_with_checks(functools.partial(Main, b="3")) + assert "partial" not in helptext + assert "Hello!" in helptext + + +def test_wraps_func() -> None: + def main(a: int, b: str) -> str: + return b * a + + @functools.wraps(main) + def wrapper(*args, **kwargs) -> int: + return kwargs["a"] + + assert tyro.cli(wrapper, args=["--a", "3", "--b", "hi"]) == 3 + with pytest.raises(SystemExit): + tyro.cli(wrapper, args=["--a", "3"]) + + +def test_wraps_partial_func_helptext() -> None: + def main(a: int, b: str) -> str: + """Hello! + + Args: + a: Argument. + b: Argument. + """ + return b * a + + @functools.wraps(main) + def wrapper(*args, **kwargs) -> int: + return kwargs["a"] + + assert tyro.cli(functools.partial(wrapper, a=3), args=["--b", "hi"]) == 3 + + helptext = get_helptext_with_checks(functools.partial(wrapper, b="3")) + assert "wraps" not in helptext + assert "Hello!" in helptext + assert "Argument." in helptext + + +def test_wraps_partial_class_helptext() -> None: + class Main: + """Hello!""" + + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + @functools.wraps(Main) + def wrapper(*args, **kwargs) -> int: + return kwargs["a"] + + assert tyro.cli(functools.partial(wrapper, a=3), args=["--b", "hi"]) == 3 + + helptext = get_helptext_with_checks(functools.partial(wrapper, b="3")) + assert "wraps" not in helptext + assert "Hello!" in helptext + + +@dataclasses.dataclass +class WrappedDataclass: + """Hello!""" + + a: int + b: str + """Second field.""" + + +def test_wraps_partial_dataclass() -> None: + @functools.wraps(WrappedDataclass) + def wrapper(*args, **kwargs) -> str: + return kwargs["a"] * kwargs["b"] + + assert tyro.cli(functools.partial(wrapper, a=3), args=["--b", "hi"]) == "hihihi" + assert ( + tyro.cli( + functools.partial(functools.partial(wrapper, a=3), b="hello"), + args=["--b", "hi"], + ) + == "hihihi" + ) + assert ( + tyro.cli(functools.partial(functools.partial(wrapper, a=3), b="hello"), args=[]) + == "hellohellohello" + ) + + helptext = get_helptext_with_checks(functools.partial(wrapper, b="3")) + assert "wraps" not in helptext + assert "Hello!" in helptext + assert "Second field." in helptext diff --git a/tests/test_py311_generated/test_generics_and_serialization_generated.py b/tests/test_py311_generated/test_generics_and_serialization_generated.py new file mode 100644 index 000000000..d832a3f31 --- /dev/null +++ b/tests/test_py311_generated/test_generics_and_serialization_generated.py @@ -0,0 +1,458 @@ +import contextlib +import dataclasses +import enum +import io +from typing import Annotated, Generic, List, NewType, Tuple, Type, TypeVar + +import pytest +import yaml + +import tyro + +T = TypeVar("T") + + +def _check_serialization_identity(cls: Type[T], instance: T) -> None: + assert tyro.extras.from_yaml(cls, tyro.extras.to_yaml(instance)) == instance + + +ScalarType = TypeVar("ScalarType") + + +def test_tuple_generic_variable() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + assert tyro.cli( + TupleGenericVariable[int], args=["--xyz", "1", "2", "3"] + ) == TupleGenericVariable((1, 2, 3)) + + +def test_tuple_generic_variable_newtype() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + SpecialInt = NewType("SpecialInt", int) + assert tyro.cli( + TupleGenericVariable[SpecialInt], args=["--xyz", "1", "2", "3"] + ) == TupleGenericVariable((SpecialInt(1), SpecialInt(2), SpecialInt(3))) + + +def test_tuple_generic_variable_newtype_container() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + SpecialInt = NewType("SpecialInt", Tuple[int]) + assert tyro.cli( + TupleGenericVariable[SpecialInt], args=["--xyz", "1", "2", "3"] + ) == TupleGenericVariable((SpecialInt((1,)), SpecialInt((2,)), SpecialInt((3,)))) + + +def test_tuple_generic_variable_more_newtype() -> None: + @dataclasses.dataclass + class TupleGenericVariable(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + SpecialInt = NewType("SpecialInt", int) + SpecialTuple = NewType("SpecialTuple", TupleGenericVariable[SpecialInt]) + assert tyro.cli(SpecialTuple, args=["--xyz", "1", "2", "3"]) == SpecialTuple( + TupleGenericVariable((SpecialInt(1), SpecialInt(2), SpecialInt(3))) + ) + + +def test_tuple_generic_helptext() -> None: + @dataclasses.dataclass + class TupleGenericVariableHelptext(Generic[ScalarType]): + """Helptext!""" + + xyz: Tuple[ScalarType, ...] + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(TupleGenericVariableHelptext[int], args=["--help"]) + helptext = f.getvalue() + assert "Helptext!" in helptext + + +def test_tuple_generic_no_helptext() -> None: + @dataclasses.dataclass + class TupleGenericVariableNoHelptext(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(TupleGenericVariableNoHelptext[int], args=["--help"]) + helptext = f.getvalue() + assert "Helptext!" not in helptext + + # Check that we don't accidentally grab docstrings from the generic alias! + assert "The central part of internal API" not in helptext + + +def test_tuple_generic_fixed() -> None: + @dataclasses.dataclass + class TupleGenericFixed(Generic[ScalarType]): + xyz: Tuple[ScalarType, ScalarType, ScalarType] + + assert tyro.cli( + TupleGenericFixed[int], args=["--xyz", "1", "2", "3"] + ) == TupleGenericFixed((1, 2, 3)) + + +class CoordinateFrame(enum.Enum): + WORLD = enum.auto() + CAMERA = enum.auto() + + +@dataclasses.dataclass +class Point3(Generic[ScalarType]): + x: ScalarType + y: ScalarType + z: ScalarType + frame: CoordinateFrame + + +def test_simple_generic() -> None: + @dataclasses.dataclass + class SimpleGeneric: + point_continuous: Point3[float] + point_discrete: Point3[int] + + parsed_instance = tyro.cli( + SimpleGeneric, + args=[ + "--point-continuous.x", + "1.2", + "--point-continuous.y", + "2.2", + "--point-continuous.z", + "3.2", + "--point-continuous.frame", + "WORLD", + "--point-discrete.x", + "1", + "--point-discrete.y", + "2", + "--point-discrete.z", + "3", + "--point-discrete.frame", + "WORLD", + ], + ) + assert parsed_instance == SimpleGeneric( + Point3(1.2, 2.2, 3.2, CoordinateFrame.WORLD), + Point3(1, 2, 3, CoordinateFrame.WORLD), + ) + _check_serialization_identity(SimpleGeneric, parsed_instance) + + with pytest.raises(SystemExit): + # Accidentally pass in floats instead of ints for discrete + tyro.cli( + SimpleGeneric, + args=[ + "--point-continuous.x", + "1.2", + "--point-continuous.y", + "2.2", + "--point-continuous.z", + "3.2", + "--point-continuous.frame", + "WORLD", + "--point-discrete.x", + "1.5", + "--point-discrete.y", + "2.5", + "--point-discrete.z", + "3.5", + "--point-discrete.frame", + "WORLD", + ], + ) + + +def test_multilevel_generic() -> None: + @dataclasses.dataclass + class Triangle(Generic[ScalarType]): + a: Point3[ScalarType] + b: Point3[ScalarType] + c: Point3[ScalarType] + + parsed_instance = tyro.cli( + Triangle[float], + args=[ + "--a.x", + "1.0", + "--a.y", + "1.2", + "--a.z", + "1.3", + "--a.frame", + "WORLD", + "--b.x", + "1.0", + "--b.y", + "1.2", + "--b.z", + "1.3", + "--b.frame", + "WORLD", + "--c.x", + "1.0", + "--c.y", + "1.2", + "--c.z", + "1.3", + "--c.frame", + "WORLD", + ], + ) + assert parsed_instance == Triangle( + Point3(1.0, 1.2, 1.3, CoordinateFrame.WORLD), + Point3(1.0, 1.2, 1.3, CoordinateFrame.WORLD), + Point3(1.0, 1.2, 1.3, CoordinateFrame.WORLD), + ) + _check_serialization_identity(Triangle[float], parsed_instance) + + +def test_multilevel_generic_no_helptext() -> None: + @dataclasses.dataclass + class LineSegment(Generic[ScalarType]): + a: Point3[ScalarType] + b: Point3[ScalarType] + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(LineSegment[int], args=["--help"]) + helptext = f.getvalue() + + # Check that we don't accidentally grab docstrings from the generic alias! + assert "The central part of internal API" not in helptext + + +def test_generic_nested_dataclass() -> None: + @dataclasses.dataclass + class Child: + a: int + b: int + + T = TypeVar("T") + + @dataclasses.dataclass + class DataclassGeneric(Generic[T]): + child: T + + parsed_instance = tyro.cli( + DataclassGeneric[Child], args=["--child.a", "5", "--child.b", "7"] + ) + assert parsed_instance == DataclassGeneric(Child(5, 7)) + + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity(DataclassGeneric[Child], parsed_instance) + + +def test_generic_nested_dataclass_helptext() -> None: + @dataclasses.dataclass + class Child: + a: int + b: int + + T = TypeVar("T") + + @dataclasses.dataclass + class DataclassGeneric(Generic[T]): + child: T + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(DataclassGeneric[Child], args=["--help"]) + helptext = f.getvalue() + + # Check that we don't accidentally grab docstrings from the generic alias! + assert "The central part of internal API" not in helptext + + +def test_generic_subparsers() -> None: + @dataclasses.dataclass + class CommandOne: + a: int + + @dataclasses.dataclass + class CommandTwo: + b: int + + T1 = TypeVar("T1") + T2 = TypeVar("T2") + + @dataclasses.dataclass + class Subparser(Generic[T1, T2]): + command: T1 | T2 + + parsed_instance = tyro.cli( + Subparser[CommandOne, CommandTwo], + args="command:command-one --command.a 5".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(CommandOne(5)) # type: ignore + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[CommandOne, CommandTwo], parsed_instance + ) + + parsed_instance = tyro.cli( + Subparser[CommandOne, CommandTwo], + args="command:command-two --command.b 7".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(CommandTwo(7)) # type: ignore + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[CommandOne, CommandTwo], parsed_instance + ) + + +def test_generic_subparsers_in_container() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class Command(Generic[T]): + a: List[T] + + T1 = TypeVar("T1") + T2 = TypeVar("T2") + + @dataclasses.dataclass + class Subparser(Generic[T1, T2]): + command: T1 | T2 + + parsed_instance = tyro.cli( + Subparser[Command[int], Command[float]], + args="command:command-int --command.a 5 3".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(Command([5, 3])) and isinstance( # type: ignore + parsed_instance.command.a[0], int + ) + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[Command[int], Command[float]], parsed_instance + ) + + parsed_instance = tyro.cli( + Subparser[Command[int], Command[float]], + args="command:command-float --command.a 7 2".split(" "), + ) + # Not supported in mypy. + assert parsed_instance == Subparser(Command([7.0, 2.0])) and isinstance( # type: ignore + parsed_instance.command.a[0], float + ) + # Local generics will break. + with pytest.raises(yaml.constructor.ConstructorError): + _check_serialization_identity( + Subparser[Command[int], Command[float]], parsed_instance + ) + + +def test_serialize_missing() -> None: + @dataclasses.dataclass + class TupleGenericVariableMissing(Generic[ScalarType]): + xyz: Tuple[ScalarType, ...] + + x = TupleGenericVariableMissing[int](xyz=(tyro.MISSING, tyro.MISSING)) + assert tyro.extras.to_yaml(x).count("!missing") == 2 + _check_serialization_identity(TupleGenericVariableMissing[int], x) + assert ( + tyro.extras.from_yaml( + TupleGenericVariableMissing[int], tyro.extras.to_yaml(x) + ).xyz[0] + is tyro.MISSING + ) + + +def test_generic_inherited_type_narrowing() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class ActualParentClass(Generic[T]): + x: T # Documentation 1 + + # Documentation 2 + y: T + + z: T = 3 # type: ignore + """Documentation 3""" + + @dataclasses.dataclass + class ChildClass(ActualParentClass[int]): + a: int = 7 + + def main(x: ActualParentClass[int] = ChildClass(5, 5)) -> ActualParentClass: + return x + + assert tyro.cli(main, args="--x.x 3".split(" ")) == ChildClass(3, 5, 3) + + +def test_pculbertson() -> None: + # https://github.com/brentyi/tyro/issues/7 + + @dataclasses.dataclass(frozen=True) + class TypeA: + data: int + + @dataclasses.dataclass + class TypeB: + data: int + + @dataclasses.dataclass + class Wrapper: + subclass: TypeA | TypeB = TypeA(1) + + wrapper1 = Wrapper() # Create Wrapper object. + assert wrapper1 == tyro.extras.from_yaml(Wrapper, tyro.extras.to_yaml(wrapper1)) + + +def test_annotated() -> None: + # https://github.com/brentyi/tyro/issues/7 + + @dataclasses.dataclass(frozen=True) + class TypeA: + data: int + + @dataclasses.dataclass + class Wrapper: + subclass: Annotated[TypeA, int] = TypeA(1) + + wrapper1 = Wrapper() # Create Wrapper object. + assert wrapper1 == tyro.extras.from_yaml(Wrapper, tyro.extras.to_yaml(wrapper1)) + + +@dataclasses.dataclass +class TypeA: + data: int + + +@dataclasses.dataclass +class TypeASubclass(TypeA): + pass + + +@dataclasses.dataclass +class Wrapper: + subclass: TypeA + + +def test_superclass() -> None: + # https://github.com/brentyi/tyro/issues/7 + + wrapper1 = Wrapper(TypeASubclass(3)) # Create Wrapper object. + assert wrapper1 == tyro.extras.from_yaml(Wrapper, tyro.extras.to_yaml(wrapper1)) diff --git a/tests/test_py311_generated/test_helptext_generated.py b/tests/test_py311_generated/test_helptext_generated.py new file mode 100644 index 000000000..075fdac46 --- /dev/null +++ b/tests/test_py311_generated/test_helptext_generated.py @@ -0,0 +1,975 @@ +import dataclasses +import enum +import json +import os +import pathlib +from collections.abc import Callable +from typing import ( + Annotated, + Any, + Dict, + Generic, + List, + Literal, + NotRequired, + Optional, + Tuple, + TypedDict, + TypeVar, + cast, +) + +from helptext_utils import get_helptext_with_checks +from torch import nn + +import tyro + + +def test_helptext() -> None: + @dataclasses.dataclass + class Helptext: + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: Annotated[int, "ignored"] + + z: int = 3 + """Documentation 3""" + + helptext = get_helptext_with_checks(Helptext) + assert cast(str, helptext) in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_sphinx_autodoc_style() -> None: + @dataclasses.dataclass + class Helptext: + """This docstring should be printed as a description.""" + + x: int #: Documentation 1 + + #:Documentation 2 + y: Annotated[int, "ignored"] + z: int = 3 + + helptext = get_helptext_with_checks(Helptext) + assert cast(str, helptext) in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert ": Documentation 1" not in helptext + assert "Documentation 2 (required)" in helptext + assert ":Documentation 2" not in helptext + + # :Documentation 2 should not be applied to `z`. + assert helptext.count("Documentation 2") == 1 + + +def test_helptext_from_class_docstring() -> None: + @dataclasses.dataclass + class Helptext2: + """This docstring should be printed as a description. + + Attributes: + x: Documentation 1 + y: Documentation 2 + z: Documentation 3 + """ + + x: int + y: Annotated[int, "ignored"] + z: int = 3 + + helptext = get_helptext_with_checks(Helptext2) + assert "This docstring should be printed as a description" in helptext + assert "Attributes" not in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_from_class_docstring_args() -> None: + @dataclasses.dataclass + class Helptext3: + """This docstring should be printed as a description. + + Args: + x: Documentation 1 + y: Documentation 2 + z: Documentation 3 + """ + + x: int + y: Annotated[int, "ignored"] + z: int = 3 + + helptext = get_helptext_with_checks(Helptext3) + assert "This docstring should be printed as a description" in helptext + assert "Args" not in helptext + assert "x INT" in helptext + assert "y INT" in helptext + assert "z INT" in helptext + assert "Documentation 1 (required)" in helptext + assert "Documentation 2 (required)" in helptext + assert "Documentation 3 (default: 3)" in helptext + + +def test_helptext_inherited() -> None: + class UnrelatedParentClass: + pass + + @dataclasses.dataclass + class ActualParentClass: + """This docstring should be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + # fmt: off + + z: int = 3 + def some_method(self) -> None: # noqa + """Coverage stress test.""" + # fmt: on + + @dataclasses.dataclass + class ChildClass(UnrelatedParentClass, ActualParentClass): + pass + + helptext = get_helptext_with_checks(ChildClass) + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + + +def test_helptext_inherited_default_override() -> None: + @dataclasses.dataclass + class ParentClass: + """This docstring should __not__ be printed as a description.""" + + x: int # Documentation 1 + + # Documentation 2 + y: int + + # fmt: off + + z: int = 3 + def some_method(self) -> None: # noqa + """Coverage stress test.""" + # fmt: on + + @dataclasses.dataclass + class ChildClass(ParentClass): + """This docstring should be printed as a description.""" + + def main(x: ParentClass = ChildClass(x=5, y=5)) -> Any: + """Main function.""" + return x + + helptext = get_helptext_with_checks(main) + assert "Main function." in helptext + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "__not__" not in helptext + assert helptext.count("should be printed") == 1 + + +def test_helptext_nested() -> None: + """For nested classes, we should pull helptext from the outermost docstring if + possible. The class docstring can be used as a fallback.""" + + class Inner: + def __init__(self, a: int): + """Something + + Args: + a (int): Hello world! + """ + + def main_with_docstring(a: Inner) -> None: + """main_with_docstring. + + Args: + a: Documented in function. + """ + + def main_no_docstring(a: Inner) -> None: + """main_no_docstring.""" + + helptext = get_helptext_with_checks(main_with_docstring) + assert "Documented in function" in helptext and str(Inner.__doc__) not in helptext + assert "main_with_docstring." in helptext + assert "Args:" not in helptext + assert "Args:" not in helptext + assert "Hello world!" in helptext + + helptext = get_helptext_with_checks(main_no_docstring) + assert "Something" in helptext + assert "main_no_docstring." in helptext + assert "Args:" not in helptext + assert "Hello world!" in helptext + + +def test_helptext_defaults() -> None: + class Color(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + @dataclasses.dataclass + class HelptextWithVariousDefaults: + x: pathlib.Path = pathlib.Path("/some/path/to/a/file") + y: Color = Color.RED + z: str = "%" + + helptext = get_helptext_with_checks(HelptextWithVariousDefaults) + assert "show this help message and exit" in helptext + assert "--x PATH" in helptext + assert "(default: /some/path/to/a/file)" in helptext + assert "--y {RED,GREEN,BLUE}" in helptext + assert "(default: RED)" in helptext + assert "--z STR" in helptext + assert "(default: %)" in helptext + + +def test_multiline_helptext() -> None: + @dataclasses.dataclass + class HelptextMultiline: + x: int # Documentation 1 + + # This comment should be ignored! + + # Documentation 2 + # Next line of documentation 2 + y: int + + z: int = 3 + """Documentation 3 + Next line of documentation 3""" + + helptext = get_helptext_with_checks(HelptextMultiline) + assert "Documentation 1 (required)" in helptext + assert "Documentation 2" in helptext + assert "documentation 2" in helptext + assert "Documentation 3" in helptext + assert "documentation 3" in helptext + + +def test_grouped_helptext() -> None: + @dataclasses.dataclass + class HelptextGrouped: + x: int # Documentation 1 + # Description of both y and z. + y: int + z: int = 3 + + helptext = get_helptext_with_checks(HelptextGrouped) + assert "Documentation 1 (required)" in helptext + assert "Description of both y and z. (required)" in helptext + assert "Description of both y and z. (default: 3)" in helptext + + +def test_none_default_value_helptext() -> None: + @dataclasses.dataclass + class Config: + x: Optional[int] = None + """An optional variable.""" + + helptext = get_helptext_with_checks(Config) + assert "--x {None}|INT" in helptext + assert "An optional variable. (default: None)" in helptext + + +def test_helptext_hard_bool() -> None: + @dataclasses.dataclass + class HelptextHardString: + # fmt: off + x: bool = ( + False + ) + """Helptext. 2% milk.""" + # fmt: on + + # Note that the percent symbol needs some extra handling in argparse. + # https://stackoverflow.com/questions/21168120/python-argparse-errors-with-in-help-string + + helptext = get_helptext_with_checks(HelptextHardString) + assert "--x" in helptext + assert "2% milk." in helptext + + +def test_helptext_with_inheritance() -> None: + @dataclasses.dataclass + class Parent: + # fmt: off + x: str = ( + "This docstring may be tougher to parse!" + ) + """Helptext.""" + # fmt: on + + @dataclasses.dataclass + class Child(Parent): + pass + + helptext = get_helptext_with_checks(Child) + assert "--x STR" in helptext + assert "Helptext." in helptext + assert "(default: 'This docstring" in helptext + + +def test_helptext_with_inheritance_overriden() -> None: + @dataclasses.dataclass + class Parent2: + # fmt: off + x: str = ( + "This docstring may be tougher to parse!" + ) + """Helptext.""" + # fmt: on + + @dataclasses.dataclass + class Child2(Parent2): + # fmt: off + x: str = ( + "This docstring may be tougher to parse?" + ) + """Helptext!""" + # fmt: on + + helptext = get_helptext_with_checks(Child2) + assert "--x STR" in helptext + assert "Helptext! (default: 'This" in helptext + + +def test_tuple_helptext() -> None: + @dataclasses.dataclass + class TupleHelptext: + x: Tuple[int, str, float] + + helptext = get_helptext_with_checks(TupleHelptext) + assert "--x INT STR FLOAT" in helptext + + +def test_tuple_helptext_defaults() -> None: + @dataclasses.dataclass + class TupleHelptextDefaults: + x: Tuple[int, str, str] = (5, "hello world", "hello") + + helptext = get_helptext_with_checks(TupleHelptextDefaults) + assert "--x INT STR STR" in helptext + assert "(default: 5 'hello world' hello)" in helptext + + +def test_generic_helptext() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class GenericTupleHelptext(Generic[T]): + x: T + + helptext = get_helptext_with_checks(GenericTupleHelptext[int]) + assert "--x INT" in helptext + + +def test_generic_tuple_helptext() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class GenericTupleHelptext(Generic[T]): + x: Tuple[T, T, T] + + helptext = get_helptext_with_checks(GenericTupleHelptext[int]) + assert "--x INT INT INT" in helptext + + +def test_generic_list_helptext() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class GenericTupleHelptext(Generic[T]): + x: List[T] + + helptext = get_helptext_with_checks(GenericTupleHelptext[int]) + assert "--x [INT [INT ...]]" in helptext + + +def test_literal_helptext() -> None: + @dataclasses.dataclass + class LiteralHelptext: + x: Literal[1, 2, 3] + """A number.""" + + helptext = get_helptext_with_checks(LiteralHelptext) + assert "--x {1,2,3}" in helptext + assert "A number. (required)" in helptext + + +def test_optional_literal_helptext() -> None: + @dataclasses.dataclass + class OptionalLiteralHelptext: + x: Optional[Literal[1, 2, 3]] = None + """A number.""" + + helptext = get_helptext_with_checks(OptionalLiteralHelptext) + assert "--x {None,1,2,3}" in helptext + assert "A number. (default: None)" in helptext + + +def test_multiple_subparsers_helptext() -> None: + @dataclasses.dataclass + class Subcommand1: + """2% milk.""" # % symbol is prone to bugs in argparse. + + x: int = 0 + + @dataclasses.dataclass + class Subcommand2: + y: int = 1 + + @dataclasses.dataclass + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: Subcommand1 | Subcommand2 | Subcommand3 + # Field b description. + b: Subcommand1 | Subcommand2 | Subcommand3 + # Field c description. + c: Subcommand1 | Subcommand2 | Subcommand3 = dataclasses.field( + default_factory=Subcommand3 + ) + + d: bool = False + + helptext = get_helptext_with_checks(MultipleSubparsers) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + # Not enough args for usage shortening to kick in. + assert "[OPTIONS]" not in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, args=["a:subcommand1", "b:subcommand1", "--help"] + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand3)" in helptext + + # Not enough args for usage shortening to kick in. + assert "[OPTIONS]" not in helptext + assert "[B:SUBCOMMAND1 OPTIONS]" not in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + +def test_multiple_subparsers_helptext_shortened_usage() -> None: + @dataclasses.dataclass + class Subcommand1: + """2% milk.""" # % symbol is prone to bugs in argparse. + + a: int = 0 + b: int = 0 + c: int = 0 + d: int = 0 + e: int = 0 + + @dataclasses.dataclass + class Subcommand2: + a: int = 0 + b: int = 0 + c: int = 0 + d: int = 0 + e: int = 0 + + @dataclasses.dataclass + class Subcommand3: + a: int = 0 + b: int = 0 + c: int = 0 + d: int = 0 + e: int = 0 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: Subcommand1 | Subcommand2 | Subcommand3 + # Field b description. + b: Subcommand1 | Subcommand2 | Subcommand3 + # Field c description. + c: Subcommand1 | Subcommand2 | Subcommand3 = dataclasses.field( + default_factory=Subcommand3 + ) + + d: bool = False + f: bool = False + g: bool = False + h: bool = False + i: bool = False + + helptext = get_helptext_with_checks(MultipleSubparsers) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + assert "[OPTIONS]" in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, args=["a:subcommand1", "b:subcommand1", "--help"] + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand3)" in helptext + + assert "[OPTIONS]" not in helptext + assert "[B:SUBCOMMAND1 OPTIONS]" in helptext + assert "[B:SUBCOMMAND2 OPTIONS]" not in helptext + + +def test_optional_helptext() -> None: + @dataclasses.dataclass + class OptionalHelptext: + """This docstring should be printed as a description. 2% milk.""" + + x: Optional[int] # Documentation 1 + + # Documentation 2 + y: List[Optional[int]] + + z: Optional[int] = 3 + """Documentation 3""" + + helptext = get_helptext_with_checks(OptionalHelptext) + assert cast(str, cast(str, OptionalHelptext.__doc__)[:20]) in helptext + assert "2% milk" in helptext + assert "--x {None}|INT" in helptext + assert "--y [{None}|INT [{None}|INT ...]]" in helptext + assert "[--z {None}|INT]" in helptext + + +def test_metavar_0() -> None: + def main(x: Literal[0, 1, 2, 3] | Tuple[int, int]) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3}|{INT INT}" in helptext + + +def test_metavar_1() -> None: + def main( + x: Literal[0, 1, 2, 3] | Literal["hey,there", "hello"] | List[int], + ) -> None: + pass + + # The comma formatting is unfortunate, but matches argparse's default behavior. + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3,hey,there,hello}|{[INT [INT ...]]}" in helptext + + +def test_metavar_2() -> None: + def main( + x: Tuple[ + Literal[0, 1, 2, 3], + int | str, + ], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3} INT|STR" in helptext + + +def test_metavar_3() -> None: + def main( + x: Literal[0, 1, 2, 3] | Tuple[int, int] | Tuple[str], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3}|{INT INT}|STR" in helptext + + +def test_metavar_4() -> None: + def main( + x: Literal[0, 1, 2, 3] | Tuple[int, int] | Tuple[str, str, str] | Literal[True], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x {0,1,2,3}|{INT INT}|{STR STR STR}|{True}" in helptext + + +def test_metavar_5() -> None: + def main( + x: List[Tuple[int, int] | Tuple[str, str]] = [(1, 1), (2, 2)], + ) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "[--x [{INT INT}|{STR STR} [{INT INT}|{STR STR} ...]]]" in helptext + + +def test_metavar_6() -> None: + def main(x: Dict[Tuple[int, int] | Tuple[str, str], Tuple[int, int]]) -> dict: + return x + + helptext = get_helptext_with_checks(main) + assert ( + "--x [{INT INT}|{STR STR} INT INT [{INT INT}|{STR STR} INT INT ...]]" + in helptext + ) + + +def test_comment_in_subclass_list() -> None: + @dataclasses.dataclass + class Something( + # This text should not show up in the helptext anywhere. + object, + ): + a: int + + # But this text should! + b: int + + helptext = get_helptext_with_checks(Something) + assert "This text should not" not in helptext + assert "But this text should!" in helptext + + +def test_unparsable() -> None: + class Struct: + a: int = 5 + b: str = "7" + + def main(x: Any = Struct()): + pass + + helptext = get_helptext_with_checks(main) + assert "--x {fixed}" in helptext + + def main2(x: Callable = nn.ReLU): + pass + + helptext = get_helptext_with_checks(main2) + assert "--x {fixed}" in helptext + assert "(fixed to:" in helptext + assert "torch" in helptext + + +def test_pathlike() -> None: + def main(x: os.PathLike) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--x PATH " in helptext + + +def test_nested_bool() -> None: + @dataclasses.dataclass + class Child: + x: bool = False + + def main(child: Child) -> None: + pass + + helptext = get_helptext_with_checks(main) + assert "--child.x | --child.no-x" in helptext + + +def test_multiple_subparsers_helptext_hyphens() -> None: + @dataclasses.dataclass + class SubcommandOne: + """2% milk.""" # % symbol is prone to bugs in argparse. + + arg_x: int = 0 + arg_flag: bool = False + + @dataclasses.dataclass + class SubcommandTwo: + arg_y: int = 1 + + @dataclasses.dataclass + class SubcommandThree: + arg_z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: SubcommandOne | SubcommandTwo | SubcommandThree + # Field b description. + b: SubcommandOne | SubcommandTwo | SubcommandThree + # Field c description. + c: SubcommandOne | SubcommandTwo | SubcommandThree = dataclasses.field( + default_factory=SubcommandThree + ) + + helptext = get_helptext_with_checks(MultipleSubparsers) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, args=["a:subcommand-one", "b:subcommand-one", "--help"] + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand-three)" in helptext + assert "--b.arg-x" in helptext + assert "--b.no-arg-flag" in helptext + assert "--b.arg-flag" in helptext + + +def test_multiple_subparsers_helptext_underscores() -> None: + @dataclasses.dataclass + class SubcommandOne: + """2% milk.""" # % symbol is prone to bugs in argparse. + + arg_x: int = 0 + arg_flag: bool = False + + @dataclasses.dataclass + class SubcommandTwo: + arg_y: int = 1 + + @dataclasses.dataclass + class SubcommandThree: + arg_z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + # Field a description. + a: SubcommandOne | SubcommandTwo | SubcommandThree + # Field b description. + b: SubcommandOne | SubcommandTwo | SubcommandThree + # Field c description. + c: SubcommandOne | SubcommandTwo | SubcommandThree = dataclasses.field( + default_factory=SubcommandThree + ) + + helptext = get_helptext_with_checks(MultipleSubparsers, use_underscores=True) + + assert "2% milk." in helptext + assert "Field a description." in helptext + assert "Field b description." not in helptext + assert "Field c description." not in helptext + + helptext = get_helptext_with_checks( + MultipleSubparsers, + args=["a:subcommand_one", "b:subcommand_one", "--help"], + use_underscores=True, + ) + + assert "2% milk." in helptext + assert "Field a description." not in helptext + assert "Field b description." not in helptext + assert "Field c description." in helptext + assert "(default: c:subcommand_three)" in helptext + assert "--b.arg_x" in helptext + assert "--b.no_arg_flag" in helptext + assert "--b.arg_flag" in helptext + + +def test_subparsers_wrapping() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CheckoutCompletion: + """Help message.""" + + y: int + + help = get_helptext_with_checks(A | CheckoutCompletion) # type: ignore + assert help.count("checkout-completion") == 3 + + +def test_subparsers_wrapping1() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CheckoutCompletio: + """Help message.""" + + y: int + + help = get_helptext_with_checks(A | CheckoutCompletio) # type: ignore + assert help.count("checkout-completio") == 3 + + +def test_subparsers_wrapping2() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CheckoutCompletionn: + """Help message.""" + + y: int + + help = get_helptext_with_checks(A | CheckoutCompletionn) # type: ignore + assert help.count("checkout-completionn") == 3 + + +def test_subparsers_wrapping3() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: int + + @dataclasses.dataclass + class CmdCheckout012: + """Help message.""" + + y: int + + help = get_helptext_with_checks(A | CmdCheckout012) # type: ignore + assert help.count("cmd-checkout012") == 3 + + +def test_tuple_default() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: Tuple[str, str] = ("hello", "world") + + help = get_helptext_with_checks(A) + assert "STR STR" in help + assert "hello world" in help + assert "('hello', 'world')" not in help + + +def test_argconf_constructor() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: Annotated[ + Tuple[str, str], tyro.conf.arg(constructor=lambda x: ("a", "b")) + ] = ("hello", "world") + + help = get_helptext_with_checks(A) + assert "STR STR" not in help + # Unlike case above, should not be converted to 'hello world'. + assert "('hello', 'world')" in help # JSON special case. + + +def test_argconf_constructor_json_special_case() -> None: + @dataclasses.dataclass + class A: + """Help message.""" + + x: Annotated[ + Tuple[str, str], tyro.conf.arg(constructor=json.loads, metavar="JSON") + ] = ("hello", "world") + y: Annotated[ + Tuple[int, int], tyro.conf.arg(constructor=json.loads, metavar="JSON") + ] = (3, 5) + + help = get_helptext_with_checks(A) + assert "STR STR" not in help + assert "JSON" in help + assert '["hello", "world"]' in help # JSON special case. + assert "[3, 5]" in help, help # JSON special case. + + +def test_optional_group() -> None: + def f( + x: Annotated[ + Tuple[str, str], tyro.conf.arg(constructor=json.loads, metavar="JSON") + ] = ("hello", "world"), + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f, default=3) + assert 'default if used: ["hello", "world"]' in help + assert "default if used: 3" in help + + +def test_append_fixed() -> None: + def f( + x: tyro.conf.UseAppendAction[Tuple[str, str]], + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f) + assert "repeatable" not in help, help + + +def test_append_good() -> None: + def f( + x: tyro.conf.UseAppendAction[Tuple[str, ...]], + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f) + assert "repeatable" in help, help + + +def test_append_with_default() -> None: + def f( + x: tyro.conf.UseAppendAction[Tuple[str, ...]] = ("hello world", "hello"), + y: int = 3, + ) -> int: + del x, y + return 5 + + help = get_helptext_with_checks(f) + assert "repeatable, appends to: 'hello world' hello" in help, help + + +def test_typeddict_exclude() -> None: + class Special(TypedDict): + x: NotRequired[int] + + help = get_helptext_with_checks(Special) + assert "unset by default" in help, help diff --git a/tests/test_py311_generated/test_initvar_min_py38_generated.py b/tests/test_py311_generated/test_initvar_min_py38_generated.py new file mode 100644 index 000000000..721462716 --- /dev/null +++ b/tests/test_py311_generated/test_initvar_min_py38_generated.py @@ -0,0 +1,27 @@ +import dataclasses + +import tyro + + +def test_dataclass_init_var() -> None: + @dataclasses.dataclass + class DataclassWithInitVar: + i: dataclasses.InitVar[int] + x: str + + def __post_init__(self, i: int) -> None: + self.x += str(i) + + # We can directly pass a dataclass to `tyro.cli()`: + assert ( + tyro.cli( + DataclassWithInitVar, + args=[ + "--i", + "5", + "--x", + "5", + ], + ).x + == "55" + ) diff --git a/tests/test_py311_generated/test_is_nested_type_generated.py b/tests/test_py311_generated/test_is_nested_type_generated.py new file mode 100644 index 000000000..a586756f4 --- /dev/null +++ b/tests/test_py311_generated/test_is_nested_type_generated.py @@ -0,0 +1,52 @@ +import dataclasses +import pathlib +from typing import Any, Dict, List, Tuple + +from tyro._fields import MISSING_NONPROP, is_nested_type + + +def test_is_nested_type_simple(): + assert not is_nested_type(int, MISSING_NONPROP) + assert not is_nested_type(bool, MISSING_NONPROP) + assert not is_nested_type(str, MISSING_NONPROP) + assert not is_nested_type(pathlib.Path, MISSING_NONPROP) + + +def test_is_nested_type_containers(): + assert not is_nested_type(List[int], MISSING_NONPROP) + assert not is_nested_type(List[bool], MISSING_NONPROP) + assert not is_nested_type(List[str], MISSING_NONPROP) + assert not is_nested_type(List[pathlib.Path], MISSING_NONPROP) + + +@dataclasses.dataclass +class Color: + r: int + g: int + b: int + + +def test_is_nested_type_actually_nested(): + assert is_nested_type(Color, Color(255, 0, 0)) + + +def test_is_nested_type_actually_nested_narrowing(): + assert is_nested_type(Any, Color(255, 0, 0)) + assert is_nested_type(object, Color(255, 0, 0)) + assert not is_nested_type(int, Color(255, 0, 0)) + + +def test_is_nested_type_actually_nested_in_container(): + assert is_nested_type(Tuple[Color, Color], MISSING_NONPROP) + assert is_nested_type(Tuple[object, ...], (Color(255, 0, 0),)) + assert is_nested_type(Tuple[Any, ...], (Color(255, 0, 0),)) + assert is_nested_type(tuple, (Color(255, 0, 0),)) + assert not is_nested_type(tuple, (1, 2, 3)) + assert is_nested_type(tuple, (1, Color(255, 0, 0), 3)) + assert is_nested_type(List[Any], [Color(255, 0, 0)]) + + +def test_nested_dict(): + assert is_nested_type(Dict[str, int], {"x": 5}) + assert is_nested_type(dict, {"x": 5}) + assert is_nested_type(Any, {"x": 5}) diff --git a/tests/test_py311_generated/test_missing_generated.py b/tests/test_py311_generated/test_missing_generated.py new file mode 100644 index 000000000..3f7f9d428 --- /dev/null +++ b/tests/test_py311_generated/test_missing_generated.py @@ -0,0 +1,82 @@ +"""Tests for tyro.MISSING.""" + +import contextlib +import dataclasses +import io +from typing import Tuple + +import pytest + +import tyro + + +def test_missing() -> None: + def main(a: int = 5, b: int = tyro.MISSING, c: int = 3) -> Tuple[int, int, int]: + return a, b, c + + target = io.StringIO() + with pytest.raises(SystemExit), contextlib.redirect_stderr(target): + tyro.cli(main, args=[]) + message = target.getvalue() + assert "Required options" in message + assert "Argument helptext" in message + assert "--a INT" not in message + assert "--b INT" in message + assert "(required)" in message + assert tyro.cli(main, args=["--b", "7"]) == (5, 7, 3) + + +def test_missing_dataclass() -> None: + @dataclasses.dataclass + class Args2: + a: int = 5 + b: int = tyro.MISSING + c: int = 3 + + with pytest.raises(SystemExit): + tyro.cli(Args2, args=[]) + assert tyro.cli(Args2, args=["--b", "7"]) == Args2(5, 7, 3) + + +def test_missing_default() -> None: + @dataclasses.dataclass + class Args2: + a: int + b: int + c: int + + with pytest.raises(SystemExit): + tyro.cli( + Args2, + args=[], + default=Args2(5, tyro.MISSING, 3), + ) + assert tyro.cli( + Args2, + args=["--b", "7"], + default=Args2(5, tyro.MISSING, 3), + ) == Args2(5, 7, 3) + + +def test_missing_nested_default() -> None: + @dataclasses.dataclass + class Child: + a: int = 5 + b: int = 3 + c: int = 7 + + @dataclasses.dataclass + class Parent: + child: Child + + with pytest.raises(SystemExit): + tyro.cli( + Parent, + args=[], + default=Parent(child=tyro.MISSING), + ) + assert tyro.cli( + Parent, + args=["--child.a", "5", "--child.b", "7", "--child.c", "3"], + default=Parent(child=tyro.MISSING), + ) == Parent(Child(5, 7, 3)) diff --git a/tests/test_py311_generated/test_missing_optional_packages_generated.py b/tests/test_py311_generated/test_missing_optional_packages_generated.py new file mode 100644 index 000000000..084e6e7f9 --- /dev/null +++ b/tests/test_py311_generated/test_missing_optional_packages_generated.py @@ -0,0 +1,33 @@ +import builtins +import dataclasses +from typing import Any + +import pytest + + +# https://stackoverflow.com/questions/60227582/making-a-python-test-think-an-installed-package-is-not-available +@pytest.fixture +def hide_optional_packages(monkeypatch: Any) -> None: + import_orig = builtins.__import__ + + def mocked_import(name, *args, **kwargs): + if name in ("attr", "pydantic", "flax", "omegaconf"): + raise ImportError() + return import_orig(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", mocked_import) + + +def test_missing_optional_packages(hide_optional_packages: Any) -> None: + @dataclasses.dataclass + class Extra: + a: int = 3 + b: str = "5" + + def main(x: int, extra: Extra) -> int: + del extra + return x + + import tyro + + assert tyro.cli(main, args=["--x", "5"]) == 5 diff --git a/tests/test_py311_generated/test_mixed_unions_generated.py b/tests/test_py311_generated/test_mixed_unions_generated.py new file mode 100644 index 000000000..c361de491 --- /dev/null +++ b/tests/test_py311_generated/test_mixed_unions_generated.py @@ -0,0 +1,97 @@ +"""Tests for unsupported union types. + +Unions like `int | str` or `SomeDataclassA | SomeDataclassB` are OK (note that the latter +will produce a pair of subcommands); when we write things like `int | SomeDataclassA` +handling gets more complicated but should still be supported! +""" + +import dataclasses +from typing import Any, Dict, List, Tuple + +import pytest + +import tyro + + +def test_subparser_strip_non_nested() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + # Note that we add [int, str] to the annotation here... this should be ignored. + bc: int | str | DefaultHTTPServer | DefaultSMTPServer = dataclasses.field( + default_factory=lambda: DefaultHTTPServer(5) + ) + + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "5"] + ) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-smtp-server", "--bc.z", "3"] + ) == DefaultSubparser(x=1, bc=DefaultSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "8"] + ) + == tyro.cli( + DefaultSubparser, + args=[], + default=DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)), + ) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_strip_nested() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + # Note that we add [int, str] to the annotation here... this should be ignored. + bc: int | str | DefaultHTTPServer | DefaultSMTPServer = 5 + + assert ( + tyro.cli(DefaultSubparser, args=["--x", "1", "bc:int", "5"]) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=5) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:str", "five"] + ) == DefaultSubparser(x=1, bc="five") + + +def test_with_fancy_types() -> None: + @dataclasses.dataclass + class Args: + y: int + + def main(x: Tuple[int, ...] | List[str] | Args | Dict[str, int]) -> Any: + return x + + assert tyro.cli(main, args="x:tuple-int-ellipsis 1 2 3".split(" ")) == (1, 2, 3) + assert tyro.cli(main, args="x:list-str 1 2 3".split(" ")) == ["1", "2", "3"] + assert tyro.cli(main, args="x:args --x.y 5".split(" ")) == Args(5) + assert tyro.cli(main, args="x:dict-str-int 1 2 3 4".split(" ")) == {"1": 2, "3": 4} diff --git a/tests/test_py311_generated/test_nested_generated.py b/tests/test_py311_generated/test_nested_generated.py new file mode 100644 index 000000000..5fb5b31be --- /dev/null +++ b/tests/test_py311_generated/test_nested_generated.py @@ -0,0 +1,1199 @@ +import dataclasses +from typing import ( + Annotated, + Any, + Final, + Generic, + Literal, + Mapping, + NewType, + Optional, + Tuple, + TypeVar, +) + +import pytest +from frozendict import frozendict # type: ignore +from helptext_utils import get_helptext_with_checks + +import tyro + + +def test_nested() -> None: + @dataclasses.dataclass + class B: + y: int + + @dataclasses.dataclass + class Nested: + x: int + b: B + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_annotated() -> None: + @dataclasses.dataclass + class B: + y: int + + @dataclasses.dataclass + class Nested: + x: int + b: Annotated[B, "this should be ignored"] + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_final() -> None: + @dataclasses.dataclass + class B: + y: int + + @dataclasses.dataclass + class Nested: + x: int + b: Final[B] # type: ignore + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_accidental_underscores() -> None: + @dataclasses.dataclass + class B: + arg_name: str + + @dataclasses.dataclass + class Nested: + x: int + child_struct: B + + assert ( + tyro.cli(Nested, args=["--x", "1", "--child-struct.arg-name", "three_five"]) + == tyro.cli(Nested, args=["--x", "1", "--child_struct.arg_name", "three_five"]) + == tyro.cli(Nested, args=["--x", "1", "--child_struct.arg-name", "three_five"]) + == tyro.cli(Nested, args=["--x", "1", "--child_struct.arg_name=three_five"]) + == Nested(x=1, child_struct=B(arg_name="three_five")) + ) + with pytest.raises(SystemExit): + tyro.cli(Nested, args=["--x", "1"]) + + +def test_nested_default() -> None: + @dataclasses.dataclass(frozen=True) + class B: + y: int = 1 + + @dataclasses.dataclass + class Nested: + x: int = 2 + b: B = B() + + assert tyro.cli(Nested, args=[], default=Nested(x=1, b=B(y=2))) == Nested( + x=1, b=B(y=2) + ) + + +def test_nested_default_alternate() -> None: + @dataclasses.dataclass + class B: + y: int = 3 + + @dataclasses.dataclass + class Nested: + x: int + b: B + + assert ( + Nested(x=1, b=B(y=3)) + == tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) + == tyro.cli(Nested, args=[], default=Nested(x=1, b=B(y=3))) + ) + assert tyro.cli(Nested, args=["--x", "1"]) == Nested(x=1, b=B(y=3)) + + +def test_default_nested() -> None: + @dataclasses.dataclass(frozen=True) + class B: + y: int = 3 + + @dataclasses.dataclass(frozen=True) + class Nested: + x: int + b: B = B(y=5) + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + assert tyro.cli(Nested, args=["--x", "1"]) == Nested(x=1, b=B(y=5)) + + +def test_double_default_nested() -> None: + @dataclasses.dataclass(frozen=True) + class Child: + y: int + + @dataclasses.dataclass(frozen=True) + class Parent: + c: Child + + @dataclasses.dataclass(frozen=True) + class Grandparent: + x: int + b: Parent = Parent(Child(y=5)) + + assert tyro.cli(Grandparent, args=["--x", "1", "--b.c.y", "3"]) == Grandparent( + x=1, b=Parent(Child(y=3)) + ) + assert tyro.cli(Grandparent, args=["--x", "1"]) == Grandparent( + x=1, b=Parent(Child(y=5)) + ) + + +def test_default_factory_nested() -> None: + @dataclasses.dataclass + class B: + y: int = 3 + + @dataclasses.dataclass + class Nested: + x: int + b: B = dataclasses.field(default_factory=lambda: B(y=5)) + + assert tyro.cli(Nested, args=["--x", "1", "--b.y", "3"]) == Nested(x=1, b=B(y=3)) + assert tyro.cli(Nested, args=["--x", "1"]) == Nested(x=1, b=B(y=5)) + + +def test_optional_nested() -> None: + @dataclasses.dataclass + class OptionalNestedChild: + y: int + z: int + + @dataclasses.dataclass + class OptionalNested: + x: int + b: Optional[OptionalNestedChild] = None + + assert tyro.cli(OptionalNested, args=["--x", "1"]) == OptionalNested(x=1, b=None) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, args=["--x", "1", "b:optional-nested-child", "--b.y", "3"] + ) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, args=["--x", "1", "b:optional-nested-child", "--b.z", "3"] + ) + + assert tyro.cli( + OptionalNested, + args=["--x", "1", "b:optional-nested-child", "--b.y", "2", "--b.z", "3"], + ) == OptionalNested(x=1, b=OptionalNestedChild(y=2, z=3)) + + +def test_optional_nested_newtype() -> None: + @dataclasses.dataclass + class OptionalNestedChild: + y: int + z: int + + SpecialOptionalNestedChild = NewType( + "SpecialOptionalNestedChild", OptionalNestedChild + ) + + @dataclasses.dataclass + class OptionalNested: + x: int + b: Optional[SpecialOptionalNestedChild] = None + + assert tyro.cli(OptionalNested, args=["--x", "1"]) == OptionalNested(x=1, b=None) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, + args=["--x", "1", "b:special-optional-nested-child", "--b.y", "3"], + ) + with pytest.raises(SystemExit): + tyro.cli( + OptionalNested, + args=["--x", "1", "b:special-optional-nested-child", "--b.z", "3"], + ) + + assert tyro.cli( + OptionalNested, + args=[ + "--x", + "1", + "b:special-optional-nested-child", + "--b.y", + "2", + "--b.z", + "3", + ], + ) == OptionalNested( + x=1, b=SpecialOptionalNestedChild(OptionalNestedChild(y=2, z=3)) + ) + + +def test_optional_nested_multiple() -> None: + """Adapted from: https://github.com/brentyi/tyro/issues/60""" + + @dataclasses.dataclass(frozen=True) + class OutputHeadSettings: + number_of_outputs: int = 1 + + @dataclasses.dataclass(frozen=True) + class OptimizerSettings: + name: str = "adam" + + @dataclasses.dataclass(frozen=True) + class ModelSettings: + output_head_settings: Optional[OutputHeadSettings] = None + optimizer_settings: Optional[OptimizerSettings] = None + + assert tyro.cli( + ModelSettings, + args="output-head-settings:None optimizer-settings:None".split(" "), + ) == ModelSettings(None, None) + + with pytest.raises(SystemExit): + # Order cannot be flipped, unfortunately. + tyro.cli( + ModelSettings, + args="optimizer-settings:None output-head-settings:None".split(" "), + ) + + assert tyro.cli( + ModelSettings, + args="output-head-settings:output-head-settings optimizer-settings:None".split( + " " + ), + ) == ModelSettings(OutputHeadSettings(1), None) + + assert tyro.cli( + ModelSettings, + args=( + "output-head-settings:output-head-settings" + " --output-head-settings.number-of-outputs 5 optimizer-settings:None".split( + " " + ) + ), + ) == ModelSettings(OutputHeadSettings(5), None) + + assert tyro.cli( + tyro.conf.OmitSubcommandPrefixes[ + tyro.conf.ConsolidateSubcommandArgs[ModelSettings] + ], + args=("output-head-settings None" " --number-of-outputs 5".split(" ")), + ) == ModelSettings(OutputHeadSettings(5), None) + + assert tyro.cli( + tyro.conf.OmitSubcommandPrefixes[ + tyro.conf.ConsolidateSubcommandArgs[ModelSettings] + ], + args=( + "output-head-settings" + " optimizer-settings --name sgd --number-of-outputs 5".split(" ") + ), + ) == ModelSettings(OutputHeadSettings(5), OptimizerSettings("sgd")) + + +def test_subparser() -> None: + @dataclasses.dataclass + class HTTPServer: + y: int + + @dataclasses.dataclass + class SMTPServer: + z: int + + @dataclasses.dataclass + class Subparser: + x: int + bc: HTTPServer | SMTPServer + + assert tyro.cli( + Subparser, args=["--x", "1", "bc:http-server", "--bc.y", "3"] + ) == Subparser(x=1, bc=HTTPServer(y=3)) + assert tyro.cli( + Subparser, args=["--x", "1", "bc:smtp-server", "--bc.z", "3"] + ) == Subparser(x=1, bc=SMTPServer(z=3)) + + with pytest.raises(SystemExit): + # Missing subcommand. + tyro.cli(Subparser, args=["--x", "1"]) + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli(Subparser, args=["--x", "1", "bc:http-server", "--bc.z", "3"]) + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli(Subparser, args=["--x", "1", "bc:smtp-server", "--bc.y", "3"]) + + +def test_subparser_root() -> None: + @dataclasses.dataclass + class HTTPServer: + y: int + + @dataclasses.dataclass + class SMTPServer: + z: int + + @dataclasses.dataclass + class Subparser: + x: int + bc: HTTPServer | SMTPServer + + assert tyro.cli( + HTTPServer | SMTPServer, + args=["http-server", "--y", "3"], # type: ignore + ) == HTTPServer(y=3) + + +def test_subparser_with_default() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + bc: DefaultHTTPServer | DefaultSMTPServer = dataclasses.field( + default_factory=lambda: DefaultHTTPServer(5) + ) + + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "5"] + ) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-smtp-server", "--bc.z", "3"] + ) == DefaultSubparser(x=1, bc=DefaultSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "8"] + ) + == tyro.cli( + DefaultSubparser, + args=[], + default=DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)), + ) + == DefaultSubparser(x=1, bc=DefaultHTTPServer(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_with_default_and_newtype() -> None: + @dataclasses.dataclass + class DefaultHTTPServer_: + y: int + + DefaultHTTPServer__ = NewType("DefaultHTTPServer__", DefaultHTTPServer_) + DefaultHTTPServer = NewType("DefaultHTTPServer", DefaultHTTPServer__) # type: ignore + # ^nesting NewType is not technically allowed and pyright will complain, + # but we should try to be robust to it anyways. + + def make_http_server(y: int) -> DefaultHTTPServer: + return DefaultHTTPServer(DefaultHTTPServer__(DefaultHTTPServer_(y))) + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + bc: DefaultHTTPServer | DefaultSMTPServer = dataclasses.field( + default_factory=lambda: make_http_server(5) + ) + + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "5"] + ) + == tyro.cli(DefaultSubparser, args=["--x", "1"]) + == DefaultSubparser(x=1, bc=make_http_server(y=5)) + ) + assert tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-smtp-server", "--bc.z", "3"] + ) == DefaultSubparser(x=1, bc=DefaultSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultSubparser, args=["--x", "1", "bc:default-http-server", "--bc.y", "8"] + ) + == tyro.cli( + DefaultSubparser, + args=[], + default=DefaultSubparser(x=1, bc=make_http_server(y=8)), + ) + == DefaultSubparser(x=1, bc=make_http_server(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_with_default_alternate() -> None: + @dataclasses.dataclass + class DefaultInstanceHTTPServer: + y: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSMTPServer: + z: int = 0 + + @dataclasses.dataclass + class DefaultInstanceSubparser: + x: int + bc: DefaultInstanceHTTPServer | DefaultInstanceSMTPServer + + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "5"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=[], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) + == tyro.cli( + DefaultInstanceSubparser, + args=["bc:default-instance-http-server"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)) + ) + assert tyro.cli( + DefaultInstanceSubparser, + args=["bc:default-instance-smtp-server", "--bc.z", "3"], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=5)), + ) == DefaultInstanceSubparser(x=1, bc=DefaultInstanceSMTPServer(z=3)) + assert ( + tyro.cli( + DefaultInstanceSubparser, + args=["--x", "1", "bc:default-instance-http-server", "--bc.y", "8"], + ) + == tyro.cli( + DefaultInstanceSubparser, + args=[], + default=DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)), + ) + == DefaultInstanceSubparser(x=1, bc=DefaultInstanceHTTPServer(y=8)) + ) + + with pytest.raises(SystemExit): + tyro.cli(DefaultInstanceSubparser, args=["--x", "1", "b", "--bc.z", "3"]) + with pytest.raises(SystemExit): + tyro.cli(DefaultInstanceSubparser, args=["--x", "1", "c", "--bc.y", "3"]) + + +def test_subparser_with_default_bad() -> None: + @dataclasses.dataclass + class DefaultHTTPServer: + y: int + + @dataclasses.dataclass + class DefaultSMTPServer: + z: int + + @dataclasses.dataclass + class DefaultSubparser: + x: int + bc: DefaultHTTPServer | DefaultSMTPServer = dataclasses.field( + default_factory=lambda: 5 # type: ignore + ) + + # Tolerate bad static types: https://github.com/brentyi/tyro/issues/20 + # Should give us a bunch of warnings! + with pytest.warns(UserWarning): + assert tyro.cli(DefaultSubparser, args=["--x", "1"]) == DefaultSubparser( + 1, + 5, # type: ignore + ) + + +def test_subparser_with_default_bad_alt() -> None: + @dataclasses.dataclass + class A: + a: int + + @tyro.conf.configure(tyro.conf.subcommand(name="bbbb")) + @dataclasses.dataclass + class B: + b: int + + @dataclasses.dataclass + class C: + c: int + + with pytest.warns(UserWarning): + assert tyro.cli( + A | Annotated[B, None], # type: ignore + default=C(3), + args=["c", "--c", "2"], + ) == C(2) + + +def test_optional_subparser() -> None: + @dataclasses.dataclass + class OptionalHTTPServer: + y: int + + @dataclasses.dataclass + class OptionalSMTPServer: + z: int + + @dataclasses.dataclass + class OptionalSubparser: + x: int + bc: Optional[OptionalHTTPServer | OptionalSMTPServer] + + assert tyro.cli( + OptionalSubparser, args=["--x", "1", "bc:optional-http-server", "--bc.y", "3"] + ) == OptionalSubparser(x=1, bc=OptionalHTTPServer(y=3)) + assert tyro.cli( + OptionalSubparser, args=["--x", "1", "bc:optional-smtp-server", "--bc.z", "3"] + ) == OptionalSubparser(x=1, bc=OptionalSMTPServer(z=3)) + assert tyro.cli( + OptionalSubparser, args=["--x", "1", "bc:None"] + ) == OptionalSubparser(x=1, bc=None) + + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli( + OptionalSubparser, + args=["--x", "1", "bc:optional-http-server", "--bc.z", "3"], + ) + with pytest.raises(SystemExit): + # Wrong field. + tyro.cli( + OptionalSubparser, + args=["--x", "1", "bc:optional-smtp-server", "--bc.y", "3"], + ) + + +def test_post_init_default() -> None: + @dataclasses.dataclass + class DataclassWithDynamicDefault: + x: int = 3 + y: Optional[int] = None + + def __post_init__(self): + # If unspecified, set y = x. + if self.y is None: + self.y = self.x + + @dataclasses.dataclass + class NoDefaultPostInitArgs: + inner: DataclassWithDynamicDefault + + @dataclasses.dataclass + class DefaultFactoryPostInitArgs: + inner: DataclassWithDynamicDefault = dataclasses.field( + default_factory=DataclassWithDynamicDefault + ) + + assert ( + tyro.cli(NoDefaultPostInitArgs, args=["--inner.x", "5"]).inner + == tyro.cli(DefaultFactoryPostInitArgs, args=["--inner.x", "5"]).inner + == DataclassWithDynamicDefault(x=5, y=5) + ) + + +def test_multiple_subparsers() -> None: + @dataclasses.dataclass + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass + class Subcommand2: + y: int = 1 + + @dataclasses.dataclass + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + a: Subcommand1 | Subcommand2 | Subcommand3 + b: Subcommand1 | Subcommand2 | Subcommand3 + c: Subcommand1 | Subcommand2 | Subcommand3 + + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=[]) + + assert tyro.cli( + MultipleSubparsers, args="a:subcommand1 b:subcommand2 c:subcommand3".split(" ") + ) == MultipleSubparsers(Subcommand1(), Subcommand2(), Subcommand3()) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand1 --a.x 5 b:subcommand2 --b.y 7 c:subcommand3 --c.z 3".split( + " " + ), + ) == MultipleSubparsers(Subcommand1(x=5), Subcommand2(y=7), Subcommand3(z=3)) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand2 --a.y 5 b:subcommand1 --b.x 7 c:subcommand3 --c.z 3".split( + " " + ), + ) == MultipleSubparsers(Subcommand2(y=5), Subcommand1(x=7), Subcommand3(z=3)) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand3 --a.z 5 b:subcommand1 --b.x 7 c:subcommand3 --c.z 3".split( + " " + ), + ) == MultipleSubparsers(Subcommand3(z=5), Subcommand1(x=7), Subcommand3(z=3)) + + +def test_multiple_subparsers_with_default() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: int = 1 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass + class MultipleSubparsers: + a: Subcommand1 | Subcommand2 | Subcommand3 = Subcommand1(tyro.MISSING) + b: Subcommand1 | Subcommand2 | Subcommand3 = Subcommand2(7) + c: Subcommand1 | Subcommand2 | Subcommand3 = Subcommand3(3) + + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=[], + ) + + assert tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "--a.x", "5"], + ) == MultipleSubparsers(Subcommand1(x=5), Subcommand2(y=7), Subcommand3(z=3)) + + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand1 --a.x 3".split(" "), + ) == MultipleSubparsers(Subcommand1(x=3), Subcommand2(y=7), Subcommand3(z=3)) + + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=[], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=[ + "a:subcommand1", + ], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + with pytest.raises(SystemExit): + tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2", "c:subcommand3"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) + assert tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2", "c:subcommand3", "--c.z", "3"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) == MultipleSubparsers(Subcommand1(x=0), Subcommand2(y=1), Subcommand3(z=3)) + assert tyro.cli( + MultipleSubparsers, + args=["a:subcommand1", "b:subcommand2", "c:subcommand2"], + default=MultipleSubparsers( + Subcommand1(), + Subcommand2(), + Subcommand3(tyro.MISSING), + ), + ) == MultipleSubparsers(Subcommand1(x=0), Subcommand2(y=1), Subcommand2(y=1)) + + +def test_nested_subparsers_with_default() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: Subcommand1 | Subcommand3 + + @dataclasses.dataclass(frozen=True) + class MultipleSubparsers: + a: Subcommand1 | Subcommand2 = Subcommand2(Subcommand1(tyro.MISSING)) + + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=[]) + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=["a:subcommand2"]) + + assert tyro.cli( + MultipleSubparsers, args="a:subcommand1 --a.x 3".split(" ") + ) == MultipleSubparsers(Subcommand1(3)) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand2 a.y:subcommand3 --a.y.z 2".split(" ") + ) == MultipleSubparsers(Subcommand2(Subcommand3())) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand2 a.y:subcommand3 --a.y.z 7".split(" ") + ) == MultipleSubparsers(Subcommand2(Subcommand3(7))) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand2 a.y:subcommand1 --a.y.x 7".split(" ") + ) == MultipleSubparsers(Subcommand2(Subcommand1(7))) + + +def test_nested_subparsers_multiple() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: Subcommand1 | Subcommand3 + + @dataclasses.dataclass(frozen=True) + class MultipleSubparsers: + a: Subcommand1 | Subcommand2 + b: Subcommand1 | Subcommand2 + + with pytest.raises(SystemExit): + tyro.cli(MultipleSubparsers, args=[]) + assert tyro.cli( + MultipleSubparsers, args="a:subcommand1 b:subcommand1".split(" ") + ) == MultipleSubparsers(Subcommand1(), Subcommand1()) + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand1 b:subcommand2 b.y:subcommand1".split(" "), + ) == MultipleSubparsers(Subcommand1(), Subcommand2(Subcommand1())) + assert tyro.cli( + MultipleSubparsers, + args="a:subcommand2 a.y:subcommand1 b:subcommand2 b.y:subcommand1".split(" "), + ) == MultipleSubparsers(Subcommand2(Subcommand1()), Subcommand2(Subcommand1())) + assert tyro.cli( + MultipleSubparsers, + args=( + "a:subcommand2 a.y:subcommand1 --a.y.x 3 b:subcommand2 b.y:subcommand1" + " --b.y.x 7".split(" ") + ), + ) == MultipleSubparsers(Subcommand2(Subcommand1(3)), Subcommand2(Subcommand1(7))) + + +def test_nested_subparsers_multiple_in_container() -> None: + @dataclasses.dataclass(frozen=True) + class Subcommand1: + x: int = 0 + + @dataclasses.dataclass(frozen=True) + class Subcommand3: + z: int = 2 + + @dataclasses.dataclass(frozen=True) + class Subcommand2: + y: Subcommand1 | Subcommand3 + + @dataclasses.dataclass(frozen=True) + class MultipleSubparsers: + a: Subcommand1 | Subcommand2 + b: Subcommand1 | Subcommand2 + + @dataclasses.dataclass(frozen=True) + class Root: + inner: MultipleSubparsers + + with pytest.raises(SystemExit): + tyro.cli(Root, args=[]) + assert tyro.cli( + Root, args="inner.a:subcommand1 inner.b:subcommand1".split(" ") + ) == Root(MultipleSubparsers(Subcommand1(), Subcommand1())) + assert tyro.cli( + Root, + args="inner.a:subcommand1 inner.b:subcommand2 inner.b.y:subcommand1".split(" "), + ) == Root(MultipleSubparsers(Subcommand1(), Subcommand2(Subcommand1()))) + assert tyro.cli( + Root, + args=( + "inner.a:subcommand2 inner.a.y:subcommand1 inner.b:subcommand2" + " inner.b.y:subcommand1".split(" ") + ), + ) == Root( + MultipleSubparsers(Subcommand2(Subcommand1()), Subcommand2(Subcommand1())) + ) + assert tyro.cli( + Root, + args=( + "inner.a:subcommand2 inner.a.y:subcommand1 --inner.a.y.x 3" + " inner.b:subcommand2 inner.b.y:subcommand1 --inner.b.y.x 7".split(" ") + ), + ) == Root( + MultipleSubparsers(Subcommand2(Subcommand1(3)), Subcommand2(Subcommand1(7))) + ) + + +def test_tuple_nesting() -> None: + @dataclasses.dataclass(frozen=True) + class Color: + r: int + g: int + b: int + + @dataclasses.dataclass(frozen=True) + class Location: + x: float + y: float + z: float + + def main(x: Tuple[Tuple[Color], Location, float]): + return x + + assert tyro.cli( + main, + args=( + "--x.0.0.r 255 --x.0.0.g 0 --x.0.0.b 0 --x.1.x 5.0 --x.1.y 0.0" + " --x.1.z 2.0 --x.2 4.0".split(" ") + ), + ) == ((Color(255, 0, 0),), Location(5.0, 0.0, 2.0), 4.0) + + +def test_tuple_nesting_union() -> None: + @dataclasses.dataclass(frozen=True) + class Color: + r: int + g: int + b: int + + @dataclasses.dataclass(frozen=True) + class Location: + x: float + y: float + z: float + + def main(x: Tuple[Tuple[Color], Location, float] | Tuple[Color, Color]): + return x + + assert tyro.cli( + main, + args=( + "x:tuple-tuple-color-location-float --x.0.0.r 255 --x.0.0.g 0 --x.0.0.b 0" + " --x.1.x 5.0 --x.1.y 0.0 --x.1.z 2.0 --x.2 4.0".split(" ") + ), + ) == ((Color(255, 0, 0),), Location(5.0, 0.0, 2.0), 4.0) + + +def test_generic_subparsers() -> None: + T = TypeVar("T") + + @dataclasses.dataclass + class A(Generic[T]): + x: T + + def main(x: A[int] | A[float]) -> Any: + return x + + assert tyro.cli(main, args="x:a-float --x.x 3.2".split(" ")) == A(3.2) + assert tyro.cli(main, args="x:a-int --x.x 3".split(" ")) == A(3) + + def main_with_default(x: A[int] | A[float] = A(5)) -> Any: + return x + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main_with_default, args=[]) + + +def test_generic_inherited() -> None: + class UnrelatedParentClass: + pass + + T = TypeVar("T") + + @dataclasses.dataclass + class ActualParentClass(Generic[T]): + x: T # Documentation 1 + + # Documentation 2 + y: T + + z: T = 3 # type: ignore + """Documentation 3""" + + @dataclasses.dataclass + class ChildClass(UnrelatedParentClass, ActualParentClass[int]): + pass + + assert tyro.cli( + ChildClass, args=["--x", "1", "--y", "2", "--z", "3"] + ) == ChildClass(x=1, y=2, z=3) + + +def test_subparser_in_nested() -> None: + @dataclasses.dataclass + class A: + a: int + + @dataclasses.dataclass + class B: + b: int + + @dataclasses.dataclass + class Nested2: + subcommand: A | B + + @dataclasses.dataclass + class Nested1: + nested2: Nested2 + + @dataclasses.dataclass + class Parent: + nested1: Nested1 + + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:a --nested1.nested2.subcommand.a 3".split(" "), + ) == Parent(Nested1(Nested2(A(3)))) + assert tyro.cli( + Parent, + args="nested1.nested2.subcommand:b --nested1.nested2.subcommand.b 7".split(" "), + ) == Parent(Nested1(Nested2(B(7)))) + + +def test_frozen_dict() -> None: + def main( + x: Mapping[str, float] = frozendict( # type: ignore + { + "num_epochs": 20, + "batch_size": 64, + } + ), + ): + return x + + assert hash(tyro.cli(main, args="--x.num-epochs 10".split(" "))) == hash( + frozendict({"num_epochs": 10, "batch_size": 64}) # type: ignore + ) + + +def test_nested_in_subparser() -> None: + # https://github.com/brentyi/tyro/issues/9 + @dataclasses.dataclass(frozen=True) + class Subtype: + data: int = 1 + + @dataclasses.dataclass(frozen=True) + class TypeA: + subtype: Subtype = Subtype(1) + + @dataclasses.dataclass(frozen=True) + class TypeB: + subtype: Subtype = Subtype(2) + + @dataclasses.dataclass(frozen=True) + class Wrapper: + supertype: TypeA | TypeB = TypeA() + + assert tyro.cli(Wrapper, args=[]) == Wrapper() + assert ( + tyro.cli(Wrapper, args="supertype:type-a --supertype.subtype.data 1".split(" ")) + == Wrapper() + ) + + +def test_nested_in_subparser_override_with_default() -> None: + @dataclasses.dataclass(frozen=True) + class Mnist: + binary: bool = False + """Set to load binary version of MNIST dataset.""" + + @dataclasses.dataclass + class ImageNet: + subset: Literal[50, 100, 1000] + """Choose between ImageNet-50, ImageNet-100, ImageNet-1000, etc.""" + + # Possible optimizer configurations. + + Selector = tyro.extras.subcommand_type_from_defaults( + { + "m": Mnist(), + "i": ImageNet(50), + } + ) + + @dataclasses.dataclass(frozen=True) + class DatasetContainer: + dataset: Selector = Mnist() # type: ignore + + @dataclasses.dataclass + class Adam: + learning_rate: float = 1e-3 + betas: Tuple[float, float] = (0.9, 0.999) + container: DatasetContainer = DatasetContainer() + + @dataclasses.dataclass + class Sgd: + learning_rate: float = 3e-4 + container: DatasetContainer = DatasetContainer() + + # Train script. + + Optimizers = tyro.extras.subcommand_type_from_defaults( + {"adam": Adam(container=DatasetContainer(ImageNet(50))), "sgd": Sgd()}, + prefix_names=False, + ) + + @tyro.conf.configure(tyro.conf.OmitSubcommandPrefixes) + def train( + optimizer: Optimizers = Adam(container=DatasetContainer(ImageNet(50))), # type: ignore + ) -> Adam | Sgd: + return optimizer + + assert tyro.cli(train, args=[]) == Adam(container=DatasetContainer(ImageNet(50))) + assert tyro.cli(train, args=["adam"]) == Adam( + container=DatasetContainer(ImageNet(50)) + ) + assert tyro.cli(train, args=["sgd"]) == Sgd(container=DatasetContainer(Mnist())) + + +def test_underscore_prefix() -> None: + """https://github.com/brentyi/tyro/issues/77""" + + @dataclasses.dataclass + class PrivateConfig: + pass + + @dataclasses.dataclass + class BaseConfig: + _private: PrivateConfig = dataclasses.field( + default_factory=lambda: PrivateConfig() + ) + + @dataclasses.dataclass + class Level1(BaseConfig): + pass + + @dataclasses.dataclass + class Level2(BaseConfig): + child: Level1 = dataclasses.field(default_factory=lambda: Level1()) + + @dataclasses.dataclass + class Level3(BaseConfig): + child: Level2 = dataclasses.field(default_factory=lambda: Level2()) + + tyro.cli(Level3, args=[]) + + +def test_subcommand_dict_helper() -> None: + def checkout(branch: str) -> str: + """Check out a branch.""" + return branch + + def commit(message: str, all: bool = False) -> Tuple[str, bool]: + """Make a commit.""" + return (message, all) + + # If we only get one, we unfortunately can't form subcommands. This is + # because unions in Python require at least 2 types. + with pytest.raises(AssertionError): + tyro.extras.subcommand_cli_from_dict( + { + "commit": commit, + }, + args="--message hello --all".split(" "), + ) + + assert ( + tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + }, + args="checkout --branch main".split(" "), + ) + == "main" + ) + assert tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + }, + args="commit --message hello".split(" "), + ) == ("hello", False) + assert tyro.extras.subcommand_cli_from_dict( + { + "checkout": checkout, + "commit": commit, + }, + args="commit --message hello --all".split(" "), + ) == ("hello", True) + + +def test_subcommand_by_type_tree() -> None: + @dataclasses.dataclass(frozen=True) + class A: + a: int | str + + @dataclasses.dataclass + class Args: + inner: Annotated[A, tyro.conf.subcommand(name="alt", default=A(5))] | A = A( + "hello" + ) + + assert tyro.cli(Args, args=[]) == Args(A("hello")) + assert "default: inner:alt" in get_helptext_with_checks(Args) diff --git a/tests/test_py311_generated/test_nested_in_containers_generated.py b/tests/test_py311_generated/test_nested_in_containers_generated.py new file mode 100644 index 000000000..819998176 --- /dev/null +++ b/tests/test_py311_generated/test_nested_in_containers_generated.py @@ -0,0 +1,379 @@ +import dataclasses +import enum +from typing import Any, Dict, Generic, List, Set, Tuple, TypeVar + +import pytest + +import tyro + + +@dataclasses.dataclass(frozen=True) +class Color: + r: int + g: int + b: int + + +def test_nested_tuple_fixed_single() -> None: + def main(x: Tuple[Color]) -> Any: + return x + + assert tyro.cli(main, args="--x.0.r 255 --x.0.g 127 --x.0.b 5".split(" ")) == ( + Color(255, 127, 5), + ) + + +def test_nested_tuple_fixed_two() -> None: + def main(x: Tuple[Color, Color]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 255 --x.0.g 127 --x.0.b 5 --x.1.r 255 --x.1.g 127 --x.1.b 0" + ).split(" "), + ) == ( + Color(255, 127, 5), + Color(255, 127, 0), + ) + + +def test_nested_tuple_fixed_three() -> None: + def main(x: Tuple[Color, int, Color]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 255 --x.0.g 127 --x.0.b 5 --x.1 94709 --x.2.r 255 --x.2.g 127" + " --x.2.b 0" + ).split(" "), + ) == ( + Color(255, 127, 5), + 94709, + Color(255, 127, 0), + ) + + +def test_nested_tuple_recursive() -> None: + def main(x: Tuple[Color, Tuple[Color, Color]]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 255 --x.0.g 127 --x.0.b 5 --x.1.0.r 255 --x.1.0.g 127 --x.1.0.b 0" + " --x.1.1.r 255 --x.1.1.g 127 --x.1.1.b 0" + ).split(" "), + ) == ( + Color(255, 127, 5), + ( + Color(255, 127, 0), + Color(255, 127, 0), + ), + ) + + +def test_tuple_bad() -> None: + # Unable to infer input length. + def main(x: Tuple[Color, ...]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_set_bad() -> None: + # Unable to infer input length. + def main(x: Set[Color]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_set_ok() -> None: + def main(x: Set[Color] = {Color(255, 0, 0)}) -> Any: + return x + + assert tyro.cli(main, args=[]) == {Color(255, 0, 0)} + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == {Color(127, 0, 0)} + + +def test_list_bad() -> None: + # Unable to infer input length. + def main(x: List[Color]) -> None: + pass + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_list_ok() -> None: + def main(x: List[Color] = [Color(255, 0, 0)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [Color(255, 0, 0)] + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == [Color(127, 0, 0)] + + +def test_list_object() -> None: + def main(x: List[object] = [Color(255, 0, 0)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [Color(255, 0, 0)] + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == [Color(127, 0, 0)] + + +def test_list_any() -> None: + def main(x: List[Any] = [Color(255, 0, 0)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [Color(255, 0, 0)] + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == [Color(127, 0, 0)] + + +def test_tuple_in_list() -> None: + def main(x: List[Tuple[Color]] = [(Color(255, 0, 0),)]) -> Any: + return x + + assert tyro.cli(main, args=[]) == [(Color(255, 0, 0),)] + assert tyro.cli(main, args="--x.0.0.r 127".split(" ")) == [(Color(127, 0, 0),)] + + +def test_tuple_variable() -> None: + def main(x: Tuple[Color, ...] = (Color(255, 0, 0), Color(255, 0, 127))) -> Any: + return x + + assert tyro.cli(main, args=[]) == (Color(255, 0, 0), Color(255, 0, 127)) + assert tyro.cli(main, args="--x.0.r 127".split(" ")) == ( + Color(127, 0, 0), + Color(255, 0, 127), + ) + + +def test_dict_bad() -> None: + def main(x: Dict[str, Color]) -> Any: + return x + + with pytest.raises(tyro.UnsupportedTypeAnnotationError): + tyro.cli(main, args=[]) + + +def test_dict_ok() -> None: + def main( + x: Dict[str, Color] = { + "red": Color(255, 0, 0), + "green": Color(0, 255, 0), + "blue": Color(0, 0, 255), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])["green"] == Color(0, 255, 0) + assert tyro.cli(main, args="--x.green.g 127".split(" "))["green"] == Color( + 0, 127, 0 + ) + + +def test_dict_key_int() -> None: + def main( + x: Dict[int, Color] = { + 0: Color(255, 0, 0), + 1: Color(0, 255, 0), + 2: Color(0, 0, 255), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])[1] == Color(0, 255, 0) + assert tyro.cli(main, args="--x.1.g 127".split(" "))[1] == Color(0, 127, 0) + + +def test_dict_key_enum() -> None: + class ColorType(enum.Enum): + RED = enum.auto() + GREEN = enum.auto() + BLUE = enum.auto() + + def main( + x: Dict[ColorType, Color] = { + ColorType.RED: Color(255, 0, 0), + ColorType.GREEN: Color(0, 255, 0), + ColorType.BLUE: Color(0, 0, 255), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])[ColorType.GREEN] == Color(0, 255, 0) + assert tyro.cli(main, args="--x.GREEN.g 127".split(" "))[ColorType.GREEN] == Color( + 0, 127, 0 + ) + + +def test_dict_nested() -> None: + def main( + x: Dict[str, Tuple[Color, int]] = { + # For each color: RGB and xterm color code. + "red": (Color(255, 0, 0), 9), + "green": (Color(0, 255, 0), 10), + "blue": (Color(0, 0, 255), 12), + }, + ) -> Any: + return x + + assert tyro.cli(main, args=[])["green"] == (Color(0, 255, 0), 10) + assert tyro.cli(main, args="--x.green.0.g 127 --x.green.1 2".split(" "))[ + "green" + ] == (Color(0, 127, 0), 2) + + +def test_generic_in_tuple() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main(x: Tuple[GenericColor[float], GenericColor[int]]) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 0.5 --x.0.g 0.2 --x.0.b 0.3 --x.1.r 25 --x.1.g 2 --x.1.b 3".split( + " " + ) + ), + ) == (GenericColor(0.5, 0.2, 0.3), GenericColor(25, 2, 3)) + + +def test_generic_in_tuple_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Tuple[GenericColor[float], GenericColor[int]] = ( + GenericColor(0.5, 0.2, 0.3), + GenericColor[int](25, 2, 3), # The subscript should be optional. + ), + ) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 0.5 --x.0.g 0.2 --x.0.b 0.3 --x.1.r 25 --x.1.g 2 --x.1.b 3".split( + " " + ) + ), + ) == (GenericColor(0.5, 0.2, 0.3), GenericColor(25, 2, 3)) + + +def test_generic_in_variable_tuple_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Tuple[GenericColor, ...] = ( + GenericColor[float](0.5, 0.2, 0.3), + GenericColor[int](25, 2, 3), + ), + ) -> Any: + return x + + assert tyro.cli( + main, + args=( + "--x.0.r 0.5 --x.0.g 0.9 --x.0.b 0.3 --x.1.r 25 --x.1.g 2 --x.1.b 3".split( + " " + ) + ), + ) == (GenericColor(0.5, 0.9, 0.3), GenericColor(25, 2, 3)) + + +def test_generic_in_dict_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Dict[str, GenericColor] = { + "float": GenericColor(0.5, 0.2, 0.3), + "int": GenericColor[int](25, 2, 3), + }, + ) -> Any: + return x + + assert tyro.cli( + main, + args="--x.float.g 0.1".split(" "), + )["float"] == GenericColor(0.5, 0.1, 0.3) + assert tyro.cli( + main, + args="--x.int.g 0".split(" "), + ) == {"float": GenericColor(0.5, 0.2, 0.3), "int": GenericColor(25, 0, 3)} + + +def test_generic_in_double_nested_dict_with_default() -> None: + ScalarType = TypeVar("ScalarType", int, float) + + @dataclasses.dataclass + class GenericColor(Generic[ScalarType]): + r: ScalarType + g: ScalarType + b: ScalarType + + def main( + x: Dict[str, Dict[str, GenericColor]] = { + "hello": { + "float": GenericColor(0.5, 0.2, 0.3), + "int": GenericColor[int](25, 2, 3), + } + }, + ) -> Any: + return x + + assert tyro.cli( + main, + args="--x.hello.float.g 0.1".split(" "), + )["hello"]["float"] == GenericColor(0.5, 0.1, 0.3) + assert tyro.cli( + main, + args="--x.hello.int.g 0".split(" "), + ) == { + "hello": {"float": GenericColor(0.5, 0.2, 0.3), "int": GenericColor(25, 0, 3)} + } + + +def test_double_nested_dict_with_inferred_type() -> None: + def main( + x: Dict[str, Any] = { + "hello": { + "a": Color(5, 2, 3), + "b": Color(25, 2, 3), + } + }, + ) -> Any: + return x + + assert tyro.cli( + main, + args="--x.hello.a.g 1".split(" "), + )["hello"]["a"] == Color(5, 1, 3) diff --git a/tests/test_py311_generated/test_new_style_annotations_generics_generated.py b/tests/test_py311_generated/test_new_style_annotations_generics_generated.py new file mode 100644 index 000000000..a1d5203c5 --- /dev/null +++ b/tests/test_py311_generated/test_new_style_annotations_generics_generated.py @@ -0,0 +1,54 @@ +from __future__ import annotations # Should enable support for all versions of Python. + +from typing import Any, Literal, Optional + +import pytest + +import tyro + + +def test_list() -> None: + def main(x: list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_tuple() -> None: + def main(x: tuple[bool, str]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == (True, "False") + + +def test_tuple_variable() -> None: + def main(x: tuple[bool | str, ...]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "Wrong"]) == (True, "Wrong") + + +def test_super_nested() -> None: + def main( + x: Optional[ + list[ + tuple[ + Optional[int], + Literal[3, 4], + tuple[int, int] | tuple[str, str], + ] + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) diff --git a/tests/test_py311_generated/test_new_style_annotations_min_py310_generated.py b/tests/test_py311_generated/test_new_style_annotations_min_py310_generated.py new file mode 100644 index 000000000..dc17e3db4 --- /dev/null +++ b/tests/test_py311_generated/test_new_style_annotations_min_py310_generated.py @@ -0,0 +1,64 @@ +from typing import Any, Literal + +import pytest + +import tyro + + +def test_union_direct(): + assert tyro.cli(int | str, args=["5"]) == 5 + assert tyro.cli(int | str, args=["five"]) == "five" + + +def test_union_basic(): + def main(x: int | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_union_with_list(): + def main(x: int | str | list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + assert tyro.cli(main, args=["--x", "True"]) == "True" + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_union_literal(): + def main(x: Literal[1, 2] | Literal[3, 4, 5] | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == "6" + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_super_nested(): + def main( + x: None + | list[ + tuple[ + None | int, + Literal[3, 4], + tuple[int, int] | tuple[str, str], + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) diff --git a/tests/test_py311_generated/test_new_style_annotations_min_py312_generated.py b/tests/test_py311_generated/test_new_style_annotations_min_py312_generated.py new file mode 100644 index 000000000..d04c73f3b --- /dev/null +++ b/tests/test_py311_generated/test_new_style_annotations_min_py312_generated.py @@ -0,0 +1,54 @@ +# mypy: ignore-errors +# +# PEP 695 isn't yet supported in mypy. (April 4, 2024) +from dataclasses import dataclass + +import tyro + + +def test_simple_generic(): + @dataclass(frozen=True) + class Container[T]: + a: T + b: T + + assert tyro.cli(Container[int], args="--a 1 --b 2".split(" ")) == Container(1, 2) + + +type X = int +type Y = list[int] +type Z = Inner[int] + + +@dataclass(frozen=True) +class Inner[T]: + a: T + b: T + + +def test_generic_with_type_statement_0(): + @dataclass(frozen=True) + class Container[T]: + a: T + b: T + + assert tyro.cli(Container[X], args="--a 1 --b 2".split(" ")) == Container(1, 2) + + +def test_generic_with_type_statement_1(): + @dataclass(frozen=True) + class Container[T]: + a: tuple[X, ...] + b: T + + assert tyro.cli(Container[Y], args="--a 1 --b 2".split(" ")) == Container((1,), [2]) + + +def test_generic_with_type_statement_2(): + @dataclass(frozen=True) + class Container[T]: + a: Z + + assert tyro.cli(Container[Y], args="--a.a 1 --a.b 2".split(" ")) == Container( + Inner(1, 2) + ) diff --git a/tests/test_py311_generated/test_new_style_annotations_min_py39_generated.py b/tests/test_py311_generated/test_new_style_annotations_min_py39_generated.py new file mode 100644 index 000000000..0fe423b59 --- /dev/null +++ b/tests/test_py311_generated/test_new_style_annotations_min_py39_generated.py @@ -0,0 +1,69 @@ +import dataclasses +from typing import Any, Literal, Optional + +import pytest + +import tyro + + +def test_list() -> None: + def main(x: list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_tuple() -> None: + def main(x: tuple[bool, str]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "False"]) == (True, "False") + + +def test_tuple_nested() -> None: + @dataclasses.dataclass + class Args: + a: int + + def main(x: tuple[Args, Args]) -> Any: + return x + + assert tyro.cli(main, args=["--x.0.a", "3", "--x.1.a", "4"]) == (Args(3), Args(4)) + + +def test_tuple_variable() -> None: + def main(x: tuple[bool | str, ...]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "True", "Wrong"]) == (True, "Wrong") + + +def test_super_nested() -> None: + def main( + x: Optional[ + list[ + tuple[ + Optional[int], + Literal[3, 4], + tuple[int, int] | tuple[str, str], + ] + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) + + +def test_tuple_direct() -> None: + assert tyro.cli(tuple[int, ...], args="1 2".split(" ")) == (1, 2) # type: ignore + assert tyro.cli(tuple[int, int], args="1 2".split(" ")) == (1, 2) # type: ignore diff --git a/tests/test_py311_generated/test_new_style_annotations_unions_generated.py b/tests/test_py311_generated/test_new_style_annotations_unions_generated.py new file mode 100644 index 000000000..ccb1920ea --- /dev/null +++ b/tests/test_py311_generated/test_new_style_annotations_unions_generated.py @@ -0,0 +1,61 @@ +from __future__ import annotations # Should enable support for all versions of Python. + +from typing import Any, Literal + +import pytest + +import tyro + + +def test_union_basic(): + def main(x: int | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_union_with_list(): + def main(x: int | str | list[bool]) -> Any: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == 6 + assert tyro.cli(main, args=["--x", "five"]) == "five" + assert tyro.cli(main, args=["--x", "True"]) == "True" + assert tyro.cli(main, args=["--x", "True", "False"]) == [True, False] + + +def test_union_literal(): + def main(x: Literal[1, 2] | Literal[3, 4, 5] | str) -> int | str: + return x + + assert tyro.cli(main, args=["--x", "5"]) == 5 + assert tyro.cli(main, args=["--x", "6"]) == "6" + assert tyro.cli(main, args=["--x", "five"]) == "five" + + +def test_super_nested(): + def main( + x: None + | list[ + tuple[ + None | int, + Literal[3, 4], + tuple[int, int] | tuple[str, str], + ] + ] = None, + ) -> Any: + return x + + assert tyro.cli(main, args=[]) is None + assert tyro.cli(main, args="--x None".split(" ")) is None + assert tyro.cli(main, args="--x None 3 2 2".split(" ")) == [(None, 3, (2, 2))] + assert tyro.cli(main, args="--x 2 3 x 2".split(" ")) == [(2, 3, ("x", "2"))] + assert tyro.cli(main, args="--x 2 3 x 2 2 3 1 2".split(" ")) == [ + (2, 3, ("x", "2")), + (2, 3, (1, 2)), + ] + with pytest.raises(SystemExit): + tyro.cli(main, args=["--help"]) diff --git a/tests/test_py311_generated/test_optional_annotated_generated.py b/tests/test_py311_generated/test_optional_annotated_generated.py new file mode 100644 index 000000000..751e612e6 --- /dev/null +++ b/tests/test_py311_generated/test_optional_annotated_generated.py @@ -0,0 +1,55 @@ +"""Test adapted from @KolinGuo: https://github.com/brentyi/tyro/issues/156""" + +import dataclasses +from typing import Annotated, Optional, Tuple + +import tyro + + +def detect_usb_device( + vendor_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--vid"])] = None, + product_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--pid"])] = "", +) -> Tuple[Optional[str], Optional[str]]: + """ + Detect connected USB device by vendor_id + product_id. + + :param vendor_id: vendor_id of a specific USB device. + :param product_id: product_id of a specific USB device. + """ + return vendor_id, product_id + + +def test_detect_usb_device() -> None: + assert tyro.cli(detect_usb_device, args=["--vid", "0x1234", "--pid", "0x5678"]) == ( + "0x1234", + "0x5678", + ) + assert tyro.cli( + detect_usb_device, args=["--vendor-id", "0x1234", "--product-id", "0x5678"] + ) == ( + "0x1234", + "0x5678", + ) + assert tyro.cli(detect_usb_device, args=[]) == (None, "") + + +@dataclasses.dataclass +class DeviceStruct: + vendor_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--vid"])] = None + product_id: Annotated[Optional[str], tyro.conf.arg(aliases=["--pid"])] = "" + + +def test_detect_usb_device_dataclass() -> None: + assert tyro.cli( + DeviceStruct, args=["--vid", "0x1234", "--pid", "0x5678"] + ) == DeviceStruct( + "0x1234", + "0x5678", + ) + assert tyro.cli( + DeviceStruct, args=["--vendor-id", "0x1234", "--product-id", "0x5678"] + ) == DeviceStruct( + "0x1234", + "0x5678", + ) + assert tyro.cli(DeviceStruct, args=[]) == DeviceStruct(None, "") diff --git a/tests/test_py311_generated/test_partial_generated.py b/tests/test_py311_generated/test_partial_generated.py new file mode 100644 index 000000000..89fe78552 --- /dev/null +++ b/tests/test_py311_generated/test_partial_generated.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import dataclasses +import functools + +from helptext_utils import get_helptext_with_checks + +import tyro + + +def test_partial_func() -> None: + def main(a: int, b: str) -> str: + return b * a + + assert tyro.cli(functools.partial(main, a=3), args=["--b", "hi"]) == "hihihi" + + +def test_partial_class() -> None: + class Main: + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + assert tyro.cli(functools.partial(Main, a=3), args=["--b", "hi"]).inner == "hihihi" + + +def test_partial_helptext_func() -> None: + def main(a: int, b: str) -> str: + """Hello!""" + return b * a + + helptext = get_helptext_with_checks(functools.partial(main, b="3")) + assert "partial" not in helptext + assert "Hello!" in helptext + + +def test_partial_helptext_class() -> None: + class Main: + """Hello!""" + + def __init__(self, a: int, b: str) -> None: + self.inner = b * a + + helptext = get_helptext_with_checks(functools.partial(Main, b="3")) + assert "partial" not in helptext + assert "Hello!" in helptext + + +def test_partial_helptext_dataclass() -> None: + @dataclasses.dataclass + class Config: + a: int + b: int + """Hello!""" + c: int + + ConfigWithDefaults = functools.partial(functools.partial(Config, a=3), c=5) + assert tyro.cli(ConfigWithDefaults, args=["--b", "4"]) == Config(3, 4, 5) + + helptext = get_helptext_with_checks(ConfigWithDefaults) + assert "partial" not in helptext + assert "Hello!" in helptext diff --git a/tests/test_py311_generated/test_positional_min_py38_generated.py b/tests/test_py311_generated/test_positional_min_py38_generated.py new file mode 100644 index 000000000..15f8f7911 --- /dev/null +++ b/tests/test_py311_generated/test_positional_min_py38_generated.py @@ -0,0 +1,118 @@ +from typing import List, Optional, Tuple + +import pytest + +import tyro + + +def test_positional(): + def main( + x: int, + y: int, + /, + # Note: it's generally a bad idea to have a mutable object (like a list) as a + # default value. But it should still work. + z: List[int] = [1, 2, 3], + ) -> Tuple[int, int, int]: + """main. + + Args: + x: x + y: y + z: z + + Returns: + Tuple[int, int, int]: Output. + """ + return (x, y, z[0]) + + assert tyro.cli(main, args="1 2 --z 3".split(" ")) == (1, 2, 3) + with pytest.raises(SystemExit): + assert tyro.cli(main, args="--x 1 --y 2 --z 3".split(" ")) == (1, 2, 3) + + +def test_nested_positional(): + class A: + def __init__(self, a: int, hello_world: int, /, c: int): + self.hello_world = hello_world + + def nest1(a: int, b: int, thing: A, /, c: int) -> A: + return thing + + assert isinstance(tyro.cli(nest1, args="0 1 2 3 4 --c 4".split(" ")), A) + assert tyro.cli(nest1, args="0 1 2 3 4 --c 4".split(" ")).hello_world == 3 + with pytest.raises(SystemExit): + tyro.cli(nest1, args="0 1 2 3 4 4 --c 4".split(" ")) + + +def test_nested_positional_alt(): + class B: + def __init__(self, a: int, b: int, /, c: int): + pass + + def nest2(a: int, b: int, /, thing: B, c: int): + return thing + + assert isinstance(tyro.cli(nest2, args="0 1 2 3 --thing.c 4 --c 4".split(" ")), B) + with pytest.raises(SystemExit): + tyro.cli(nest2, args="0 1 2 3 4 --thing.c 4 --c 4".split(" ")) + + +def test_positional_with_underscores(): + """Hyphen replacement works a bit different for positional arguments.""" + + def main(a_multi_word_input: int, /) -> int: + return a_multi_word_input + + assert tyro.cli(main, args=["5"]) == 5 + + +def test_positional_booleans(): + """Make sure that flag behavior is disabled for positional booleans.""" + + def main( + flag1: bool, + flag2: bool = True, + flag3: bool = False, + /, + ) -> Tuple[bool, bool, bool]: + return flag1, flag2, flag3 + + assert tyro.cli(main, args=["True"]) == (True, True, False) + assert tyro.cli(main, args=["True", "False"]) == (True, False, False) + assert tyro.cli(main, args=["False", "False", "True"]) == (False, False, True) + + with pytest.raises(SystemExit): + tyro.cli(main, args=["hmm"]) + with pytest.raises(SystemExit): + tyro.cli(main, args=["true"]) + with pytest.raises(SystemExit): + tyro.cli(main, args=["True", "false"]) + + +def test_optional_list(): + def main(a: Optional[List[int]], /) -> Optional[List[int]]: + return a + + assert tyro.cli(main, args=["None"]) is None + assert tyro.cli(main, args=["1", "2"]) == [1, 2] + assert tyro.cli(main, args=[]) == [] + with pytest.raises(SystemExit): + tyro.cli(main, args=["hm"]) + + +def test_optional_list_with_default(): + def main(a: Optional[List[int]] = None, /) -> Optional[List[int]]: + return a + + assert tyro.cli(main, args=["None"]) is None + assert tyro.cli(main, args=["5", "5"]) == [5, 5] + with pytest.raises(SystemExit): + tyro.cli(main, args=["None", "5"]) + + +def test_positional_tuple(): + def main(x: Tuple[int, int], y: Tuple[str, str], /): + return x, y + + assert tyro.cli(main, args="1 2 3 4".split(" ")) == ((1, 2), ("3", "4")) diff --git a/tests/test_py311_generated/test_pydantic_generated.py b/tests/test_py311_generated/test_pydantic_generated.py new file mode 100644 index 000000000..c54ef87df --- /dev/null +++ b/tests/test_py311_generated/test_pydantic_generated.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +import contextlib +import io +import pathlib +from typing import Annotated, cast + +import pytest +from pydantic import BaseModel, Field, v1 + +import tyro +import tyro._strings + + +def test_pydantic() -> None: + class ManyTypesA(BaseModel): + i: int + s: str = "hello" + f: float = Field(default_factory=lambda: 3.0) + p: pathlib.Path + ignored: float = Field(default_factory=lambda: 3.0, init_var=False) + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesA, + args=[ + "--i", + "5", + "--p", + "~", + ], + ) == ManyTypesA(i=5, s="hello", f=3.0, p=pathlib.Path("~")) + + +def test_pydantic_v1() -> None: + class ManyTypesA(v1.BaseModel): + i: int + s: str = "hello" + f: float = v1.Field(default_factory=lambda: 3.0) + p: pathlib.Path + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesA, + args=[ + "--i", + "5", + "--p", + "~", + ], + ) == ManyTypesA(i=5, s="hello", f=3.0, p=pathlib.Path("~")) + + +def test_pydantic_helptext() -> None: + class Helptext(BaseModel): + """This docstring should be printed as a description.""" + + x: int = Field(description="Documentation 1") + + y: int = Field(description="Documentation 2") + + z: int = Field(description="Documentation 3") + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + assert tyro._strings.strip_ansi_sequences(cast(str, Helptext.__doc__)) in helptext + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_pydantic_suppress_base_model_helptext() -> None: + class Helptext(BaseModel): + x: int = Field(description="Documentation 1") + + y: int = Field(description="Documentation 2") + + z: int = Field(description="Documentation 3") + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + + assert "Create a new model by parsing and validating" not in helptext + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +class HelptextWithFieldDocstring(BaseModel): + """This docstring should be printed as a description.""" + + x: int + """Documentation 1""" + + y: int = Field(description="Documentation 2") + + z: int = Field(description="Documentation 3") + + +def test_pydantic_field_helptext_from_docstring() -> None: + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(HelptextWithFieldDocstring, args=["--help"]) + helptext = f.getvalue() + assert ( + tyro._strings.strip_ansi_sequences( + cast(str, HelptextWithFieldDocstring.__doc__) + ) + in helptext + ) + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_pydantic_positional_annotation() -> None: + class AnnotatedAsPositional(BaseModel): + name: tyro.conf.Positional[str] + """This is annotated as a positional argument.""" + + result = tyro.cli(AnnotatedAsPositional, args=["myname"]) + assert isinstance(result, AnnotatedAsPositional) + + +def test_pydantic_alias() -> None: + class AliasCfg(BaseModel): + alias: Annotated[str, tyro.conf.arg(aliases=["-a"])] + + assert tyro.cli(AliasCfg, args=["--alias", "3"]) == AliasCfg(alias="3") + assert tyro.cli(AliasCfg, args=["-a", "3"]) == AliasCfg(alias="3") + + +def test_pydantic_default_instance() -> None: + class Inside(BaseModel): + x: int = 1 + + class Outside(BaseModel): + i: "Inside" = Inside(x=2) + + assert tyro.cli(Outside, args=[]).i.x == 2, ( + "Expected x value from the default instance", + ) + assert tyro.cli(Outside, args=["--i.x", "3"]).i.x == 3 + + +def test_pydantic_nested_default_instance() -> None: + class Inside(BaseModel): + x: int = 1 + + class Middle(BaseModel): + i: Inside + + class Outside(BaseModel): + m: Middle = Middle(i=Inside(x=2)) + + assert tyro.cli(Outside, args=[]).m.i.x == 2, ( + "Expected x value from the default instance", + ) + assert tyro.cli(Outside, args=["--m.i.x", "3"]).m.i.x == 3 + + +# Updating forward references in Pydantic v1 requires that these classes are global. + + +class InsideV1(v1.BaseModel): + x: int = 1 + + +class MiddleV1(v1.BaseModel): + i: InsideV1 + + +def test_pydantic_v1_nested_default_instance() -> None: + class Outside(v1.BaseModel): + m: MiddleV1 = MiddleV1(i=InsideV1(x=2)) + + assert tyro.cli(Outside, args=[]).m.i.x == 2, ( + "Expected x value from the default instance", + ) + assert tyro.cli(Outside, args=["--m.i.x", "3"]).m.i.x == 3 diff --git a/tests/test_py311_generated/test_pydantic_with_newtype_generated.py b/tests/test_py311_generated/test_pydantic_with_newtype_generated.py new file mode 100644 index 000000000..44ecb7300 --- /dev/null +++ b/tests/test_py311_generated/test_pydantic_with_newtype_generated.py @@ -0,0 +1,36 @@ +from typing import Annotated, NewType, Tuple + +import pydantic +from pydantic import BaseModel + +import tyro + +Microliter = NewType("Microliter", int) + + +class Measurements(BaseModel): + single: Microliter = pydantic.Field(10) + renamed_single: Annotated[Microliter, tyro.conf.arg(name="other_single")] = ( + pydantic.Field(10) + ) + pair: Tuple[Microliter, Microliter] = pydantic.Field((20, 30)) + + +IncorrectMeasurements = NewType("IncorrectMeasurements", Measurements) + + +def test_pydantic_with_newtype(): + assert tyro.cli( + IncorrectMeasurements, args="--single 1 --pair 2 3".split(" ") + ) == Measurements( + single=Microliter(1), + renamed_single=Microliter(10), + pair=(Microliter(2), Microliter(3)), + ) + assert tyro.cli( + IncorrectMeasurements, args="--single 1 --other-single 5".split(" ") + ) == Measurements( + single=Microliter(1), + renamed_single=Microliter(5), + pair=(Microliter(20), Microliter(30)), + ) diff --git a/tests/test_py311_generated/test_self_type_generated.py b/tests/test_py311_generated/test_self_type_generated.py new file mode 100644 index 000000000..021ef5b00 --- /dev/null +++ b/tests/test_py311_generated/test_self_type_generated.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from typing import Self + +import pytest + +import tyro + + +class SomeClass: + def __init__(self, a: int, b: int) -> None: + self.a = a + self.b = b + + def method1(self, x: Self) -> None: + self.effect = x + + @classmethod + def method2(cls, x: Self) -> SomeClass: + return x + + # Self is not valid in static methods. + # https://peps.python.org/pep-0673/#valid-locations-for-self + # + # @staticmethod + # def method3(x: Self) -> TestClass: + # return x + + +class SomeSubclass(SomeClass): ... + + +def test_method() -> None: + x = SomeClass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method1, args=[]) + + assert tyro.cli(x.method1, args="--x.a 3 --x.b 3".split(" ")) is None + assert x.effect.a == 3 and x.effect.b == 3 + assert isinstance(x, SomeClass) + + +def test_classmethod() -> None: + x = SomeClass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method2, args=[]) + with pytest.raises(SystemExit): + tyro.cli(SomeClass.method2, args=[]) + + y = tyro.cli(x.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + y = tyro.cli(SomeClass.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + +def test_subclass_method() -> None: + x = SomeSubclass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method1, args=[]) + + assert tyro.cli(x.method1, args="--x.a 3 --x.b 3".split(" ")) is None + assert x.effect.a == 3 and x.effect.b == 3 + assert isinstance(x, SomeSubclass) + + y = tyro.cli(x.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + +def test_subclass_classmethod() -> None: + x = SomeSubclass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method2, args=[]) + with pytest.raises(SystemExit): + tyro.cli(SomeSubclass.method2, args=[]) + + y = tyro.cli(x.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + y = tyro.cli(SomeSubclass.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) diff --git a/tests/test_py311_generated/test_strings_generated.py b/tests/test_py311_generated/test_strings_generated.py new file mode 100644 index 000000000..ff7669a75 --- /dev/null +++ b/tests/test_py311_generated/test_strings_generated.py @@ -0,0 +1,70 @@ +from tyro import _strings + + +def test_words_from_name(): + assert _strings.hyphen_separated_from_camel_case("MyHTTPServer") == "my-http-server" + assert ( + _strings.hyphen_separated_from_camel_case("my-http-server") == "my-http-server" + ) + + +def test_make_field_name(): + assert _strings.make_field_name(["hello", "world"]) == "hello.world" + assert _strings.make_field_name(["hello_world", "world"]) == "hello-world.world" + assert ( + _strings.make_field_name(["hello_world", "___hello_world"]) + == "hello-world.___hello-world" + ) + assert ( + _strings.make_field_name(["hello_world", "---hello_world"]) + == "hello-world.---hello-world" + ) + + +def test_postprocess_helptext(): + assert _strings.remove_single_line_breaks("hello world") == "hello world" + assert _strings.remove_single_line_breaks("hello\nworld") == "hello world" + assert _strings.remove_single_line_breaks("hello \nworld") == "hello world" + assert _strings.remove_single_line_breaks("hello\n\nworld") == "hello\n\nworld" + assert ( + _strings.remove_single_line_breaks( + "a paragraph:\nSentence one.\nSentence two.\nSentence three.\n" + ) + == "a paragraph: Sentence one. Sentence two. Sentence three." + ) + assert ( + _strings.remove_single_line_breaks( + "a bulleted list:\n" + "- The first problem.\n" + "- The second problem.\n" + "- The third problem.\n" + ) + == "a bulleted list:\n" + "- The first problem.\n" + "- The second problem.\n" + "- The third problem." + ) + assert ( + _strings.remove_single_line_breaks( + "an indented list:\n" + " The first problem.\n" + " The second problem.\n" + " The third problem.\n" + ) + == "an indented list:\n" + " The first problem.\n" + " The second problem.\n" + " The third problem." + ) + assert ( + _strings.remove_single_line_breaks( + "a numbered list:\n" + "1. The first problem.\n" + "2. The second problem.\n" + "3. The third problem.\n" + ) + == "a numbered list:\n" + "1. The first problem.\n" + "2. The second problem.\n" + "3. The third problem." + ) diff --git a/tests/test_py311_generated/test_tuple_with_subcommands_generated.py b/tests/test_py311_generated/test_tuple_with_subcommands_generated.py new file mode 100644 index 000000000..9419bce56 --- /dev/null +++ b/tests/test_py311_generated/test_tuple_with_subcommands_generated.py @@ -0,0 +1,109 @@ +# mypy: disable-error-code="call-overload,misc" +# +# Mypy errors from passing union types directly into tyro.cli() as Type[T]. We would +# benefit from TypeForm[T]: https://github.com/python/mypy/issues/9773 +"""Tests adapted from https://github.com/brentyi/tyro/issues/89, which catches edge +cases when combining nested tuple types, renamed arguments, and subcommands. + +Largely written by @wookayin. +""" + +import dataclasses +from pathlib import Path +from typing import Annotated, Generic, Tuple, TypeVar + +import tyro.conf + +T = TypeVar("T") + + +@dataclasses.dataclass +class Checkout(Generic[T]): + """Check out a branch.""" + + branch: T + + +@dataclasses.dataclass +class Commit: + """Commit something.""" + + input: tyro.conf.Positional[Path] + + +@dataclasses.dataclass +class Arg: + verbose: bool = True + + +def test_case1() -> None: + o = tyro.cli( + Checkout[str] | Commit, + args=["commit", "./path.txt"], + ) + assert o == Commit(input=Path("path.txt")) + + +def test_case2() -> None: + arg, action = tyro.cli( + Tuple[ + Arg, + Annotated[ + Checkout[str] | Commit, + tyro.conf.arg(name=""), + ], + ], + args=["commit", "./path.txt"], + ) + + assert isinstance(arg, Arg) + assert isinstance(action, Commit) + assert action.input == Path("./path.txt") + + +def test_case3() -> None: + o = tyro.cli( + Tuple[ + Annotated[ + Arg, + tyro.conf.arg(name=""), + ], + Annotated[ + Annotated[Checkout[str], tyro.conf.subcommand(name="checkout")] + | Annotated[Commit, tyro.conf.subcommand(name="commit")], + tyro.conf.arg(name=""), + ], + ], + args=["commit", "./path.txt"], + ) + assert o == (Arg(), Commit(Path("./path.txt"))) + + +def test_case4() -> None: + o = tyro.cli( + Tuple[ + Annotated[ + Annotated[Checkout[str], tyro.conf.subcommand(name="checkout")] + | Annotated[Commit, tyro.conf.subcommand(name="commit")], + tyro.conf.arg(name=""), + ] + ], + args=["commit", "./path.txt"], + ) + assert o == (Commit(Path("./path.txt")),) + + +def test_case5() -> None: + assert tyro.cli( + tyro.conf.OmitArgPrefixes[ + Tuple[ + Annotated[ + Checkout[str], + tyro.conf.subcommand(prefix_name=False), + ] + | Annotated[Commit, tyro.conf.subcommand(prefix_name=False)], + Arg, + ] + ], + args=["--no-verbose", "checkout-str", "--branch", "branch"], + ) == (Checkout("branch"), Arg(False)) diff --git a/tests/test_py311_generated/test_typeddict_readonly_min_py38_generated.py b/tests/test_py311_generated/test_typeddict_readonly_min_py38_generated.py new file mode 100644 index 000000000..b91cb7cc9 --- /dev/null +++ b/tests/test_py311_generated/test_typeddict_readonly_min_py38_generated.py @@ -0,0 +1,23 @@ +from typing import TypedDict + +import pytest +from typing_extensions import ReadOnly + +import tyro + + +def test_read_only_typeddict() -> None: + class ManyTypesTypedDict(TypedDict): + i: ReadOnly[int] # type: ignore + s: ReadOnly[ReadOnly[str]] # type: ignore + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) diff --git a/tests/test_py311_generated/test_union_from_mapping_generated.py b/tests/test_py311_generated/test_union_from_mapping_generated.py new file mode 100644 index 000000000..2eb8ff42b --- /dev/null +++ b/tests/test_py311_generated/test_union_from_mapping_generated.py @@ -0,0 +1,49 @@ +import dataclasses +from typing import TYPE_CHECKING, Optional + +import tyro + + +@dataclasses.dataclass +class A: + x: int + + +def test_union_from_mapping(): + base_configs = { + "one": A(1), + "two": A(2), + "three": A(3), + } + ConfigUnion = tyro.extras.subcommand_type_from_defaults(base_configs) + + assert tyro.cli(ConfigUnion, args="one".split(" ")) == A(1) + assert tyro.cli(ConfigUnion, args="two".split(" ")) == A(2) + assert tyro.cli(ConfigUnion, args="two --x 4".split(" ")) == A(4) + assert tyro.cli(ConfigUnion, args="three".split(" ")) == A(3) + + +def test_union_from_mapping_in_function(): + base_configs = { + "one": A(1), + "two": A(2), + "three": A(3), + } + + if TYPE_CHECKING: + ConfigUnion = A + else: + ConfigUnion = tyro.extras.subcommand_type_from_defaults(base_configs) + + def main(config: ConfigUnion, flag: bool = False) -> Optional[A]: # type: ignore + if flag: + return config + return None + + assert tyro.cli(main, args="--flag config:one".split(" ")) == A(1) + assert tyro.cli(main, args="--flag config:one --config.x 3".split(" ")) == A(3) + assert tyro.cli(main, args="config:one --config.x 1".split(" ")) is None + + assert tyro.cli(main, args="--flag config:two".split(" ")) == A(2) + assert tyro.cli(main, args="--flag config:two --config.x 3".split(" ")) == A(3) + assert tyro.cli(main, args="config:two --config.x 1".split(" ")) is None diff --git a/tests/test_py311_generated/test_unsafe_cache_generated.py b/tests/test_py311_generated/test_unsafe_cache_generated.py new file mode 100644 index 000000000..35cb41741 --- /dev/null +++ b/tests/test_py311_generated/test_unsafe_cache_generated.py @@ -0,0 +1,31 @@ +from tyro import _unsafe_cache + + +def test_unsafe_cache(): + x = 0 + + @_unsafe_cache.unsafe_cache(maxsize=2) + def f(dummy: int): + nonlocal x + x += 1 + + f(0) + f(0) + f(0) + assert x == 1 + f(1) + f(1) + f(1) + assert x == 2 + f(0) + f(0) + f(0) + assert x == 2 + f(2) + f(2) + f(2) + assert x == 3 + f(0) + f(0) + f(0) + assert x == 4 diff --git a/tests/test_py311_generated/test_unsupported_but_should_work_generated.py b/tests/test_py311_generated/test_unsupported_but_should_work_generated.py new file mode 100644 index 000000000..74327ac80 --- /dev/null +++ b/tests/test_py311_generated/test_unsupported_but_should_work_generated.py @@ -0,0 +1,59 @@ +"""Tests for features that are not officially features, but should work. + +Includes things like omegaconf.MISSING, attrs, etc, which mostly work but either likely +have corner cases or just seem sketchy. +""" + +from typing import Tuple + +import omegaconf +import pytest + +import tyro +import tyro._strings + + +def test_omegaconf_missing(): + """Passing in a omegaconf.MISSING default; this will mark an argument as required.""" + + def main( + required_a: int, + optional: int = 3, + required_b: int = None, # type: ignore + ) -> Tuple[int, int, int]: + return (required_a, optional, required_b) # type: ignore + + assert tyro.cli( + main, args="--required-a 3 --optional 4 --required-b 5".split(" ") + ) == (3, 4, 5) + assert tyro.cli(main, args="--required-a 3 --required-b 5".split(" ")) == ( + 3, + 3, + 5, + ) + + with pytest.raises(SystemExit): + tyro.cli(main, args="--required-a 3 --optional 4") + with pytest.raises(SystemExit): + tyro.cli(main, args="--required-a 3") + + def main2( + required_a: int, + optional: int = 3, + required_b: int = omegaconf.MISSING, + ) -> Tuple[int, int, int]: + return (required_a, optional, required_b) + + assert tyro.cli( + main2, args="--required-a 3 --optional 4 --required-b 5".split(" ") + ) == (3, 4, 5) + assert tyro.cli(main2, args="--required-a 3 --required-b 5".split(" ")) == ( + 3, + 3, + 5, + ) + + with pytest.raises(SystemExit): + tyro.cli(main2, args="--required-a 3 --optional 4") + with pytest.raises(SystemExit): + tyro.cli(main2, args="--required-a 3") diff --git a/tests/test_pydantic.py b/tests/test_pydantic.py new file mode 100644 index 000000000..5a588e71c --- /dev/null +++ b/tests/test_pydantic.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import contextlib +import io +import pathlib +from typing import cast + +import pytest +from pydantic import BaseModel, Field, v1 +from typing_extensions import Annotated + +import tyro +import tyro._strings + + +def test_pydantic() -> None: + class ManyTypesA(BaseModel): + i: int + s: str = "hello" + f: float = Field(default_factory=lambda: 3.0) + p: pathlib.Path + ignored: float = Field(default_factory=lambda: 3.0, init_var=False) + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesA, + args=[ + "--i", + "5", + "--p", + "~", + ], + ) == ManyTypesA(i=5, s="hello", f=3.0, p=pathlib.Path("~")) + + +def test_pydantic_v1() -> None: + class ManyTypesA(v1.BaseModel): + i: int + s: str = "hello" + f: float = v1.Field(default_factory=lambda: 3.0) + p: pathlib.Path + + # We can directly pass a dataclass to `tyro.cli()`: + assert tyro.cli( + ManyTypesA, + args=[ + "--i", + "5", + "--p", + "~", + ], + ) == ManyTypesA(i=5, s="hello", f=3.0, p=pathlib.Path("~")) + + +def test_pydantic_helptext() -> None: + class Helptext(BaseModel): + """This docstring should be printed as a description.""" + + x: int = Field(description="Documentation 1") + + y: int = Field(description="Documentation 2") + + z: int = Field(description="Documentation 3") + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + assert tyro._strings.strip_ansi_sequences(cast(str, Helptext.__doc__)) in helptext + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_pydantic_suppress_base_model_helptext() -> None: + class Helptext(BaseModel): + x: int = Field(description="Documentation 1") + + y: int = Field(description="Documentation 2") + + z: int = Field(description="Documentation 3") + + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(Helptext, args=["--help"]) + helptext = f.getvalue() + + assert "Create a new model by parsing and validating" not in helptext + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +class HelptextWithFieldDocstring(BaseModel): + """This docstring should be printed as a description.""" + + x: int + """Documentation 1""" + + y: int = Field(description="Documentation 2") + + z: int = Field(description="Documentation 3") + + +def test_pydantic_field_helptext_from_docstring() -> None: + f = io.StringIO() + with pytest.raises(SystemExit): + with contextlib.redirect_stdout(f): + tyro.cli(HelptextWithFieldDocstring, args=["--help"]) + helptext = f.getvalue() + assert ( + tyro._strings.strip_ansi_sequences( + cast(str, HelptextWithFieldDocstring.__doc__) + ) + in helptext + ) + + assert "Documentation 1" in helptext + assert "Documentation 2" in helptext + assert "Documentation 3" in helptext + + +def test_pydantic_positional_annotation() -> None: + class AnnotatedAsPositional(BaseModel): + name: tyro.conf.Positional[str] + """This is annotated as a positional argument.""" + + result = tyro.cli(AnnotatedAsPositional, args=["myname"]) + assert isinstance(result, AnnotatedAsPositional) + + +def test_pydantic_alias() -> None: + class AliasCfg(BaseModel): + alias: Annotated[str, tyro.conf.arg(aliases=["-a"])] + + assert tyro.cli(AliasCfg, args=["--alias", "3"]) == AliasCfg(alias="3") + assert tyro.cli(AliasCfg, args=["-a", "3"]) == AliasCfg(alias="3") + + +def test_pydantic_default_instance() -> None: + class Inside(BaseModel): + x: int = 1 + + class Outside(BaseModel): + i: "Inside" = Inside(x=2) + + assert tyro.cli(Outside, args=[]).i.x == 2, ( + "Expected x value from the default instance", + ) + assert tyro.cli(Outside, args=["--i.x", "3"]).i.x == 3 + + +def test_pydantic_nested_default_instance() -> None: + class Inside(BaseModel): + x: int = 1 + + class Middle(BaseModel): + i: Inside + + class Outside(BaseModel): + m: Middle = Middle(i=Inside(x=2)) + + assert tyro.cli(Outside, args=[]).m.i.x == 2, ( + "Expected x value from the default instance", + ) + assert tyro.cli(Outside, args=["--m.i.x", "3"]).m.i.x == 3 + + +# Updating forward references in Pydantic v1 requires that these classes are global. + + +class InsideV1(v1.BaseModel): + x: int = 1 + + +class MiddleV1(v1.BaseModel): + i: InsideV1 + + +def test_pydantic_v1_nested_default_instance() -> None: + class Outside(v1.BaseModel): + m: MiddleV1 = MiddleV1(i=InsideV1(x=2)) + + assert tyro.cli(Outside, args=[]).m.i.x == 2, ( + "Expected x value from the default instance", + ) + assert tyro.cli(Outside, args=["--m.i.x", "3"]).m.i.x == 3 diff --git a/tests/test_pydantic_with_newtype.py b/tests/test_pydantic_with_newtype.py new file mode 100644 index 000000000..2b4d71d30 --- /dev/null +++ b/tests/test_pydantic_with_newtype.py @@ -0,0 +1,37 @@ +from typing import NewType, Tuple + +import pydantic +from pydantic import BaseModel +from typing_extensions import Annotated + +import tyro + +Microliter = NewType("Microliter", int) + + +class Measurements(BaseModel): + single: Microliter = pydantic.Field(10) + renamed_single: Annotated[Microliter, tyro.conf.arg(name="other_single")] = ( + pydantic.Field(10) + ) + pair: Tuple[Microliter, Microliter] = pydantic.Field((20, 30)) + + +IncorrectMeasurements = NewType("IncorrectMeasurements", Measurements) + + +def test_pydantic_with_newtype(): + assert tyro.cli( + IncorrectMeasurements, args="--single 1 --pair 2 3".split(" ") + ) == Measurements( + single=Microliter(1), + renamed_single=Microliter(10), + pair=(Microliter(2), Microliter(3)), + ) + assert tyro.cli( + IncorrectMeasurements, args="--single 1 --other-single 5".split(" ") + ) == Measurements( + single=Microliter(1), + renamed_single=Microliter(5), + pair=(Microliter(20), Microliter(30)), + ) diff --git a/tests/test_self_type.py b/tests/test_self_type.py new file mode 100644 index 000000000..e1fd2270c --- /dev/null +++ b/tests/test_self_type.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import pytest +from typing_extensions import Self + +import tyro + + +class SomeClass: + def __init__(self, a: int, b: int) -> None: + self.a = a + self.b = b + + def method1(self, x: Self) -> None: + self.effect = x + + @classmethod + def method2(cls, x: Self) -> SomeClass: + return x + + # Self is not valid in static methods. + # https://peps.python.org/pep-0673/#valid-locations-for-self + # + # @staticmethod + # def method3(x: Self) -> TestClass: + # return x + + +class SomeSubclass(SomeClass): ... + + +def test_method() -> None: + x = SomeClass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method1, args=[]) + + assert tyro.cli(x.method1, args="--x.a 3 --x.b 3".split(" ")) is None + assert x.effect.a == 3 and x.effect.b == 3 + assert isinstance(x, SomeClass) + + +def test_classmethod() -> None: + x = SomeClass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method2, args=[]) + with pytest.raises(SystemExit): + tyro.cli(SomeClass.method2, args=[]) + + y = tyro.cli(x.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + y = tyro.cli(SomeClass.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + +def test_subclass_method() -> None: + x = SomeSubclass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method1, args=[]) + + assert tyro.cli(x.method1, args="--x.a 3 --x.b 3".split(" ")) is None + assert x.effect.a == 3 and x.effect.b == 3 + assert isinstance(x, SomeSubclass) + + y = tyro.cli(x.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + +def test_subclass_classmethod() -> None: + x = SomeSubclass(0, 0) + with pytest.raises(SystemExit): + tyro.cli(x.method2, args=[]) + with pytest.raises(SystemExit): + tyro.cli(SomeSubclass.method2, args=[]) + + y = tyro.cli(x.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) + + y = tyro.cli(SomeSubclass.method2, args="--x.a 3 --x.b 3".split(" ")) + assert y.a == 3 + assert y.b == 3 + assert isinstance(y, SomeClass) diff --git a/tests/test_strings.py b/tests/test_strings.py index 1f3635c69..ff7669a75 100644 --- a/tests/test_strings.py +++ b/tests/test_strings.py @@ -1,4 +1,4 @@ -import dcargs._strings as _strings +from tyro import _strings def test_words_from_name(): @@ -6,4 +6,65 @@ def test_words_from_name(): assert ( _strings.hyphen_separated_from_camel_case("my-http-server") == "my-http-server" ) - assert _strings.hyphen_separated_from_camel_case("MyHttpServer") == "my-http-server" + + +def test_make_field_name(): + assert _strings.make_field_name(["hello", "world"]) == "hello.world" + assert _strings.make_field_name(["hello_world", "world"]) == "hello-world.world" + assert ( + _strings.make_field_name(["hello_world", "___hello_world"]) + == "hello-world.___hello-world" + ) + assert ( + _strings.make_field_name(["hello_world", "---hello_world"]) + == "hello-world.---hello-world" + ) + + +def test_postprocess_helptext(): + assert _strings.remove_single_line_breaks("hello world") == "hello world" + assert _strings.remove_single_line_breaks("hello\nworld") == "hello world" + assert _strings.remove_single_line_breaks("hello \nworld") == "hello world" + assert _strings.remove_single_line_breaks("hello\n\nworld") == "hello\n\nworld" + assert ( + _strings.remove_single_line_breaks( + "a paragraph:\nSentence one.\nSentence two.\nSentence three.\n" + ) + == "a paragraph: Sentence one. Sentence two. Sentence three." + ) + assert ( + _strings.remove_single_line_breaks( + "a bulleted list:\n" + "- The first problem.\n" + "- The second problem.\n" + "- The third problem.\n" + ) + == "a bulleted list:\n" + "- The first problem.\n" + "- The second problem.\n" + "- The third problem." + ) + assert ( + _strings.remove_single_line_breaks( + "an indented list:\n" + " The first problem.\n" + " The second problem.\n" + " The third problem.\n" + ) + == "an indented list:\n" + " The first problem.\n" + " The second problem.\n" + " The third problem." + ) + assert ( + _strings.remove_single_line_breaks( + "a numbered list:\n" + "1. The first problem.\n" + "2. The second problem.\n" + "3. The third problem.\n" + ) + == "a numbered list:\n" + "1. The first problem.\n" + "2. The second problem.\n" + "3. The third problem." + ) diff --git a/tests/test_tuple_with_subcommands.py b/tests/test_tuple_with_subcommands.py new file mode 100644 index 000000000..88f1eba5d --- /dev/null +++ b/tests/test_tuple_with_subcommands.py @@ -0,0 +1,123 @@ +# mypy: disable-error-code="call-overload,misc" +# +# Mypy errors from passing union types directly into tyro.cli() as Type[T]. We would +# benefit from TypeForm[T]: https://github.com/python/mypy/issues/9773 +"""Tests adapted from https://github.com/brentyi/tyro/issues/89, which catches edge +cases when combining nested tuple types, renamed arguments, and subcommands. + +Largely written by @wookayin. +""" + +import dataclasses +from pathlib import Path +from typing import Generic, Tuple, TypeVar, Union + +from typing_extensions import Annotated + +import tyro.conf + +T = TypeVar("T") + + +@dataclasses.dataclass +class Checkout(Generic[T]): + """Check out a branch.""" + + branch: T + + +@dataclasses.dataclass +class Commit: + """Commit something.""" + + input: tyro.conf.Positional[Path] + + +@dataclasses.dataclass +class Arg: + verbose: bool = True + + +def test_case1() -> None: + o = tyro.cli( + Union[ + Checkout[str], + Commit, + ], + args=["commit", "./path.txt"], + ) + assert o == Commit(input=Path("path.txt")) + + +def test_case2() -> None: + arg, action = tyro.cli( + Tuple[ + Arg, + Annotated[ + Union[ + Checkout[str], + Commit, + ], + tyro.conf.arg(name=""), + ], + ], + args=["commit", "./path.txt"], + ) + + assert isinstance(arg, Arg) + assert isinstance(action, Commit) + assert action.input == Path("./path.txt") + + +def test_case3() -> None: + o = tyro.cli( + Tuple[ + Annotated[ + Arg, + tyro.conf.arg(name=""), + ], + Annotated[ + Union[ + Annotated[Checkout[str], tyro.conf.subcommand(name="checkout")], + Annotated[Commit, tyro.conf.subcommand(name="commit")], + ], + tyro.conf.arg(name=""), + ], + ], + args=["commit", "./path.txt"], + ) + assert o == (Arg(), Commit(Path("./path.txt"))) + + +def test_case4() -> None: + o = tyro.cli( + Tuple[ + Annotated[ + Union[ + Annotated[Checkout[str], tyro.conf.subcommand(name="checkout")], + Annotated[Commit, tyro.conf.subcommand(name="commit")], + ], + tyro.conf.arg(name=""), + ] + ], + args=["commit", "./path.txt"], + ) + assert o == (Commit(Path("./path.txt")),) + + +def test_case5() -> None: + assert tyro.cli( + tyro.conf.OmitArgPrefixes[ + Tuple[ + Union[ + Annotated[ + Checkout[str], + tyro.conf.subcommand(prefix_name=False), + ], + Annotated[Commit, tyro.conf.subcommand(prefix_name=False)], + ], + Arg, + ] + ], + args=["--no-verbose", "checkout-str", "--branch", "branch"], + ) == (Checkout("branch"), Arg(False)) diff --git a/tests/test_typeddict_readonly_min_py38.py b/tests/test_typeddict_readonly_min_py38.py new file mode 100644 index 000000000..ae26777bf --- /dev/null +++ b/tests/test_typeddict_readonly_min_py38.py @@ -0,0 +1,21 @@ +import pytest +from typing_extensions import ReadOnly, TypedDict + +import tyro + + +def test_read_only_typeddict() -> None: + class ManyTypesTypedDict(TypedDict): + i: ReadOnly[int] # type: ignore + s: ReadOnly[ReadOnly[str]] # type: ignore + + assert tyro.cli( + ManyTypesTypedDict, + args="--i 5 --s 5".split(" "), + ) == dict(i=5, s="5") + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--i 5".split(" ")) + + with pytest.raises(SystemExit): + tyro.cli(ManyTypesTypedDict, args="--s 5".split(" ")) diff --git a/tests/test_union_from_mapping.py b/tests/test_union_from_mapping.py new file mode 100644 index 000000000..2eb8ff42b --- /dev/null +++ b/tests/test_union_from_mapping.py @@ -0,0 +1,49 @@ +import dataclasses +from typing import TYPE_CHECKING, Optional + +import tyro + + +@dataclasses.dataclass +class A: + x: int + + +def test_union_from_mapping(): + base_configs = { + "one": A(1), + "two": A(2), + "three": A(3), + } + ConfigUnion = tyro.extras.subcommand_type_from_defaults(base_configs) + + assert tyro.cli(ConfigUnion, args="one".split(" ")) == A(1) + assert tyro.cli(ConfigUnion, args="two".split(" ")) == A(2) + assert tyro.cli(ConfigUnion, args="two --x 4".split(" ")) == A(4) + assert tyro.cli(ConfigUnion, args="three".split(" ")) == A(3) + + +def test_union_from_mapping_in_function(): + base_configs = { + "one": A(1), + "two": A(2), + "three": A(3), + } + + if TYPE_CHECKING: + ConfigUnion = A + else: + ConfigUnion = tyro.extras.subcommand_type_from_defaults(base_configs) + + def main(config: ConfigUnion, flag: bool = False) -> Optional[A]: # type: ignore + if flag: + return config + return None + + assert tyro.cli(main, args="--flag config:one".split(" ")) == A(1) + assert tyro.cli(main, args="--flag config:one --config.x 3".split(" ")) == A(3) + assert tyro.cli(main, args="config:one --config.x 1".split(" ")) is None + + assert tyro.cli(main, args="--flag config:two".split(" ")) == A(2) + assert tyro.cli(main, args="--flag config:two --config.x 3".split(" ")) == A(3) + assert tyro.cli(main, args="config:two --config.x 1".split(" ")) is None diff --git a/tests/test_unsafe_cache.py b/tests/test_unsafe_cache.py new file mode 100644 index 000000000..35cb41741 --- /dev/null +++ b/tests/test_unsafe_cache.py @@ -0,0 +1,31 @@ +from tyro import _unsafe_cache + + +def test_unsafe_cache(): + x = 0 + + @_unsafe_cache.unsafe_cache(maxsize=2) + def f(dummy: int): + nonlocal x + x += 1 + + f(0) + f(0) + f(0) + assert x == 1 + f(1) + f(1) + f(1) + assert x == 2 + f(0) + f(0) + f(0) + assert x == 2 + f(2) + f(2) + f(2) + assert x == 3 + f(0) + f(0) + f(0) + assert x == 4 diff --git a/tests/test_unsupported_but_should_work.py b/tests/test_unsupported_but_should_work.py new file mode 100644 index 000000000..74327ac80 --- /dev/null +++ b/tests/test_unsupported_but_should_work.py @@ -0,0 +1,59 @@ +"""Tests for features that are not officially features, but should work. + +Includes things like omegaconf.MISSING, attrs, etc, which mostly work but either likely +have corner cases or just seem sketchy. +""" + +from typing import Tuple + +import omegaconf +import pytest + +import tyro +import tyro._strings + + +def test_omegaconf_missing(): + """Passing in a omegaconf.MISSING default; this will mark an argument as required.""" + + def main( + required_a: int, + optional: int = 3, + required_b: int = None, # type: ignore + ) -> Tuple[int, int, int]: + return (required_a, optional, required_b) # type: ignore + + assert tyro.cli( + main, args="--required-a 3 --optional 4 --required-b 5".split(" ") + ) == (3, 4, 5) + assert tyro.cli(main, args="--required-a 3 --required-b 5".split(" ")) == ( + 3, + 3, + 5, + ) + + with pytest.raises(SystemExit): + tyro.cli(main, args="--required-a 3 --optional 4") + with pytest.raises(SystemExit): + tyro.cli(main, args="--required-a 3") + + def main2( + required_a: int, + optional: int = 3, + required_b: int = omegaconf.MISSING, + ) -> Tuple[int, int, int]: + return (required_a, optional, required_b) + + assert tyro.cli( + main2, args="--required-a 3 --optional 4 --required-b 5".split(" ") + ) == (3, 4, 5) + assert tyro.cli(main2, args="--required-a 3 --required-b 5".split(" ")) == ( + 3, + 3, + 5, + ) + + with pytest.raises(SystemExit): + tyro.cli(main2, args="--required-a 3 --optional 4") + with pytest.raises(SystemExit): + tyro.cli(main2, args="--required-a 3")