Change ``frame`` and ``scope`` of ``NamedExpr`` for certain parents (#1221)

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

3 files changed, 155 insertions, 0 deletions

diff --git a/ChangeLog b/ChangeLog
index 98b014a5..f52f6846 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -12,6 +12,10 @@ What's New in astroid 2.8.4?
 ============================
 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``.
 
 
 What's New in astroid 2.8.3?
diff --git a/astroid/nodes/node_classes.py b/astroid/nodes/node_classes.py
index c27b7689..cb1d09c9 100644
--- a/astroid/nodes/node_classes.py
+++ b/astroid/nodes/node_classes.py
@@ -4225,6 +4225,46 @@ class NamedExpr(mixins.AssignTypeMixin, NodeNG):
         self.target = target
         self.value = value
 
+    def frame(self):
+        """The first parent frame node.
+
+        A frame node is a :class:`Module`, :class:`FunctionDef`,
+        or :class:`ClassDef`.
+
+        :returns: The first parent frame node.
+        """
+        if not self.parent:
+            raise ParentMissingError(target=self)
+
+        # For certain parents NamedExpr evaluate to the scope of the parent
+        if isinstance(self.parent, (Arguments, Keyword, Comprehension)):
+            if not self.parent.parent:
+                raise ParentMissingError(target=self.parent)
+            if not self.parent.parent.parent:
+                raise ParentMissingError(target=self.parent.parent)
+            return self.parent.parent.parent.frame()
+
+        return self.parent.frame()
+
+    def scope(self) -> "LocalsDictNodeNG":
+        """The first parent node defining a new scope.
+        These can be Module, FunctionDef, ClassDef, Lambda, or GeneratorExp nodes.
+
+        :returns: The first parent scope node.
+        """
+        if not self.parent:
+            raise ParentMissingError(target=self)
+
+        # For certain parents NamedExpr evaluate to the scope of the parent
+        if isinstance(self.parent, (Arguments, Keyword, Comprehension)):
+            if not self.parent.parent:
+                raise ParentMissingError(target=self.parent)
+            if not self.parent.parent.parent:
+                raise ParentMissingError(target=self.parent.parent)
+            return self.parent.parent.parent.scope()
+
+        return self.parent.scope()
+
 
 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 05c8b4ad..9b1fee88 100644
--- a/tests/unittest_nodes.py
+++ b/tests/unittest_nodes.py
@@ -682,6 +682,117 @@ class NameNodeTest(unittest.TestCase):
             builder.parse(code)
 
 
+@pytest.mark.skipif(not PY38_PLUS, reason="needs assignment expressions")
+class TestNamedExprNode:
+    """Tests for the NamedExpr node"""
+
+    @staticmethod
+    def test_frame() -> None:
+        """Test if the frame of NamedExpr is correctly set for certain types
+        of parent nodes.
+        """
+        module = builder.parse(
+            """
+            def func(var_1):
+                pass
+
+            def func_two(var_2, var_2 = (named_expr_1 := "walrus")):
+                pass
+
+            class MyBaseClass:
+                pass
+
+            class MyInheritedClass(MyBaseClass, var_3=(named_expr_2 := "walrus")):
+                pass
+
+            VAR = lambda y = (named_expr_3 := "walrus"): print(y)
+
+            def func_with_lambda(
+                var_5 = (
+                    named_expr_4 := lambda y = (named_expr_5 := "walrus"): y
+                    )
+                ):
+                pass
+
+            COMPREHENSION = [y for i in (1, 2) if (y := i ** 2)]
+        """
+        )
+        function = module.body[0]
+        assert function.args.frame() == function
+
+        function_two = module.body[1]
+        assert function_two.args.args[0].frame() == function_two
+        assert function_two.args.args[1].frame() == function_two
+        assert function_two.args.defaults[0].frame() == module
+
+        inherited_class = module.body[3]
+        assert inherited_class.keywords[0].frame() == inherited_class
+        assert inherited_class.keywords[0].value.frame() == module
+
+        lambda_assignment = module.body[4].value
+        assert lambda_assignment.args.args[0].frame() == lambda_assignment
+        assert lambda_assignment.args.defaults[0].frame() == module
+
+        lambda_named_expr = module.body[5].args.defaults[0]
+        assert lambda_named_expr.value.args.defaults[0].frame() == module
+
+        comprehension = module.body[6].value
+        assert comprehension.generators[0].ifs[0].frame() == module
+
+    @staticmethod
+    def test_scope() -> None:
+        """Test if the scope of NamedExpr is correctly set for certain types
+        of parent nodes.
+        """
+        module = builder.parse(
+            """
+            def func(var_1):
+                pass
+
+            def func_two(var_2, var_2 = (named_expr_1 := "walrus")):
+                pass
+
+            class MyBaseClass:
+                pass
+
+            class MyInheritedClass(MyBaseClass, var_3=(named_expr_2 := "walrus")):
+                pass
+
+            VAR = lambda y = (named_expr_3 := "walrus"): print(y)
+
+            def func_with_lambda(
+                var_5 = (
+                    named_expr_4 := lambda y = (named_expr_5 := "walrus"): y
+                    )
+                ):
+                pass
+
+            COMPREHENSION = [y for i in (1, 2) if (y := i ** 2)]
+        """
+        )
+        function = module.body[0]
+        assert function.args.scope() == function
+
+        function_two = module.body[1]
+        assert function_two.args.args[0].scope() == function_two
+        assert function_two.args.args[1].scope() == function_two
+        assert function_two.args.defaults[0].scope() == module
+
+        inherited_class = module.body[3]
+        assert inherited_class.keywords[0].scope() == inherited_class
+        assert inherited_class.keywords[0].value.scope() == module
+
+        lambda_assignment = module.body[4].value
+        assert lambda_assignment.args.args[0].scope() == lambda_assignment
+        assert lambda_assignment.args.defaults[0].scope()
+
+        lambda_named_expr = module.body[5].args.defaults[0]
+        assert lambda_named_expr.value.args.defaults[0].scope() == module
+
+        comprehension = module.body[6].value
+        assert comprehension.generators[0].ifs[0].scope() == module
+
+
 class AnnAssignNodeTest(unittest.TestCase):
     def test_primitive(self) -> None:
         code = textwrap.dedent(