Skip to content

Commit b1027d4

Browse files
gh-138151: Fix annotationlib handling of multiple nonlocals (#138164)
1 parent d590685 commit b1027d4

File tree

3 files changed

+45
-13
lines changed

3 files changed

+45
-13
lines changed

Lib/annotationlib.py

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,9 @@ def __init__(
8585
# These are always set to None here but may be non-None if a ForwardRef
8686
# is created through __class__ assignment on a _Stringifier object.
8787
self.__globals__ = None
88+
# This may be either a cell object (for a ForwardRef referring to a single name)
89+
# or a dict mapping cell names to cell objects (for a ForwardRef containing references
90+
# to multiple names).
8891
self.__cell__ = None
8992
self.__extra_names__ = None
9093
# These are initially None but serve as a cache and may be set to a non-None
@@ -117,7 +120,7 @@ def evaluate(
117120
is_forwardref_format = True
118121
case _:
119122
raise NotImplementedError(format)
120-
if self.__cell__ is not None:
123+
if isinstance(self.__cell__, types.CellType):
121124
try:
122125
return self.__cell__.cell_contents
123126
except ValueError:
@@ -160,11 +163,18 @@ def evaluate(
160163

161164
# Type parameters exist in their own scope, which is logically
162165
# between the locals and the globals. We simulate this by adding
163-
# them to the globals.
164-
if type_params is not None:
166+
# them to the globals. Similar reasoning applies to nonlocals stored in cells.
167+
if type_params is not None or isinstance(self.__cell__, dict):
165168
globals = dict(globals)
169+
if type_params is not None:
166170
for param in type_params:
167171
globals[param.__name__] = param
172+
if isinstance(self.__cell__, dict):
173+
for cell_name, cell_value in self.__cell__.items():
174+
try:
175+
globals[cell_name] = cell_value.cell_contents
176+
except ValueError:
177+
pass
168178
if self.__extra_names__:
169179
locals = {**locals, **self.__extra_names__}
170180

@@ -202,7 +212,7 @@ def evaluate(
202212
except Exception:
203213
return self
204214
else:
205-
new_locals.transmogrify()
215+
new_locals.transmogrify(self.__cell__)
206216
return result
207217

208218
def _evaluate(self, globalns, localns, type_params=_sentinel, *, recursive_guard):
@@ -274,7 +284,7 @@ def __hash__(self):
274284
self.__forward_module__,
275285
id(self.__globals__), # dictionaries are not hashable, so hash by identity
276286
self.__forward_is_class__,
277-
self.__cell__,
287+
tuple(sorted(self.__cell__.items())) if isinstance(self.__cell__, dict) else self.__cell__,
278288
self.__owner__,
279289
tuple(sorted(self.__extra_names__.items())) if self.__extra_names__ else None,
280290
))
@@ -642,13 +652,15 @@ def __missing__(self, key):
642652
self.stringifiers.append(fwdref)
643653
return fwdref
644654

645-
def transmogrify(self):
655+
def transmogrify(self, cell_dict):
646656
for obj in self.stringifiers:
647657
obj.__class__ = ForwardRef
648658
obj.__stringifier_dict__ = None # not needed for ForwardRef
649659
if isinstance(obj.__ast_node__, str):
650660
obj.__arg__ = obj.__ast_node__
651661
obj.__ast_node__ = None
662+
if cell_dict is not None and obj.__cell__ is None:
663+
obj.__cell__ = cell_dict
652664

653665
def create_unique_name(self):
654666
name = f"__annotationlib_name_{self.next_id}__"
@@ -712,7 +724,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
712724

713725
globals = _StringifierDict({}, format=format)
714726
is_class = isinstance(owner, type)
715-
closure = _build_closure(
727+
closure, _ = _build_closure(
716728
annotate, owner, is_class, globals, allow_evaluation=False
717729
)
718730
func = types.FunctionType(
@@ -756,7 +768,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
756768
is_class=is_class,
757769
format=format,
758770
)
759-
closure = _build_closure(
771+
closure, cell_dict = _build_closure(
760772
annotate, owner, is_class, globals, allow_evaluation=True
761773
)
762774
func = types.FunctionType(
@@ -774,7 +786,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
774786
except Exception:
775787
pass
776788
else:
777-
globals.transmogrify()
789+
globals.transmogrify(cell_dict)
778790
return result
779791

780792
# Try again, but do not provide any globals. This allows us to return
@@ -786,7 +798,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
786798
is_class=is_class,
787799
format=format,
788800
)
789-
closure = _build_closure(
801+
closure, cell_dict = _build_closure(
790802
annotate, owner, is_class, globals, allow_evaluation=False
791803
)
792804
func = types.FunctionType(
@@ -797,7 +809,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
797809
kwdefaults=annotate.__kwdefaults__,
798810
)
799811
result = func(Format.VALUE_WITH_FAKE_GLOBALS)
800-
globals.transmogrify()
812+
globals.transmogrify(cell_dict)
801813
if _is_evaluate:
802814
if isinstance(result, ForwardRef):
803815
return result.evaluate(format=Format.FORWARDREF)
@@ -822,14 +834,16 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
822834

823835
def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluation):
824836
if not annotate.__closure__:
825-
return None
837+
return None, None
826838
freevars = annotate.__code__.co_freevars
827839
new_closure = []
840+
cell_dict = {}
828841
for i, cell in enumerate(annotate.__closure__):
829842
if i < len(freevars):
830843
name = freevars[i]
831844
else:
832845
name = "__cell__"
846+
cell_dict[name] = cell
833847
new_cell = None
834848
if allow_evaluation:
835849
try:
@@ -850,7 +864,7 @@ def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluat
850864
stringifier_dict.stringifiers.append(fwdref)
851865
new_cell = types.CellType(fwdref)
852866
new_closure.append(new_cell)
853-
return tuple(new_closure)
867+
return tuple(new_closure), cell_dict
854868

855869

856870
def _stringify_single(anno):

Lib/test/test_annotationlib.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1194,6 +1194,21 @@ class RaisesAttributeError:
11941194
},
11951195
)
11961196

1197+
def test_nonlocal_in_annotation_scope(self):
1198+
class Demo:
1199+
nonlocal sequence_b
1200+
x: sequence_b
1201+
y: sequence_b[int]
1202+
1203+
fwdrefs = get_annotations(Demo, format=Format.FORWARDREF)
1204+
1205+
self.assertIsInstance(fwdrefs["x"], ForwardRef)
1206+
self.assertIsInstance(fwdrefs["y"], ForwardRef)
1207+
1208+
sequence_b = list
1209+
self.assertIs(fwdrefs["x"].evaluate(), list)
1210+
self.assertEqual(fwdrefs["y"].evaluate(), list[int])
1211+
11971212
def test_raises_error_from_value(self):
11981213
# test that if VALUE is the only supported format, but raises an error
11991214
# that error is propagated from get_annotations
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
In :mod:`annotationlib`, improve evaluation of forward references to
2+
nonlocal variables that are not yet defined when the annotations are
3+
initially evaluated.

0 commit comments

Comments
 (0)