nonlocal in locals

astroid/rebuilder.py

@@ -61,6 +61,11 @@
         self._manager = manager
         self._data = data.split("\n") if data else None
         self._global_names: list[dict[str, list[nodes.Global]]] = []
+        # In _nonlocal_names,
+        # what we save is the function where the variable is created,
+        # rather than nodes.Nonlocal.
+        # We don't really need the Nonlocal statement.
+        self._nonlocal_names: list[dict[str, nodes.FunctionDef]] = []
         self._import_from_nodes: list[nodes.ImportFrom] = []
         self._delayed_assattr: list[nodes.AssignAttr] = []
         self._visit_meths: dict[type[ast.AST], Callable[[ast.AST, NodeNG], NodeNG]] = {}
@@ -451,6 +456,9 @@
         """Save assignment situation since node.parent is not available yet."""
         if self._global_names and node.name in self._global_names[-1]:
             node.root().set_local(node.name, node)
+        elif self._nonlocal_names and node.name in self._nonlocal_names[-1]:
+            function_def = self._nonlocal_names[-1][node.name]
+            function_def.set_local(node.name, node)
         else:
             assert node.parent
             assert node.name
@@ -1065,6 +1073,7 @@
     ) -> _FunctionT:
         """Visit an FunctionDef node to become astroid."""
         self._global_names.append({})
+        self._nonlocal_names.append({})
         node, doc_ast_node = self._get_doc(node)
 
         lineno = node.lineno
@@ -1113,6 +1122,7 @@
             ),
         )
         self._global_names.pop()
+        self._nonlocal_names.pop()
         parent.set_local(newnode.name, newnode)
         return newnode
 
@@ -1383,14 +1393,32 @@
 
     def visit_nonlocal(self, node: ast.Nonlocal, parent: NodeNG) -> nodes.Nonlocal:
         """Visit a Nonlocal node and return a new instance of it."""
-        return nodes.Nonlocal(
+        newnode = nodes.Nonlocal(
             names=node.names,
             lineno=node.lineno,
             col_offset=node.col_offset,
             end_lineno=node.end_lineno,
             end_col_offset=node.end_col_offset,
             parent=parent,
         )
+        names = set(newnode.names)
+        # Go through the tree and find where those names are created
+        scope = newnode
+        while len(names) != 0:
+            scope = scope.parent
+            if not scope:
+                # It's not inside a nested function or there are no variables with that name.
+                # Just ignore it as visit_global does when global is used in module scope.
+                break
+            if isinstance(scope, nodes.FunctionDef):
+                found = []
+                for name in names:
+                    if name in scope.locals:
+                        found.append(name)
+                        self._nonlocal_names[-1][name] = scope
+                for name in found:
+                    names.remove(name)
+        return newnode
 
     def visit_constant(self, node: ast.Constant, parent: NodeNG) -> nodes.Const:
         """Visit a Constant node by returning a fresh instance of Const."""

tests/test_builder.py

@@ -716,6 +716,39 @@ def test_type_comments_without_content(self) -> None:
         )
         assert node
 
+    def test_locals_with_global_and_nonlocal(self) -> None:
+        module = builder.parse(
+            """
+            x1 = 1                   # Line 2
+            def f1():                # Line 3
+                x2 = 2               # Line 4
+                def f2():            # Line 5
+                    global x1        # Line 6
+                    nonlocal x2      # Line 7
+                    x1 = 1           # Line 8
+                    x2 = 2           # Line 9
+                    x3 = 3           # Line 10
+            """
+        )
+        self.assertSetEqual(set(module.locals), {"x1", "f1"})
+        x1 = module.locals["x1"]
+        f1 = module.locals["f1"][0]
+        self.assertEqual(len(x1), 2)
+        self.assertEqual(x1[0].lineno, 2)
+        self.assertEqual(x1[1].lineno, 8)
+
+        self.assertSetEqual(set(f1.locals), {"x2", "f2"})
+        x2 = f1.locals["x2"]
+        f2 = f1.locals["f2"][0]
+        self.assertEqual(len(x2), 2)
+        self.assertEqual(x2[0].lineno, 4)
+        self.assertEqual(x2[1].lineno, 9)
+
+        self.assertSetEqual(set(f2.locals), {"x3"})
+        x3 = f2.locals["x3"]
+        self.assertEqual(len(x3), 1)
+        self.assertEqual(x3[0].lineno, 10)
+
 
 class FileBuildTest(unittest.TestCase):
     def setUp(self) -> None: