Fix inheritance edge cases for type parameter resolution (#230)

* Fix issue with generics when TypeVars are reused

* Refine

* Fix type param resolution for superclass __init__ / __new__

* pydantic test

* attrs

* Tests

* Python 3.8 hack

* Minor update for coverage

* ruff

* mypy

* upgrade ruff

* remove redundant annotations in tests

* no cover for py38 hack

* run test gen

Author: brentyi

src/tyro/_fields.py

@@ -7,6 +7,7 @@
 import dataclasses
 import functools
 import inspect
+import sys
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import docstring_parser
@@ -74,7 +75,7 @@ def make(
         helptext: Optional[str],
         call_argname_override: Optional[Any] = None,
     ):
-        # Resolve generics.
+        # Resolve type parameters.
         typ = _resolver.TypeParamResolver.concretize_type_params(typ)
 
         # Narrow types.
@@ -265,14 +266,17 @@ def _field_list_from_function(
 
     # Unwrap functools.wraps and functools.partial.
     done = False
+    functools_marker = False
     while not done:
         done = True
         if hasattr(f, "__wrapped__"):
             f = f.__wrapped__  # type: ignore
             done = False
+            functools_marker = True
         if isinstance(f, functools.partial):
             f = f.func
             done = False
+            functools_marker = True
 
     # Check for abstract classes.
     if inspect.isabstract(f):
@@ -282,11 +286,71 @@ def _field_list_from_function(
     # signature. But the docstrings may still be in the class signature itself.
     f_before_init_unwrap = f
 
+    hints = None
+
     if inspect.isclass(f):
+        signature_func = None
         if hasattr(f, "__init__") and f.__init__ is not object.__init__:
-            f = f.__init__  # type: ignore
+            signature_func = "__init__"
         elif hasattr(f, "__new__") and f.__new__ is not object.__new__:
-            f = f.__new__
+            signature_func = "__new__"
+
+        if signature_func is not None:
+            # Get the __init__ / __new__ method from the class, as well as the
+            # class that contains it.
+            #
+            # We call this the "signature function", because it's the function
+            # that we use to instantiate the class.
+            orig_cls = f
+            base_cls_with_signature = None
+            for base_cls_with_signature in inspect.getmro(f):
+                if signature_func in base_cls_with_signature.__dict__:
+                    f = getattr(base_cls_with_signature, signature_func)
+                    break
+            assert base_cls_with_signature is not None
+            assert f is not orig_cls
+
+            # For older versions of Python, the signature returned above (when
+            # passed through generics base classes) will sometimes be (*args,
+            # **kwargs).
+            #
+            # This is a hack. We can remove it if we deprecate Python 3.8 support.
+            if sys.version_info < (3, 9) and not functools_marker:  # pragma: no cover
+                params = list(inspect.signature(f).parameters.values())[1:]
+
+            # Get hints for the signature function by recursing through the
+            # inheritance tree. This is needed to correctly resolve type
+            # parameters, which can be set anywhere between the input class and
+            # the class where the __init__ or __new__ method is defined.
+            def get_hints_for_signature_func(cls):
+                typevar_context = _resolver.TypeParamResolver.get_assignment_context(
+                    cls
+                )
+                cls = typevar_context.origin_type
+                with typevar_context:
+                    if cls is base_cls_with_signature:
+                        return _resolver.get_type_hints_resolve_type_params(
+                            f, include_extras=True
+                        )
+                    for base_cls in (
+                        cls.__orig_bases__
+                        if hasattr(cls, "__orig_bases__")
+                        else cls.__bases__
+                    ):
+                        if not issubclass(
+                            _resolver.unwrap_origin_strip_extras(base_cls),
+                            base_cls_with_signature,
+                        ):
+                            continue
+                        return get_hints_for_signature_func(
+                            _resolver.TypeParamResolver.concretize_type_params(base_cls)
+                        )
+
+                assert False, (
+                    "We couldn't find the base class. This seems like a bug in tyro."
+                )
+
+            hints = get_hints_for_signature_func(orig_cls)
 
     # Get type annotations, docstrings.
     docstring = inspect.getdoc(f)
@@ -296,11 +360,13 @@ def _field_list_from_function(
             docstring_from_arg_name[param_doc.arg_name] = param_doc.description
     del docstring
 
+    # Get hints if we haven't done it already.
     # 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 UnsupportedStructTypeMessage(f"Could not get hints for {f}!")
+    if hints is None:
+        try:
+            hints = _resolver.get_type_hints_resolve_type_params(f, include_extras=True)
+        except TypeError:
+            return UnsupportedStructTypeMessage(f"Could not get hints for {f}!")
 
     field_list = []
     for param in params:

src/tyro/_resolver.py

@@ -14,6 +14,7 @@
     Callable,
     ClassVar,
     Dict,
+    Generic,
     List,
     Sequence,
     Set,
@@ -76,14 +77,14 @@ def is_dataclass(cls: Union[TypeForm, Callable]) -> bool:
     return dataclasses.is_dataclass(unwrap_origin_strip_extras(cls))  # type: ignore
 
 
-@_unsafe_cache.unsafe_cache(maxsize=1024)
+# @_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(
+    annotations = get_type_hints_resolve_type_params(
         cast(Callable, cls), include_extras=True
     )
     for field in getattr(cls, "__dataclass_fields__").values():
@@ -566,24 +567,6 @@ def resolve_generic_types(
         assert len(typevars) == len(typevar_values)
         typ = origin_cls
         type_from_typevar.update(dict(zip(typevars, typevar_values)))
-    elif (
-        # Apply some heuristics for generic types. Should revisit this.
-        hasattr(typ, "__parameters__") and hasattr(typ.__parameters__, "__len__")  # type: ignore
-    ):
-        typevars = typ.__parameters__  # type: ignore
-        typevar_values = tuple(type_from_typevar_constraints(x) for x in typevars)
-        assert len(typevars) == len(typevar_values)
-        type_from_typevar.update(dict(zip(typevars, typevar_values)))
-
-    if hasattr(typ, "__orig_bases__"):
-        bases = getattr(typ, "__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 typ, type_from_typevar
@@ -594,11 +577,56 @@ def resolve_generic_types(
         )
 
 
-def get_type_hints_with_backported_syntax(
+def get_type_hints_resolve_type_params(
+    obj: Callable[..., Any], include_extras: bool = False
+) -> Dict[str, Any]:
+    """Variant of `typing.get_type_hints()` that resolves type parameters."""
+    if not inspect.isclass(obj):
+        return {
+            k: TypeParamResolver.concretize_type_params(v)
+            for k, v in _get_type_hints_backported_syntax(
+                obj, include_extras=include_extras
+            ).items()
+        }
+
+    typevar_context = TypeParamResolver.get_assignment_context(obj)
+    obj = typevar_context.origin_type
+    with typevar_context:
+        # Only include type hints that are explicitly defined in this class.
+        # The follow loop will handle superclasses.
+        out = {
+            x: TypeParamResolver.concretize_type_params(t)
+            for x, t in _get_type_hints_backported_syntax(
+                obj, include_extras=include_extras
+            ).items()
+            if x in obj.__annotations__
+        }
+
+    # We need to recurse into base classes in order to correctly resolve superclass parameters.
+    for base in obj.__orig_bases__ if hasattr(obj, "__orig_bases__") else obj.__bases__:  # type: ignore
+        base_typevar_context = TypeParamResolver.get_assignment_context(base)
+        if get_origin(base_typevar_context.origin_type) is Generic:
+            continue
+        with base_typevar_context:
+            base_hints = get_type_hints_resolve_type_params(
+                base_typevar_context.origin_type, include_extras=include_extras
+            )
+            out.update(
+                {
+                    x: TypeParamResolver.concretize_type_params(t)
+                    for x, t in base_hints.items()
+                }
+            )
+
+    return out
+
+
+def _get_type_hints_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)
 

src/tyro/constructors/_registry.py

@@ -173,11 +173,12 @@ def get_struct_spec(cls, type_info: StructTypeInfo) -> StructConstructorSpec | N
                 f"Invalid default instance for type {type_info.type}: {type_info.default}"
             )
 
-        for registry in cls._active_registries[::-1]:
-            for spec_factory in registry._struct_rules[::-1]:
-                maybe_spec = spec_factory(type_info)
-                if maybe_spec is not None:
-                    return maybe_spec
+        with type_info._typevar_context:
+            for registry in cls._active_registries[::-1]:
+                for spec_factory in registry._struct_rules[::-1]:
+                    maybe_spec = spec_factory(type_info)
+                    if maybe_spec is not None:
+                        return maybe_spec
 
         return None
 

src/tyro/constructors/_struct_spec.py

@@ -214,7 +214,7 @@ def typeddict_rule(info: StructTypeInfo) -> StructConstructorSpec | None:
         total = getattr(cls, "__total__", True)
         assert isinstance(total, bool)
         assert not valid_default_instance or isinstance(info.default, dict)
-        for name, typ in _resolver.get_type_hints_with_backported_syntax(
+        for name, typ in _resolver.get_type_hints_resolve_type_params(
             cls, include_extras=True
         ).items():
             typ_origin = get_origin(typ)
@@ -277,7 +277,13 @@ def attrs_rule(info: StructTypeInfo) -> StructConstructorSpec | None:
             return None
 
         # Resolve forward references in-place, if any exist.
-        attr.resolve_types(info.type)
+        # attr.resolve_types(info.type)
+
+        # We'll use our own type resolution system instead of attr's. This is
+        # primarily to improve generics support.
+        our_hints = _resolver.get_type_hints_resolve_type_params(
+            info.type, include_extras=True
+        )
 
         # Handle attr classes.
         field_list = []
@@ -301,7 +307,7 @@ def attrs_rule(info: StructTypeInfo) -> StructConstructorSpec | None:
             field_list.append(
                 StructFieldSpec(
                     name=name,
-                    type=attr_field.type,
+                    type=our_hints[name],
                     default=default,
                     helptext=_docstrings.get_field_docstring(info.type, name),
                 )
@@ -353,7 +359,7 @@ def namedtuple_rule(info: StructTypeInfo) -> StructConstructorSpec | None:
         field_list = []
         field_defaults = getattr(info.type, "_field_defaults", {})
 
-        for name, typ in _resolver.get_type_hints_with_backported_syntax(
+        for name, typ in _resolver.get_type_hints_resolve_type_params(
             info.type, include_extras=True
         ).items():
             default = field_defaults.get(name, MISSING_NONPROP)
@@ -534,7 +540,7 @@ def pydantic_rule(info: StructTypeInfo) -> StructConstructorSpec | None:
         ):
             # Pydantic 1.xx
             cls_cast = info.type
-            hints = _resolver.get_type_hints_with_backported_syntax(
+            hints = _resolver.get_type_hints_resolve_type_params(
                 info.type, include_extras=True
             )
             for pd1_field in cast(Dict[str, Any], cls_cast.__fields__).values():

tests/test_attrs.py

@@ -3,11 +3,12 @@
 import contextlib
 import io
 import pathlib
-from typing import cast
+from typing import Generic, TypeVar, cast
 
 import attr
 import pytest
 from attrs import define, field
+from helptext_utils import get_helptext_with_checks
 
 import tyro
 import tyro._strings
@@ -129,3 +130,24 @@ class ManyTypesB:
         args=["--i", "5"],
         default=ManyTypesB(i=5, s="5", f=2.0),
     ) == ManyTypesB(i=5, s="5", f=2.0)
+
+
+T = TypeVar("T")
+
+
+def test_attrs_inheritance_with_same_typevar() -> None:
+    @attr.s
+    class A(Generic[T]):
+        x: T = attr.ib()
+
+    @attr.s
+    class B(A[int], Generic[T]):
+        y: T = attr.ib()
+
+    assert "INT" in get_helptext_with_checks(B[int])
+    assert "STR" not in get_helptext_with_checks(B[int])
+    assert "STR" in get_helptext_with_checks(B[str])
+    assert "INT" in get_helptext_with_checks(B[str])
+
+    assert tyro.cli(B[str], args=["--x", "1", "--y", "2"]) == B(x=1, y="2")
+    assert tyro.cli(B[int], args=["--x", "1", "--y", "2"]) == B(x=1, y=2)