Merge all of the code that calls __annotate__/evaluate_value (#86)

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

Author: msullivan

tests/test_type_eval.py

@@ -422,7 +422,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

@@ -189,84 +189,67 @@ 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__',
+    annos_ok: bool = True,
+) -> Any | None:
+    """Get the annotations on an object, substituting in type vars."""
+
+    rr = None
+    globs = 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
-        )
+        globs = af.__globals__
+        ff = types.FunctionType(af.__code__, globs, af.__name__, None, closure)
         rr = ff(annotationlib.Format.VALUE)
-
-        if rr:
-            for k, v in rr.items():
+    elif annos_ok and (rr := getattr(obj, "__annotations__", None)):
+        globs = {}
+        if mod := sys.modules.get(obj.__module__):
+            globs.update(vars(mod))
+
+    if isinstance(rr, dict) and any(isinstance(v, str) for v in rr.values()):
+        # 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
+
+        for k, v in rr.items():
+            # Eval strings
+            if isinstance(v, str):
+                v = eval(v, globs, args)
+                # Handle cases where annotation is explicitly a string,
+                # e.g.:
+                #   class Foo[X]:
+                #       x: "Foo[X | None]"
                 if isinstance(v, str):
-                    # Handle cases where annotation is explicitly a string,
-                    # e.g.:
-                    #
-                    #   class Foo[X]:
-                    #       x: "Foo[X | None]"
+                    v = eval(v, globs, args)
+            rr[k] = v
 
-                    annos[k] = eval(v, af.__globals__, boxed.str_args)
-                else:
-                    annos[k] = v
-    elif af := getattr(boxed.cls, "__annotations__", None):
-        # TODO: substitute vars in this case
-        _globals = {}
-        if mod := sys.modules.get(boxed.cls.__module__):
-            _globals.update(vars(mod))
-        _globals.update(boxed.str_args)
-
-        _locals = dict(boxed.cls.__dict__)
-        _locals.update(boxed.str_args)
-
-        for k, v in af.items():
-            if isinstance(v, str):
-                result = eval(v, _globals, _locals)
-                # Handle cases where annotation is explicitly a string
-                # e.g.
-                #   class Foo[T]:
-                #       x: "Bar[T]"
-                if isinstance(result, str):
-                    result = eval(result, _globals, _locals)
-                annos[k] = result
+    return rr
 
-            else:
-                annos[k] = v
+
+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)
 
     for name, orig in boxed.cls.__dict__.items():
         if name in EXCLUDED_ATTRIBUTES:
@@ -275,42 +258,18 @@ 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
+
+            # TODO: This annos_ok thing is a hack because processing
+            # __annotations__ on methods broke stuff and I didn't want
+            # to chase it down yet.
+            if (
+                rr := get_annotations(stuff, boxed.str_args, annos_ok=False)
+            ) is not None:
                 local_fn = make_func(orig, rr)
-            elif af := getattr(stuff, "__annotations__", None):
+            elif 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
 
@@ -39,7 +38,7 @@ def _get_bound_type_args(
     arg_types: tuple[RtType, ...],
     kwarg_types: dict[str, RtType],
 ) -> dict[str, RtType]:
-    sig = _eval_operators._resolved_function_signature(func)
+    sig = inspect.signature(func)
 
     bound = sig.bind(*arg_types, **kwarg_types)
 
@@ -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_operators.py

@@ -7,7 +7,6 @@
 import re
 import types
 import typing
-import sys
 
 from typing_extensions import _AnnotatedAlias as typing_AnnotatedAlias
 
@@ -637,7 +636,7 @@ def _callable_type_to_method(name, typ, ctx):
 
 def _function_type(func, *, receiver_type):
     root = inspect.unwrap(func)
-    sig = _resolved_function_signature(root, receiver_type)
+    sig = inspect.signature(root)
     # XXX: __type_params__!!!
 
     empty = inspect.Parameter.empty
@@ -730,94 +729,6 @@ def _create_generic_callable_lambda(
         ]
 
 
-def _resolved_function_signature(func, receiver_type=None):
-    """Get the signature of a function with type hints resolved.
-
-    This is used to deal with string annotations in the signature which are
-    generated when using __future__ import annotations.
-    """
-
-    sig = inspect.signature(func)
-
-    _globals, _locals = _get_function_hint_namespaces(func, receiver_type)
-    if hints := typing.get_type_hints(
-        func, globalns=_globals, localns=_locals, include_extras=True
-    ):
-        params = []
-        for name, param in sig.parameters.items():
-            annotation = hints.get(name, param.annotation)
-            params.append(param.replace(annotation=annotation))
-
-        return_annotation = hints.get("return", sig.return_annotation)
-        sig = sig.replace(
-            parameters=params, return_annotation=return_annotation
-        )
-
-    return sig
-
-
-def _get_class_type_hint_namespaces(
-    obj: type,
-) -> tuple[dict[str, typing.Any], dict[str, typing.Any]]:
-    globalns: dict[str, typing.Any] = {}
-    localns: dict[str, typing.Any] = {}
-
-    # Get module globals
-    if obj.__module__ and (module := sys.modules.get(obj.__module__)):
-        globalns.update(module.__dict__)
-
-    # Annotations may use typevars defined in the class
-    localns.update(obj.__dict__)
-
-    if _typing_inspect.is_generic_alias(obj):
-        # We need the origin's type vars
-        localns.update(obj.__origin__.__dict__)
-
-        # Extract type parameters from the class
-        args = typing.get_args(obj)
-        origin = typing.get_origin(obj)
-        tps = getattr(obj, '__type_params__', ()) or getattr(
-            origin, '__parameters__', ()
-        )
-        for tp, arg in zip(tps, args, strict=False):
-            localns[tp.__name__] = arg
-
-    # Add the class itself for self-references
-    localns[obj.__name__] = obj
-
-    return globalns, localns
-
-
-def _get_function_hint_namespaces(func, receiver_type=None):
-    globalns = {}
-    localns = {}
-
-    # module globals
-    module = inspect.getmodule(func)
-    if module:
-        globalns |= module.__dict__
-
-    # If no receiver was specified, this might still be a method, try to get
-    # the class from the qualname.
-    if (
-        not receiver_type
-        and (qn := getattr(func, '__qualname__', None))
-        and '.' in qn
-    ):
-        class_name = qn.rsplit('.', 1)[0]
-        receiver_type = getattr(module, class_name, None)
-
-    # Get the class's type hint namespaces
-    if receiver_type:
-        cls_globalns, cls_localns = _get_class_type_hint_namespaces(
-            receiver_type
-        )
-        globalns.update(cls_globalns)
-        localns.update(cls_localns)
-
-    return globalns, localns
-
-
 def _hint_to_member(n, t, qs, init, d, *, ctx):
     return Member[
         typing.Literal[n],