Add assignment expressions to correct ``locals`` for certain parents (#1213)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

3 files changed, 95 insertions, 2 deletions

diff --git a/ChangeLog b/ChangeLog
index f52f6846..5a81bebb 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -14,8 +14,13 @@ Release date: TBA
 
 * Fix the ``scope()`` and ``frame()`` methods of ``NamedExpr`` nodes.
   When these nodes occur in ``Arguments``, ``Keyword``  or ``Comprehension`` nodes these
-  methods now correctly point to the outer-scope of the `` FunctionDef``,
-  ``ClassDef`` or ``Comprehension``.
+  methods now correctly point to the outer-scope of the ``FunctionDef``,
+  ``ClassDef``, or ``Comprehension``.
+
+* Fix the ``set_local`` function for ``NamedExpr`` nodes.
+  When these nodes occur in ``Arguments``, ``Keyword``, or ``Comprehension`` nodes these
+  nodes are now correctly added to the locals of the ``FunctionDef``,
+  ``ClassDef``, or ``Comprehension``.
 
 
 What's New in astroid 2.8.3?
diff --git a/astroid/nodes/node_classes.py b/astroid/nodes/node_classes.py
index cb1d09c9..39151d25 100644
--- a/astroid/nodes/node_classes.py
+++ b/astroid/nodes/node_classes.py
@@ -4265,6 +4265,19 @@ class NamedExpr(mixins.AssignTypeMixin, NodeNG):
 
         return self.parent.scope()
 
+    def set_local(self, name: str, stmt: AssignName) -> None:
+        """Define that the given name is declared in the given statement node.
+        NamedExpr's in Arguments, Keyword or Comprehension are evaluated in their
+        parent's parent scope. So we add to their frame's locals.
+
+        .. seealso:: :meth:`scope`
+
+        :param name: The name that is being defined.
+
+        :param stmt: The statement that defines the given name.
+        """
+        self.frame().set_local(name, stmt)
+
 
 class Unknown(mixins.AssignTypeMixin, NodeNG):
     """This node represents a node in a constructed AST where
diff --git a/tests/unittest_nodes.py b/tests/unittest_nodes.py
index 9b1fee88..d855f510 100644
--- a/tests/unittest_nodes.py
+++ b/tests/unittest_nodes.py
@@ -1485,6 +1485,81 @@ def test_assignment_expression() -> None:
     assert second.as_string() == "b := test"
 
 
+@pytest.mark.skipif(not PY38_PLUS, reason="needs assignment expressions")
+def test_assignment_expression_in_functiondef() -> None:
+    code = """
+    def function(param = (assignment := "walrus")):
+        def inner_function(inner_param = (inner_assign := "walrus")):
+            pass
+        pass
+
+    class MyClass(attr = (assignment_two := "walrus")):
+        pass
+
+    VAR = lambda y = (assignment_three := "walrus"): print(y)
+
+    def func_with_lambda(
+        param=(named_expr_four := lambda y=(assignment_four := "walrus"): y),
+    ):
+        pass
+
+    COMPREHENSION = [y for i in (1, 2) if (assignment_five := i ** 2)]
+
+    def func():
+        var = lambda y = (assignment_six := 2): print(y)
+
+    VAR_TWO = [
+        func(assignment_seven := 2)
+        for _ in (1,)
+    ]
+
+    LAMBDA = lambda x: print(assignment_eight := x ** 2)
+
+    class SomeClass:
+        (assignment_nine := 2**2)
+    """
+    module = astroid.parse(code)
+
+    assert "assignment" in module.locals
+    assert isinstance(module.locals.get("assignment")[0], nodes.AssignName)
+    function = module.body[0]
+    assert "inner_assign" in function.locals
+    assert "inner_assign" not in module.locals
+    assert isinstance(function.locals.get("inner_assign")[0], nodes.AssignName)
+
+    assert "assignment_two" in module.locals
+    assert isinstance(module.locals.get("assignment_two")[0], nodes.AssignName)
+
+    assert "assignment_three" in module.locals
+    assert isinstance(module.locals.get("assignment_three")[0], nodes.AssignName)
+
+    assert "assignment_four" in module.locals
+    assert isinstance(module.locals.get("assignment_four")[0], nodes.AssignName)
+
+    assert "assignment_five" in module.locals
+    assert isinstance(module.locals.get("assignment_five")[0], nodes.AssignName)
+
+    func = module.body[5]
+    assert "assignment_six" in func.locals
+    assert "assignment_six" not in module.locals
+    assert isinstance(func.locals.get("assignment_six")[0], nodes.AssignName)
+
+    assert "assignment_seven" in module.locals
+    assert isinstance(module.locals.get("assignment_seven")[0], nodes.AssignName)
+
+    lambda_assign = module.body[7]
+    assert "assignment_eight" in lambda_assign.value.locals
+    assert "assignment_eight" not in module.locals
+    assert isinstance(
+        lambda_assign.value.locals.get("assignment_eight")[0], nodes.AssignName
+    )
+
+    class_assign = module.body[8]
+    assert "assignment_nine" in class_assign.locals
+    assert "assignment_nine" not in module.locals
+    assert isinstance(class_assign.locals.get("assignment_nine")[0], nodes.AssignName)
+
+
 def test_get_doc() -> None:
     node = astroid.extract_node(
         """