gh-141388: Fully support non-function callables as annotate functions

Lib/annotationlib.py

@@ -728,13 +728,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
             annotate, owner, is_class, globals, allow_evaluation=False
         )
         func = types.FunctionType(
-            annotate.__code__,
+            _get_annotate_attr(annotate, "__code__", None),
             globals,
             closure=closure,
-            argdefs=annotate.__defaults__,
-            kwdefaults=annotate.__kwdefaults__,
+            argdefs=_get_annotate_attr(annotate, "__defaults__", None),
+            kwdefaults=_get_annotate_attr(annotate, "__kwdefaults__", None),
         )
-        annos = func(Format.VALUE_WITH_FAKE_GLOBALS)
+        annos = _direct_call_annotate(func, annotate, Format.VALUE_WITH_FAKE_GLOBALS)
         if _is_evaluate:
             return _stringify_single(annos)
         return {
@@ -759,11 +759,18 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
         # reconstruct the source. But in the dictionary that we eventually return, we
         # want to return objects with more user-friendly behavior, such as an __eq__
         # that returns a bool and an defined set of attributes.
-        namespace = {**annotate.__builtins__, **annotate.__globals__}
+        annotate_globals = _get_annotate_attr(annotate, "__globals__", {})
+        annotate_code = _get_annotate_attr(annotate, "__code__", None)
+        annotate_defaults = _get_annotate_attr(annotate, "__defaults__", None)
+        annotate_kwdefaults = _get_annotate_attr(annotate, "__kwdefaults__", None)
+        namespace = {
+            **_get_annotate_attr(annotate, "__builtins__", {}),
+            **annotate_globals
+        }
         is_class = isinstance(owner, type)
         globals = _StringifierDict(
             namespace,
-            globals=annotate.__globals__,
+            globals=annotate_globals,
             owner=owner,
             is_class=is_class,
             format=format,
@@ -772,14 +779,14 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
             annotate, owner, is_class, globals, allow_evaluation=True
         )
         func = types.FunctionType(
-            annotate.__code__,
+            annotate_code,
             globals,
             closure=closure,
-            argdefs=annotate.__defaults__,
-            kwdefaults=annotate.__kwdefaults__,
+            argdefs=annotate_defaults,
+            kwdefaults=annotate_kwdefaults,
         )
         try:
-            result = func(Format.VALUE_WITH_FAKE_GLOBALS)
+            result = _direct_call_annotate(func, annotate, Format.VALUE_WITH_FAKE_GLOBALS)
         except NotImplementedError:
             # FORWARDREF and VALUE_WITH_FAKE_GLOBALS not supported, fall back to VALUE
             return annotate(Format.VALUE)
@@ -793,7 +800,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
         # a value in certain cases where an exception gets raised during evaluation.
         globals = _StringifierDict(
             {},
-            globals=annotate.__globals__,
+            globals=annotate_globals,
             owner=owner,
             is_class=is_class,
             format=format,
@@ -802,13 +809,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
             annotate, owner, is_class, globals, allow_evaluation=False
         )
         func = types.FunctionType(
-            annotate.__code__,
+            annotate_code,
             globals,
             closure=closure,
-            argdefs=annotate.__defaults__,
-            kwdefaults=annotate.__kwdefaults__,
+            argdefs=annotate_defaults,
+            kwdefaults=annotate_kwdefaults,
         )
-        result = func(Format.VALUE_WITH_FAKE_GLOBALS)
+        result = _direct_call_annotate(func, annotate, Format.VALUE_WITH_FAKE_GLOBALS)
         globals.transmogrify(cell_dict)
         if _is_evaluate:
             if isinstance(result, ForwardRef):
@@ -833,12 +840,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
 
 
 def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluation):
-    if not annotate.__closure__:
+    closure = _get_annotate_attr(annotate, "__closure__", None)
+    if not closure:
         return None, None
-    freevars = annotate.__code__.co_freevars
+    freevars = _get_annotate_attr(annotate, "__code__", None).co_freevars
     new_closure = []
     cell_dict = {}
-    for i, cell in enumerate(annotate.__closure__):
+    for i, cell in enumerate(closure):
         if i < len(freevars):
             name = freevars[i]
         else:
@@ -857,7 +865,7 @@ def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluat
                 name,
                 cell=cell,
                 owner=owner,
-                globals=annotate.__globals__,
+                globals=_get_annotate_attr(annotate, "__globals__", {}),
                 is_class=is_class,
                 stringifier_dict=stringifier_dict,
             )
@@ -879,6 +887,48 @@ def _stringify_single(anno):
         return repr(anno)
 
 
+def _get_annotate_attr(annotate, attr, default):
+    if (value := getattr(annotate, attr, None)) is not None:
+        return value
+
+    if isinstance(annotate.__call__, types.MethodType):
+        if call_func := getattr(annotate.__call__, "__func__", None):
+            return getattr(call_func, attr, default)
+    elif isinstance(annotate, type):
+        return getattr(annotate.__init__, attr, default)
+    elif (
+        (functools := sys.modules.get("functools", None))
+        and isinstance(annotate, functools.partial)
+    ):
+        return getattr(annotate.func, attr, default)
+
+    return default
+
+def _direct_call_annotate(func, annotate, format):
+    # If annotate is a method, we need to pass its self as the first param
+    if (
+        hasattr(annotate.__call__, "__func__") and
+        (self := getattr(annotate.__call__, "__self__", None))
+    ):
+        return func(self, format)
+
+    # If annotate is a class, `func` is the __init__ method, so we still need to call
+    # __new__() to create the instance
+    if isinstance(annotate, type):
+        inst = annotate.__new__(annotate)
+        func(inst, format)
+        return inst
+
+    # If annotate is a partial function, re-create it with the new function object.
+    # We could call the function directly, but then we'd have to handle placeholders,
+    # and this way should be more robust for future changes.
+    if functools := sys.modules.get("functools", None):
+        if isinstance(annotate, functools.partial):
+            return functools.partial(func, *annotate.args, **annotate.keywords)(format)
+
+    return func(format)
+
+
 def get_annotate_from_class_namespace(obj):
     """Retrieve the annotate function from a class namespace dictionary.
 

Lib/test/test_annotationlib.py

@@ -1132,6 +1132,26 @@ def __annotate__(self):
             {"x": "int"},
         )
 
+    def test_non_function_annotate(self):
+        class AnnotateCallable:
+            def __call__(self, format, /):
+                if format > 2:
+                    raise NotImplementedError
+                return {"x": int}
+
+        class OnlyAnnotate:
+            @property
+            def __annotate__(self):
+                return AnnotateCallable()
+
+        oa = OnlyAnnotate()
+        self.assertEqual(get_annotations(oa, format=Format.VALUE), {"x": int})
+        self.assertEqual(get_annotations(oa, format=Format.FORWARDREF), {"x": int})
+        self.assertEqual(
+            get_annotations(oa, format=Format.STRING),
+            {"x": "int"},
+        )
+
     def test_non_dict_annotate(self):
         class WeirdAnnotate:
             def __annotate__(self, *args, **kwargs):
@@ -1561,6 +1581,92 @@ def annotate(format, /, __Format=Format, __NotImplementedError=NotImplementedErr
 
         self.assertEqual(annotations, {"x": "str"})
 
+    def test_callable_object_annotate(self):
+        # If Format.STRING and Format.VALUE_WITH_FAKE_GLOBALS are not
+        # supported fall back to Format.VALUE and convert to strings
+        class Annotate:
+            def __call__(self, format, /):
+                return {"x": str}
+
+        for fmt in [Format.VALUE, Format.FORWARDREF, Format.STRING]:
+            self.assertEqual(
+                annotationlib.call_annotate_function(Annotate(), format=fmt),
+                {"x": str}
+            )
+
+    def test_callable_object_annotate_forwardref_value_fallback(self):
+        # If Format.STRING and Format.VALUE_WITH_FAKE_GLOBALS are not
+        # supported fall back to Format.VALUE and convert to strings
+        class Annotate:
+            def __call__(self, format, /, __Format=Format,
+                         __NotImplementedError=NotImplementedError):
+                if format == __Format.VALUE:
+                    return {"x": str}
+                else:
+                    raise __NotImplementedError(format)
+
+        annotations = annotationlib.call_annotate_function(
+            Annotate(),
+            Format.FORWARDREF,
+        )
+
+        self.assertEqual(annotations, {"x": str})
+
+    def test_callable_class_annotate_forwardref_value_fallback(self):
+        # If Format.STRING and Format.VALUE_WITH_FAKE_GLOBALS are not
+        # supported fall back to Format.VALUE and convert to strings
+        class Annotate(dict):
+            def __init__(self, format, /, __Format=Format,
+                         __NotImplementedError=NotImplementedError):
+                if format == __Format.VALUE:
+                    super().__init__({"x": int})
+                else:
+                    raise __NotImplementedError(format)
+
+        annotations = annotationlib.call_annotate_function(
+            Annotate,
+            Format.FORWARDREF,
+        )
+
+        self.assertEqual(annotations, {"x": int})
+
+    def test_callable_partial_annotate_forwardref_value_fallback(self):
+        # If Format.STRING and Format.VALUE_WITH_FAKE_GLOBALS are not
+        # supported fall back to Format.VALUE and convert to strings
+        def format(format, second, /, *, third, __Format=Format,
+                     __NotImplementedError=NotImplementedError):
+            if format == __Format.VALUE:
+                return {"x": format * second * third}
+            else:
+                raise __NotImplementedError(format)
+
+        annotations = annotationlib.call_annotate_function(
+            functools.partial(format, functools.Placeholder, 5, third=6),
+            Format.FORWARDREF,
+        )
+
+        self.assertEqual(annotations, {"x": Format.VALUE * 5 * 6})
+
+    def test_callable_object_annotate_string_fakeglobals(self):
+        # If Format.STRING is not supported but Format.VALUE_WITH_FAKE_GLOBALS is
+        # prefer that over Format.VALUE
+        class Annotate:
+            def __call__(self, format, /, __Format=Format,
+                         __NotImplementedError=NotImplementedError):
+                if format == __Format.VALUE:
+                    return {'x': str}
+                elif format == __Format.VALUE_WITH_FAKE_GLOBALS:
+                    return {'x': int}
+                else:
+                    raise __NotImplementedError(format)
+
+        annotations = annotationlib.call_annotate_function(
+            Annotate(),
+            Format.STRING,
+        )
+
+        self.assertEqual(annotations, {"x": "int"})
+
     def test_condition_not_stringified(self):
         # Make sure the first condition isn't evaluated as True by being converted
         # to a _Stringifier
@@ -1606,6 +1712,23 @@ def annotate(format, /):
             with self.assertRaises(DemoException):
                 annotationlib.call_annotate_function(annotate, format=fmt)
 
+    def test_callable_object_error_from_value_raised(self):
+        # Test that the error from format.VALUE is raised
+        # if all formats fail
+
+        class DemoException(Exception): ...
+
+        class Annotate:
+            def __call__(self, format, /):
+                if format == Format.VALUE:
+                    raise DemoException()
+                else:
+                    raise NotImplementedError(format)
+
+        for fmt in [Format.VALUE, Format.FORWARDREF, Format.STRING]:
+            with self.assertRaises(DemoException):
+                annotationlib.call_annotate_function(Annotate(), format=fmt)
+
 
 class MetaclassTests(unittest.TestCase):
     def test_annotated_meta(self):