Skip to content

Commit 1c86ad9

Browse files
committed
gh-138151: Fix annotationlib handling of multiple nonlocals
1 parent fd8f42d commit 1c86ad9

File tree

3 files changed

+44
-13
lines changed

3 files changed

+44
-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

@@ -199,7 +209,7 @@ def evaluate(
199209
except Exception:
200210
return self
201211
else:
202-
new_locals.transmogrify()
212+
new_locals.transmogrify(self.__cell__)
203213
return result
204214

205215
def _evaluate(self, globalns, localns, type_params=_sentinel, *, recursive_guard):
@@ -278,7 +288,7 @@ def __hash__(self):
278288
self.__forward_module__,
279289
id(self.__globals__), # dictionaries are not hashable, so hash by identity
280290
self.__forward_is_class__,
281-
self.__cell__,
291+
tuple(sorted(self.__cell__.items())) if isinstance(self.__cell__, dict) else self.__cell__,
282292
self.__owner__,
283293
tuple(sorted(self.__extra_names__.items())) if self.__extra_names__ else None,
284294
))
@@ -608,13 +618,15 @@ def __missing__(self, key):
608618
self.stringifiers.append(fwdref)
609619
return fwdref
610620

611-
def transmogrify(self):
621+
def transmogrify(self, cell_dict):
612622
for obj in self.stringifiers:
613623
obj.__class__ = ForwardRef
614624
obj.__stringifier_dict__ = None # not needed for ForwardRef
615625
if isinstance(obj.__ast_node__, str):
616626
obj.__arg__ = obj.__ast_node__
617627
obj.__ast_node__ = None
628+
if cell_dict is not None and obj.__cell__ is None:
629+
obj.__cell__ = cell_dict
618630

619631
def create_unique_name(self):
620632
name = f"__annotationlib_name_{self.next_id}__"
@@ -666,7 +678,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
666678
# original source.
667679
globals = _StringifierDict({}, format=format)
668680
is_class = isinstance(owner, type)
669-
closure = _build_closure(
681+
closure, _ = _build_closure(
670682
annotate, owner, is_class, globals, allow_evaluation=False
671683
)
672684
func = types.FunctionType(
@@ -710,7 +722,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
710722
is_class=is_class,
711723
format=format,
712724
)
713-
closure = _build_closure(
725+
closure, cell_dict = _build_closure(
714726
annotate, owner, is_class, globals, allow_evaluation=True
715727
)
716728
func = types.FunctionType(
@@ -725,7 +737,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
725737
except Exception:
726738
pass
727739
else:
728-
globals.transmogrify()
740+
globals.transmogrify(cell_dict)
729741
return result
730742

731743
# Try again, but do not provide any globals. This allows us to return
@@ -737,7 +749,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
737749
is_class=is_class,
738750
format=format,
739751
)
740-
closure = _build_closure(
752+
closure, cell_dict = _build_closure(
741753
annotate, owner, is_class, globals, allow_evaluation=False
742754
)
743755
func = types.FunctionType(
@@ -748,7 +760,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
748760
kwdefaults=annotate.__kwdefaults__,
749761
)
750762
result = func(Format.VALUE_WITH_FAKE_GLOBALS)
751-
globals.transmogrify()
763+
globals.transmogrify(cell_dict)
752764
if _is_evaluate:
753765
if isinstance(result, ForwardRef):
754766
return result.evaluate(format=Format.FORWARDREF)
@@ -773,14 +785,16 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
773785

774786
def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluation):
775787
if not annotate.__closure__:
776-
return None
788+
return None, None
777789
freevars = annotate.__code__.co_freevars
778790
new_closure = []
791+
cell_dict = {}
779792
for i, cell in enumerate(annotate.__closure__):
780793
if i < len(freevars):
781794
name = freevars[i]
782795
else:
783796
name = "__cell__"
797+
cell_dict[name] = cell
784798
new_cell = None
785799
if allow_evaluation:
786800
try:
@@ -801,7 +815,7 @@ def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluat
801815
stringifier_dict.stringifiers.append(fwdref)
802816
new_cell = types.CellType(fwdref)
803817
new_closure.append(new_cell)
804-
return tuple(new_closure)
818+
return tuple(new_closure), cell_dict
805819

806820

807821
def _stringify_single(anno):

Lib/test/test_annotationlib.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1186,6 +1186,20 @@ class RaisesAttributeError:
11861186
},
11871187
)
11881188

1189+
def test_nonlocal_in_annotation_scope(self):
1190+
class Demo:
1191+
nonlocal sequence_b
1192+
x: sequence_b
1193+
y: sequence_b[int]
1194+
1195+
fwdrefs = get_annotations(Demo, format=Format.FORWARDREF)
1196+
1197+
self.assertIsInstance(fwdrefs["x"], ForwardRef)
1198+
self.assertIsInstance(fwdrefs["y"], ForwardRef)
1199+
1200+
sequence_b = list
1201+
self.assertIs(fwdrefs["x"].evaluate(), list)
1202+
self.assertEqual(fwdrefs["y"].evaluate(), list[int])
11891203

11901204
class TestCallEvaluateFunction(unittest.TestCase):
11911205
def test_evaluation(self):
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)