Merge all of the code that calls __annotate__/evaluate_value

Now we only have one place where we replace the closure of a function.

Author: msullivan

tests/test_type_eval.py

@@ -404,7 +404,7 @@ def test_getmember_01():
 
 def test_getmember_02():
     class C:
-        def f[T](self, x: T) -> OnlyIntToSet[T]: ...
+        def f[TX](self, x: TX) -> OnlyIntToSet[TX]: ...
 
     m = eval_typing(GetMember[C, Literal["f"]])
     assert eval_typing(GetName[m]) == Literal["f"]

typemap/type_eval/_apply_generic.py

@@ -187,62 +187,61 @@ def make_func(
     return new_func
 
 
-def _get_closure_types(af: types.FunctionType) -> dict[str, type]:
-    # Generate a fallback mapping of closure classes.
-    # This is needed for locally defined generic types which reference
-    # themselves in their type annotations.
-    if not af.__closure__:
-        return {}
-    return {
-        name: variable.cell_contents
-        for name, variable in zip(
-            af.__code__.co_freevars, af.__closure__, strict=True
-        )
-    }
-
-
 EXCLUDED_ATTRIBUTES = typing.EXCLUDED_ATTRIBUTES - {'__init__'}  # type: ignore[attr-defined]
 
 
-def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
-    annos: dict[str, Any] = {}
-    dct: dict[str, Any] = {}
-
-    if af := typing.cast(
-        types.FunctionType, getattr(boxed.cls, "__annotate__", None)
-    ):
-        # Class has annotations, let's resolve generic arguments
-
-        closure_types = _get_closure_types(af)
-        args = tuple(
-            types.CellType(
-                boxed.cls.__dict__
-                if name == "__classdict__"
-                else boxed.str_args[name]
-                if name in boxed.str_args
-                else closure_types[name]
+def get_annotations(
+    obj: object,
+    args: dict[str, object],
+    key: str = '__annotate__',
+) -> Any | None:
+    """Get the annotations on an object, substituting in type vars."""
+
+    # Copy in any __type_params__ that aren't provided for, so that if
+    # we have to eval, we have them.
+    if params := getattr(obj, "__type_params__", None):
+        args = args.copy()
+        for param in params:
+            if str(param) not in args:
+                args[str(param)] = param
+
+    rr = None
+    if af := typing.cast(types.FunctionType, getattr(obj, key, None)):
+        # Substitute in names that are provided but keep the existing
+        # values for everything else.
+        closure = tuple(
+            types.CellType(args[name]) if name in args else orig_value
+            for name, orig_value in zip(
+                af.__code__.co_freevars, af.__closure__ or (), strict=True
             )
-            for name in af.__code__.co_freevars
         )
 
         ff = types.FunctionType(
-            af.__code__, af.__globals__, af.__name__, None, args
+            af.__code__, af.__globals__, af.__name__, None, closure
         )
         rr = ff(annotationlib.Format.VALUE)
 
-        if rr:
-            for k, v in rr.items():
-                if isinstance(v, str):
-                    # Handle cases where annotation is explicitly a string,
-                    # e.g.:
-                    #
-                    #   class Foo[X]:
-                    #       x: "Foo[X | None]"
+    if isinstance(rr, dict):
+        for k, v in rr.items():
+            if isinstance(v, str):
+                # Handle cases where annotation is explicitly a string,
+                # e.g.:
+                #
+                #   class Foo[X]:
+                #       x: "Foo[X | None]"
 
-                    annos[k] = eval(v, af.__globals__, boxed.str_args)
-                else:
-                    annos[k] = v
-    elif af := getattr(boxed.cls, "__annotations__", None):
+                rr[k] = eval(v, af.__globals__, args)
+
+    return rr
+
+
+def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
+    annos: dict[str, Any] = {}
+    dct: dict[str, Any] = {}
+
+    if (rr := get_annotations(boxed.cls, boxed.str_args)) is not None:
+        annos.update(rr)
+    elif anns := getattr(boxed.cls, "__annotations__", None):
         # TODO: substitute vars in this case
         _globals = {}
         if mod := sys.modules.get(boxed.cls.__module__):
@@ -252,7 +251,7 @@ def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
         _locals = dict(boxed.cls.__dict__)
         _locals.update(boxed.str_args)
 
-        for k, v in af.items():
+        for k, v in anns.items():
             if isinstance(v, str):
                 result = eval(v, _globals, _locals)
                 # Handle cases where annotation is explicitly a string
@@ -273,42 +272,13 @@ def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
         stuff = inspect.unwrap(orig)
 
         if isinstance(stuff, types.FunctionType):
-            local_fn: types.FunctionType | classmethod | staticmethod | None = (
-                None
-            )
-
-            if af := typing.cast(
-                types.FunctionType, getattr(stuff, "__annotate__", None)
-            ):
-                params = dict(
-                    zip(
-                        map(str, stuff.__type_params__),
-                        stuff.__type_params__,
-                        strict=True,
-                    )
-                )
-
-                closure_types = _get_closure_types(af)
-                args = tuple(
-                    types.CellType(
-                        boxed.cls.__dict__
-                        if name == "__classdict__"
-                        else params[name]
-                        if name in params
-                        else boxed.str_args[name]
-                        if name in boxed.str_args
-                        else closure_types[name]
-                    )
-                    for name in af.__code__.co_freevars
-                )
-
-                ff = types.FunctionType(
-                    af.__code__, af.__globals__, af.__name__, None, args
-                )
-                rr = ff(annotationlib.Format.VALUE)
+            local_fn: Any = None
 
+            if (rr := get_annotations(stuff, boxed.str_args)) is not None:
                 local_fn = make_func(orig, rr)
-            elif af := getattr(stuff, "__annotations__", None):
+            elif anns := getattr(stuff, "__annotations__", None):
+                # XXX: This is totally wrong; we still need to do
+                # substitute in class vars
                 local_fn = stuff
 
             if local_fn is not None:

typemap/type_eval/_eval_call.py

@@ -1,4 +1,3 @@
-import annotationlib
 import enum
 import inspect
 import types
@@ -12,7 +11,7 @@
 from . import _eval_typing
 from . import _typing_inspect
 from ._eval_operators import _callable_type_to_signature
-from ._apply_generic import substitute, _get_closure_types
+from ._apply_generic import substitute, get_annotations
 
 RtType = Any
 
@@ -188,30 +187,12 @@ def _eval_call_with_type_vars(
     vars: dict[str, RtType],
     ctx: _eval_typing.EvalContext,
 ) -> RtType:
-    try:
-        af = typing.cast(types.FunctionType, func.__annotate__)
-    except AttributeError:
-        raise ValueError("func has no __annotate__ attribute")
-    if not af:
-        raise ValueError("func has no __annotate__ attribute")
-
-    closure_types = _get_closure_types(af)
-    for name, value in closure_types.items():
-        if name not in vars:
-            vars[name] = value
-
-    af_args = tuple(
-        types.CellType(vars[name]) for name in af.__code__.co_freevars
-    )
-
-    ff = types.FunctionType(
-        af.__code__, af.__globals__, af.__name__, None, af_args
-    )
-
     old_obj = ctx.current_generic_alias
     ctx.current_generic_alias = func
     try:
-        rr = ff(annotationlib.Format.VALUE)
+        rr = get_annotations(func, vars)
+        if rr is None:
+            return Any
         return _eval_typing.eval_typing(rr["return"])
     finally:
         ctx.current_generic_alias = old_obj

typemap/type_eval/_eval_typing.py

@@ -1,5 +1,3 @@
-import annotationlib
-
 import collections.abc
 import contextlib
 import contextvars
@@ -368,13 +366,9 @@ def _eval_annotated(obj: typing_AnnotatedAlias, ctx: EvalContext):
 
 @_eval_types_impl.register
 def _eval_type_alias(obj: typing.TypeAliasType, ctx: EvalContext):
-    assert obj.__module__  # FIXME: or can this really happen?
-    func = obj.evaluate_value
-    mod = sys.modules[obj.__module__]
-    ff = types.FunctionType(
-        func.__code__, mod.__dict__, None, None, func.__closure__
+    unpacked = _apply_generic.get_annotations(
+        obj, args={}, key='evaluate_value'
     )
-    unpacked = ff(annotationlib.Format.VALUE)
     return _eval_types(unpacked, ctx)
 
 
@@ -417,32 +411,23 @@ def _eval_applied_type_alias(obj: types.GenericAlias, ctx: EvalContext):
         # Let's reconstruct it by evaluating all arguments
         return new_obj
 
-    func = obj.evaluate_value
-
-    # obj.__args__ matches the declared parameter order, but args are expected
-    # to be in the same order as func.__code__.co_freevars.
-    args_by_name = dict(
+    named_args = dict(
         zip(
             (p.__name__ for p in obj.__origin__.__type_params__),
             obj.__args__,
             strict=False,
         )
     )
 
-    args = tuple(
-        types.CellType(_eval_types(args_by_name[name], ctx))
-        for name in func.__code__.co_freevars
-    )
-    mod = sys.modules[obj.__module__]
-
     with _child_context() as child_ctx:
         child_ctx.current_generic_alias = new_obj
         if not _is_type_alias_type(new_obj):
             # Type alias types are already added in _eval_types
             child_ctx.alias_stack.add(new_obj)
 
-        ff = types.FunctionType(func.__code__, mod.__dict__, None, None, args)
-        unpacked = ff(annotationlib.Format.VALUE)
+        unpacked = _apply_generic.get_annotations(
+            obj, named_args, key='evaluate_value'
+        )
 
         child_ctx.seen[obj] = unpacked
         evaled = _eval_types(unpacked, child_ctx)