Repository: brentyi/tyro

Issue:
Attribute inheritance regression 0.9.8 vs 0.9.11
```py
# x.py
from dataclasses import dataclass
from typing import Annotated

import tyro

_dataset_map = {
    "alpaca": "tatsu-lab/alpaca",
    "alpaca_clean": "yahma/alpaca-cleaned",
    "alpaca_gpt4": "vicgalle/alpaca-gpt4",
}
_inv_dataset_map = {value: key for key, value in _dataset_map.items()}
_datasets = list(_dataset_map.keys())


HFDataset = Annotated[
    str,
    tyro.constructors.PrimitiveConstructorSpec(
        nargs=1,
        metavar="{" + ",".join(_datasets) + "}",
        instance_from_str=lambda args: _dataset_map[args[0]],
        is_instance=lambda instance: isinstance(instance, str)
        and instance in _inv_dataset_map,
        str_from_instance=lambda instance: [_inv_dataset_map[instance]],
        choices=tuple(_datasets),
    ),
    tyro.conf.arg(
        help_behavior_hint=lambda df: f"(default: {df}, run datasets.py for full options)"
    ),
]


@dataclass
class TrainConfig:
    dataset: str = "vicgalle/alpaca-gpt4"


@dataclass
class CLITrainerConfig(TrainConfig):
    dataset: HFDataset = "vicgalle/alpaca-gpt4"


if __name__ == "__main__":
    cfg = tyro.cli(CLITrainerConfig, config=(tyro.conf.ConsolidateSubcommandArgs,))
```
Something changed between the last versions and now tyro defaults to using the field types of the parent (TrainConfig) rather than CLITrainerConfig. On 0.9.8, `py x.py --help`:

```
usage: tyro_test.py [-h] [--dataset {alpaca,alpaca_clean,alpaca_gpt4}]

╭─ options ──────────────────────────────────────────────────────────────────╮
│ -h, --help        show this help message and exit                          │
│ --dataset {alpaca,alpaca_clean,alpaca_gpt4}                                │
│                   (default: alpaca_gpt4, run datasets.py for full options) │
╰────────────────────────────────────────────────────────────────────────────╯
```

While on 0.9.11, `py x.py --help`:
```
usage: tyro_test.py [-h] [--dataset STR]

╭─ options ────────────────────────────────────────────╮
│ -h, --help           show this help message and exit │
│ --dataset STR        (default: vicgalle/alpaca-gpt4) │
╰──────────────────────────────────────────────────────╯
```
This also happens for more complicated fields involving Unions: 

```py
# x.py
from dataclasses import dataclass
from typing import Annotated, Union

import tyro


@dataclass(frozen=True)
class OptimizerConfig:
    lr: float = 1e-1


@dataclass(frozen=True)
class AdamConfig(OptimizerConfig):
    adam_foo: float = 1.0


@dataclass(frozen=True)
class SGDConfig(OptimizerConfig):
    sgd_foo: float = 1.0


@dataclass
class TrainConfig:
    optimizer: OptimizerConfig = AdamConfig()


def _dummy_constructor() -> type[OptimizerConfig]:
    return Union[AdamConfig, SGDConfig]  # type: ignore


CLIOptimizerConfig = Annotated[
    OptimizerConfig,
    tyro.conf.arg(constructor_factory=_dummy_constructor),
]


@dataclass
class CLITrainerConfig(TrainConfig):
    optimizer: CLIOptimizerConfig = SGDConfig()


if __name__ == "__main__":
    cfg = tyro.cli(CLITrainerConfig, config=(tyro.conf.ConsolidateSubcommandArgs,))

```

On 0.9.8, `py x.py --help`:
```
usage: tyro_test.py [-h] [{optimizer:adam-config,optimizer:sgd-config}]

╭─ options ─────────────────────────────────────────╮
│ -h, --help        show this help message and exit │
╰───────────────────────────────────────────────────╯
╭─ optional subcommands ────────────────────────────╮
│ (default: optimizer:sgd-config)                   │
│ ──────────────────────────────────────────────    │
│ [{optimizer:adam-config,optimizer:sgd-config}]    │
│     optimizer:adam-config                         │
│     optimizer:sgd-config                          │
╰───────────────────────────────────────────────────╯
```

While on 0.9.11:
```
usage: tyro_test.py [-h] [--optimizer.lr FLOAT] [--optimizer.sgd-foo FLOAT]

╭─ options ───────────────────────────────────────────────╮
│ -h, --help              show this help message and exit │
╰─────────────────────────────────────────────────────────╯
╭─ optimizer options ─────────────────────────────────────╮
│ --optimizer.lr FLOAT    (default: 0.1)                  │
│ --optimizer.sgd-foo FLOAT                               │
│                         (default: 1.0)                  │
╰─────────────────────────────────────────────────────────╯
```
Note that in this case it does seem to be picking up correctly on the default value being changed (from parent to child), but it ignores the type `CLIOptimizerConfig` which should behave like a Union and instead behaves as if we still have an OptimizerConfig as in the parent dataclass. To make sure this isn't because of the custom constructor, with a simpler Union:

```py
from dataclasses import dataclass

import tyro


@dataclass(frozen=True)
class OptimizerConfig:
    lr: float = 1e-1


@dataclass(frozen=True)
class AdamConfig(OptimizerConfig):
    adam_foo: float = 1.0


@dataclass(frozen=True)
class SGDConfig(OptimizerConfig):
    sgd_foo: float = 1.0


@dataclass
class TrainConfig:
    optimizer: OptimizerConfig = AdamConfig()


@dataclass
class CLITrainerConfig(TrainConfig):
    optimizer: AdamConfig | SGDConfig = SGDConfig()


if __name__ == "__main__":
    cfg = tyro.cli(CLITrainerConfig, config=(tyro.conf.ConsolidateSubcommandArgs,))
```
0.9.8:
```
usage: tyro_test.py [-h] [{optimizer:adam-config,optimizer:sgd-config}]

╭─ options ─────────────────────────────────────────╮
│ -h, --help        show this help message and exit │
╰───────────────────────────────────────────────────╯
╭─ optional subcommands ────────────────────────────╮
│ (default: optimizer:sgd-config)                   │
│ ──────────────────────────────────────────────    │
│ [{optimizer:adam-config,optimizer:sgd-config}]    │
│     optimizer:adam-config                         │
│     optimizer:sgd-config                          │
╰───────────────────────────────────────────────────╯
```
0.9.11:
```
usage: tyro_test.py [-h] [--optimizer.lr FLOAT] [--optimizer.sgd-foo FLOAT]

╭─ options ───────────────────────────────────────────────╮
│ -h, --help              show this help message and exit │
╰─────────────────────────────────────────────────────────╯
╭─ optimizer options ─────────────────────────────────────╮
│ --optimizer.lr FLOAT    (default: 0.1)                  │
│ --optimizer.sgd-foo FLOAT                               │
│                         (default: 1.0)                  │
╰─────────────────────────────────────────────────────────╯
```
This was tested on python 3.12.0, which I know had some issues from a previous bug. ~~I'll try to test it on some version > 3.12.0 later today.~~ Same issue on 3.12.5.

Discussion:
@brentyi: Thanks as usual! The examples are immensely helpful; it's also cool to see `PrimitiveConstructorSpec` in use.

The problem was flawed tree traversal logic, which was introduced in our efforts to support inherited type parameters.... sorry for the inconvenience!! It should be fixed in `0.9.12`.@mirceamironenco: Wow thanks for the super fast fix. Tested it and it works great so far!

> it's also cool to see PrimitiveConstructorSpec in use.

This was a nice addition, and easy to grok.

Solution patch:
diff --git a/src/tyro/_resolver.py b/src/tyro/_resolver.py
index 0b603535..e1cb9828 100644
--- a/src/tyro/_resolver.py
+++ b/src/tyro/_resolver.py
@@ -673,7 +673,13 @@ def get_type_hints_resolve_type_params(
             for x, t in _get_type_hints_backported_syntax(
                 obj, include_extras=include_extras
             ).items()
-            if x in obj.__annotations__
+            # Only include type hints that are explicitly defined in this class.
+            #
+            # Why `cls.__dict__.__annotations__` instead of `cls.__annotations__`?
+            # Because in Python 3.8 and earlier, `cls.__annotations__`
+            # recursively merges parent class annotations.
+            # See this issue: https://github.com/python/cpython/issues/99535
+            if x in obj.__dict__.get("__annotations__", {})
         }
 
     # We need to recurse into base classes in order to correctly resolve superclass parameters.
@@ -689,6 +695,11 @@ def get_type_hints_resolve_type_params(
                 {
                     x: TypeParamResolver.concretize_type_params(t)
                     for x, t in base_hints.items()
+                    # Include type hints that are not assigned earlier in the MRO.
+                    #
+                    # This needs to be recursive (include parents of parents),
+                    # so we shouldn't filter by local __annotations__.
+                    if x not in out
                 }
             )
 


Test patch:
diff --git a/tests/test_conf.py b/tests/test_conf.py
index 8aec850f..cc019960 100644
--- a/tests/test_conf.py
+++ b/tests/test_conf.py
@@ -1794,3 +1794,94 @@ def test_suppress_in_union() -> None:
     with pytest.raises(SystemExit):
         # ConfigB has a required argument.
         assert tyro.cli(main, args=["x:config-b"])
+
+
+_dataset_map = {
+    "alpaca": "tatsu-lab/alpaca",
+    "alpaca_clean": "yahma/alpaca-cleaned",
+    "alpaca_gpt4": "vicgalle/alpaca-gpt4",
+}
+_inv_dataset_map = {value: key for key, value in _dataset_map.items()}
+_datasets = list(_dataset_map.keys())
+
+HFDataset = Annotated[
+    str,
+    tyro.constructors.PrimitiveConstructorSpec(
+        nargs=1,
+        metavar="{" + ",".join(_datasets) + "}",
+        instance_from_str=lambda args: _dataset_map[args[0]],
+        is_instance=lambda instance: isinstance(instance, str)
+        and instance in _inv_dataset_map,
+        str_from_instance=lambda instance: [_inv_dataset_map[instance]],
+        choices=tuple(_datasets),
+    ),
+    tyro.conf.arg(
+        help_behavior_hint=lambda df: f"(default: {df}, run datasets.py for full options)"
+    ),
+]
+
+
+def test_annotated_attribute_inheritance() -> None:
+    """From @mirceamironenco.
+
+    https://github.com/brentyi/tyro/issues/239"""
+
+    @dataclasses.dataclass(frozen=True)
+    class TrainConfig:
+        dataset: str = "vicgalle/alpaca-gpt4"
+
+    @dataclasses.dataclass(frozen=True)
+    class CLITrainerConfig(TrainConfig):
+        dataset: HFDataset = "vicgalle/alpaca-gpt4"
+
+    assert "{alpaca,alpaca_clean,alpaca_gpt4}" in get_helptext_with_checks(
+        CLITrainerConfig
+    )
+    assert (
+        "default: alpaca_gpt4, run datasets.py for full options"
+        in get_helptext_with_checks(CLITrainerConfig)
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class OptimizerConfig:
+    lr: float = 1e-1
+
+
+@dataclasses.dataclass(frozen=True)
+class AdamConfig(OptimizerConfig):
+    adam_foo: float = 1.0
+
+
+@dataclasses.dataclass(frozen=True)
+class SGDConfig(OptimizerConfig):
+    sgd_foo: float = 1.0
+
+
+@dataclasses.dataclass
+class TrainConfig:
+    optimizer: OptimizerConfig = AdamConfig()
+
+
+def _dummy_constructor() -> Type[OptimizerConfig]:
+    return Union[AdamConfig, SGDConfig]  # type: ignore
+
+
+CLIOptimizerConfig = Annotated[
+    OptimizerConfig,
+    tyro.conf.arg(constructor_factory=_dummy_constructor),
+]
+
+
+def test_attribute_inheritance_2() -> None:
+    """From @mirceamironenco.
+
+    https://github.com/brentyi/tyro/issues/239"""
+
+    @dataclasses.dataclass
+    class CLITrainerConfig(TrainConfig):
+        optimizer: CLIOptimizerConfig = SGDConfig()
+
+    assert "[{optimizer:adam-config,optimizer:sgd-config}]" in get_helptext_with_checks(
+        CLITrainerConfig
+    )
diff --git a/tests/test_py311_generated/test_conf_generated.py b/tests/test_py311_generated/test_conf_generated.py
index c917c389..02ec3bb4 100644
--- a/tests/test_py311_generated/test_conf_generated.py
+++ b/tests/test_py311_generated/test_conf_generated.py
@@ -1800,3 +1800,94 @@ def test_suppress_in_union() -> None:
     with pytest.raises(SystemExit):
         # ConfigB has a required argument.
         assert tyro.cli(main, args=["x:config-b"])
+
+
+_dataset_map = {
+    "alpaca": "tatsu-lab/alpaca",
+    "alpaca_clean": "yahma/alpaca-cleaned",
+    "alpaca_gpt4": "vicgalle/alpaca-gpt4",
+}
+_inv_dataset_map = {value: key for key, value in _dataset_map.items()}
+_datasets = list(_dataset_map.keys())
+
+HFDataset = Annotated[
+    str,
+    tyro.constructors.PrimitiveConstructorSpec(
+        nargs=1,
+        metavar="{" + ",".join(_datasets) + "}",
+        instance_from_str=lambda args: _dataset_map[args[0]],
+        is_instance=lambda instance: isinstance(instance, str)
+        and instance in _inv_dataset_map,
+        str_from_instance=lambda instance: [_inv_dataset_map[instance]],
+        choices=tuple(_datasets),
+    ),
+    tyro.conf.arg(
+        help_behavior_hint=lambda df: f"(default: {df}, run datasets.py for full options)"
+    ),
+]
+
+
+def test_annotated_attribute_inheritance() -> None:
+    """From @mirceamironenco.
+
+    https://github.com/brentyi/tyro/issues/239"""
+
+    @dataclasses.dataclass(frozen=True)
+    class TrainConfig:
+        dataset: str = "vicgalle/alpaca-gpt4"
+
+    @dataclasses.dataclass(frozen=True)
+    class CLITrainerConfig(TrainConfig):
+        dataset: HFDataset = "vicgalle/alpaca-gpt4"
+
+    assert "{alpaca,alpaca_clean,alpaca_gpt4}" in get_helptext_with_checks(
+        CLITrainerConfig
+    )
+    assert (
+        "default: alpaca_gpt4, run datasets.py for full options"
+        in get_helptext_with_checks(CLITrainerConfig)
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class OptimizerConfig:
+    lr: float = 1e-1
+
+
+@dataclasses.dataclass(frozen=True)
+class AdamConfig(OptimizerConfig):
+    adam_foo: float = 1.0
+
+
+@dataclasses.dataclass(frozen=True)
+class SGDConfig(OptimizerConfig):
+    sgd_foo: float = 1.0
+
+
+@dataclasses.dataclass
+class TrainConfig:
+    optimizer: OptimizerConfig = AdamConfig()
+
+
+def _dummy_constructor() -> Type[OptimizerConfig]:
+    return AdamConfig | SGDConfig  # type: ignore
+
+
+CLIOptimizerConfig = Annotated[
+    OptimizerConfig,
+    tyro.conf.arg(constructor_factory=_dummy_constructor),
+]
+
+
+def test_attribute_inheritance_2() -> None:
+    """From @mirceamironenco.
+
+    https://github.com/brentyi/tyro/issues/239"""
+
+    @dataclasses.dataclass
+    class CLITrainerConfig(TrainConfig):
+        optimizer: CLIOptimizerConfig = SGDConfig()
+
+    assert "[{optimizer:adam-config,optimizer:sgd-config}]" in get_helptext_with_checks(
+        CLITrainerConfig
+    )