Skip to content

Commit

Permalink
pythongh-125614: annotationlib: Fix bug where not all Stringifiers ar…
Browse files Browse the repository at this point in the history
…e converted
  • Loading branch information
JelleZijlstra committed Oct 17, 2024
1 parent a5a7f5e commit 301d680
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 4 deletions.
28 changes: 24 additions & 4 deletions Lib/annotationlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ class Format(enum.IntEnum):
"__globals__",
"__owner__",
"__cell__",
"__stringifier_dict__",
)


Expand Down Expand Up @@ -268,7 +269,16 @@ class _Stringifier:
# instance of the other in place.
__slots__ = _SLOTS

def __init__(self, node, globals=None, owner=None, is_class=False, cell=None):
def __init__(
self,
node,
globals=None,
owner=None,
is_class=False,
cell=None,
*,
stringifier_dict,
):
# Either an AST node or a simple str (for the common case where a ForwardRef
# represent a single name).
assert isinstance(node, (ast.AST, str))
Expand All @@ -283,6 +293,7 @@ def __init__(self, node, globals=None, owner=None, is_class=False, cell=None):
self.__globals__ = globals
self.__cell__ = cell
self.__owner__ = owner
self.__stringifier_dict__ = stringifier_dict

def __convert_to_ast(self, other):
if isinstance(other, _Stringifier):
Expand Down Expand Up @@ -317,9 +328,15 @@ def __get_ast(self):
return node

def __make_new(self, node):
return _Stringifier(
node, self.__globals__, self.__owner__, self.__forward_is_class__
stringifier = _Stringifier(
node,
self.__globals__,
self.__owner__,
self.__forward_is_class__,
stringifier_dict=self.__stringifier_dict__,
)
self.__stringifier_dict__.stringifiers.append(stringifier)
return stringifier

# Must implement this since we set __eq__. We hash by identity so that
# stringifiers in dict keys are kept separate.
Expand Down Expand Up @@ -462,6 +479,7 @@ def __missing__(self, key):
globals=self.globals,
owner=self.owner,
is_class=self.is_class,
stringifier_dict=self,
)
self.stringifiers.append(fwdref)
return fwdref
Expand Down Expand Up @@ -516,7 +534,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
name = freevars[i]
else:
name = "__cell__"
fwdref = _Stringifier(name)
fwdref = _Stringifier(name, stringifier_dict=globals)
new_closure.append(types.CellType(fwdref))
closure = tuple(new_closure)
else:
Expand Down Expand Up @@ -573,6 +591,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
owner=owner,
globals=annotate.__globals__,
is_class=is_class,
stringifier_dict=globals,
)
globals.stringifiers.append(fwdref)
new_closure.append(types.CellType(fwdref))
Expand All @@ -591,6 +610,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
result = func(Format.VALUE)
for obj in globals.stringifiers:
obj.__class__ = ForwardRef
obj.__stringifier_dict__ = None # not needed for ForwardRef
if isinstance(obj.__ast_node__, str):
obj.__arg__ = obj.__ast_node__
obj.__ast_node__ = None
Expand Down
39 changes: 39 additions & 0 deletions Lib/test/test_annotationlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,35 @@ def f(x: int, y: doesntexist):
fwdref.evaluate()
self.assertEqual(fwdref.evaluate(globals={"doesntexist": 1}), 1)

def test_nonexistent_attribute(self):
def f(x: some.module, y: some[module], z: some(module), alpha: some | obj, beta: +some, gamma: some < obj):
pass

anno = annotationlib.get_annotations(f, format=Format.FORWARDREF)
x_anno = anno["x"]
self.assertIsInstance(x_anno, ForwardRef)
self.assertEqual(x_anno, ForwardRef("some.module"))

y_anno = anno["y"]
self.assertIsInstance(y_anno, ForwardRef)
self.assertEqual(y_anno, ForwardRef("some[module]"))

z_anno = anno["z"]
self.assertIsInstance(z_anno, ForwardRef)
self.assertEqual(z_anno, ForwardRef("some(module)"))

alpha_anno = anno["alpha"]
self.assertIsInstance(alpha_anno, ForwardRef)
self.assertEqual(alpha_anno, ForwardRef("some | obj"))

beta_anno = anno["beta"]
self.assertIsInstance(beta_anno, ForwardRef)
self.assertEqual(beta_anno, ForwardRef("+some"))

gamma_anno = anno["gamma"]
self.assertIsInstance(gamma_anno, ForwardRef)
self.assertEqual(gamma_anno, ForwardRef("some < obj"))


class TestSourceFormat(unittest.TestCase):
def test_closure(self):
Expand All @@ -91,6 +120,16 @@ def inner(arg: x):
anno = annotationlib.get_annotations(inner, format=Format.STRING)
self.assertEqual(anno, {"arg": "x"})

def test_closure_undefined(self):
if False:
x = 0

def inner(arg: x):
pass

anno = annotationlib.get_annotations(inner, format=Format.STRING)
self.assertEqual(anno, {"arg": "x"})

def test_function(self):
def f(x: int, y: doesntexist):
pass
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
In the :data:`~annotationlib.Format.FORWARDREF` format of
:mod:`annotationlib`, fix bug where nested expressions were not returned as
:class:`annotionlib.ForwardRef` format.

0 comments on commit 301d680

Please sign in to comment.