Add assignment expressions to correct ``locals`` for certain parents (#1213)
Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
diff --git a/ChangeLog b/ChangeLog
index f52f684..5a81beb 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -14,8 +14,13 @@
* 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 cb1d09c..39151d2 100644
--- a/astroid/nodes/node_classes.py
+++ b/astroid/nodes/node_classes.py
@@ -4265,6 +4265,19 @@
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 9b1fee8..d855f51 100644
--- a/tests/unittest_nodes.py
+++ b/tests/unittest_nodes.py
@@ -1485,6 +1485,81 @@
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(
"""
