diff --git a/Lib/annotationlib.py b/Lib/annotationlib.py index bee019cd51591e..13590c62bee2d3 100644 --- a/Lib/annotationlib.py +++ b/Lib/annotationlib.py @@ -664,6 +664,17 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False): # possibly constants if the annotate function uses them directly). We then # convert each of those into a string to get an approximation of the # original source. + + # Attempt to call with VALUE_WITH_FAKE_GLOBALS to check if it is implemented + # See: https://github.com/python/cpython/issues/138764 + # Only fail on NotImplementedError + try: + annotate(Format.VALUE_WITH_FAKE_GLOBALS) + except NotImplementedError: + raise + except Exception: + pass + globals = _StringifierDict({}, format=format) is_class = isinstance(owner, type) closure = _build_closure( @@ -722,6 +733,10 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False): ) try: result = func(Format.VALUE_WITH_FAKE_GLOBALS) + except NotImplementedError: + # If NotImplementedError is raised, don't try to call again with + # no globals. + raise except Exception: pass else: diff --git a/Lib/test/test_annotationlib.py b/Lib/test/test_annotationlib.py index 88e0d611647f28..5a2d8a558a2891 100644 --- a/Lib/test/test_annotationlib.py +++ b/Lib/test/test_annotationlib.py @@ -8,8 +8,10 @@ import itertools import pickle from string.templatelib import Template +import types import typing import unittest +import unittest.mock from annotationlib import ( Format, ForwardRef, @@ -1206,6 +1208,84 @@ def evaluate(format, exc=NotImplementedError): ) +class TestCallAnnotateFunction(unittest.TestCase): + def _annotate_mock(self): + def annotate(format, /): + if format == Format.VALUE: + return {"x": str} + else: + raise NotImplementedError(format) + + annotate_mock = unittest.mock.MagicMock( + wraps=annotate + ) + + # Add missing magic attributes needed + required_magic = [ + "__builtins__", + "__closure__", + "__code__", + "__defaults__", + "__globals__", + "__kwdefaults__", + ] + + for attrib in required_magic: + setattr(annotate_mock, attrib, getattr(annotate, attrib)) + + return annotate_mock + + def test_user_annotate_value(self): + annotate = self._annotate_mock() + + annotations = annotationlib.call_annotate_function( + annotate, + Format.VALUE, + ) + + self.assertEqual(annotations, {"x": str}) + annotate.assert_called_once_with(Format.VALUE) + + def test_user_annotate_forwardref(self): + annotate = self._annotate_mock() + + new_annotate = None + functype = types.FunctionType + + def functiontype_mock(*args, **kwargs): + nonlocal new_annotate + new_func = unittest.mock.MagicMock(wraps=functype(*args, **kwargs)) + new_annotate = new_func + return new_func + + with unittest.mock.patch("types.FunctionType", new=functiontype_mock): + with self.assertRaises(NotImplementedError): + annotations = annotationlib.call_annotate_function( + annotate, + Format.FORWARDREF, + ) + + # Test the direct call + annotate.assert_called_once_with(Format.FORWARDREF) + + # Test the call on the function with fake globals + new_annotate.assert_called_once_with(Format.VALUE_WITH_FAKE_GLOBALS) + + def test_user_annotate_string(self): + annotate = self._annotate_mock() + + with self.assertRaises(NotImplementedError): + annotations = annotationlib.call_annotate_function( + annotate, + Format.STRING, + ) + + annotate.assert_has_calls([ + unittest.mock.call(Format.STRING), + unittest.mock.call(Format.VALUE_WITH_FAKE_GLOBALS), + ]) + + class MetaclassTests(unittest.TestCase): def test_annotated_meta(self): class Meta(type):