Repository: serge-sans-paille/beniget

Issue:
Warlus target should be stored in first non comprehension scope
While looking at https://github.com/gvanrossum/gvanrossum.github.io/blob/main/formal/build.py#L98-L108, I've noticed that beniget does not correctly understands the warlus operator. 

Discussion:

Patch:
diff --git a/beniget/beniget.py b/beniget/beniget.py
index 50f035b..a0fa8a3 100644
--- a/beniget/beniget.py
+++ b/beniget/beniget.py
@@ -383,8 +383,10 @@ class DefUseChains(ast.NodeVisitor):
             return ""
 
     def unbound_identifier(self, name, node):
-        location = self.location(node)
-        print("W: unbound identifier '{}'{}".format(name, location))
+        self.warn("unbound identifier '{}'".format(name), node)
+    
+    def warn(self, msg, node):
+        print("W: {}{}".format(msg, self.location(node)))
 
     def invalid_name_lookup(self, name, scope, precomputed_locals, local_defs):
         # We may hit the situation where we refer to a local variable which is
@@ -621,16 +623,16 @@ class DefUseChains(ast.NodeVisitor):
         assert not self._scope_depths
         assert not self._precomputed_locals
 
-    def set_definition(self, name, dnode_or_dnodes):
+    def set_definition(self, name, dnode_or_dnodes, index=-1):
         if self._deadcode:
             return
         # set the islive flag to False on killed Defs
-        for d in self._definitions[-1].get(name, ()):
+        for d in self._definitions[index].get(name, ()):
             d.islive = False
         if isinstance(dnode_or_dnodes, Def):
-            self._definitions[-1][name] = ordered_set((dnode_or_dnodes,))
+            self._definitions[index][name] = ordered_set((dnode_or_dnodes,))
         else:
-            self._definitions[-1][name] = ordered_set(dnode_or_dnodes)
+            self._definitions[index][name] = ordered_set(dnode_or_dnodes)
 
     @staticmethod
     def add_to_definition(definition, name, dnode_or_dnodes):
@@ -1025,11 +1027,19 @@ class DefUseChains(ast.NodeVisitor):
     def visit_Lambda(self, node, step=DeclarationStep):
         if step is DeclarationStep:
             dnode = self.chains.setdefault(node, Def(node))
+            for default in node.args.defaults:
+                self.visit(default).add_user(dnode)
+            self._defered.append((node,
+                                  list(self._definitions),
+                                  list(self._scopes),
+                                  list(self._scope_depths),
+                                  list(self._precomputed_locals)))
             return dnode
         elif step is DefinitionStep:
             dnode = self.chains[node]
             with self.ScopeContext(node):
-                self.visit(node.args)
+                for a in node.args.args:
+                    self.visit(a)
                 self.visit(node.body).add_user(dnode)
             return dnode
         else:
@@ -1058,10 +1068,15 @@ class DefUseChains(ast.NodeVisitor):
 
     def visit_ListComp(self, node):
         dnode = self.chains.setdefault(node, Def(node))
-
+        try:
+            _validate_comprehension(node)
+        except SyntaxError as e:
+            self.warn(str(e), node)
+            return dnode
         with self.CompScopeContext(node):
-            for comprehension in node.generators:
-                self.visit(comprehension).add_user(dnode)
+            for i, comprehension in enumerate(node.generators):
+                self.visit_comprehension(comprehension, 
+                                         is_nested=i!=0).add_user(dnode)
             self.visit(node.elt).add_user(dnode)
 
         return dnode
@@ -1070,10 +1085,15 @@ class DefUseChains(ast.NodeVisitor):
 
     def visit_DictComp(self, node):
         dnode = self.chains.setdefault(node, Def(node))
-
+        try:
+            _validate_comprehension(node)
+        except SyntaxError as e:
+            self.warn(str(e), node)
+            return dnode
         with self.CompScopeContext(node):
-            for comprehension in node.generators:
-                self.visit(comprehension).add_user(dnode)
+            for i, comprehension in enumerate(node.generators):
+                self.visit_comprehension(comprehension, 
+                                         is_nested=i!=0).add_user(dnode)
             self.visit(node.key).add_user(dnode)
             self.visit(node.value).add_user(dnode)
 
@@ -1137,27 +1157,54 @@ class DefUseChains(ast.NodeVisitor):
         self.visit(node.slice).add_user(dnode)
         return dnode
 
-    visit_Starred = visit_Await
+    def visit_Starred(self, node):
+        if isinstance(node.ctx, ast.Store):
+            return self.visit(node.value)
+        else:
+            dnode = self.chains.setdefault(node, Def(node))
+            self.visit(node.value).add_user(dnode)
+            return dnode
 
     def visit_NamedExpr(self, node):
         dnode = self.chains.setdefault(node, Def(node))
         self.visit(node.value).add_user(dnode)
-        self.visit(node.target)
+        if isinstance(node.target, ast.Name):
+            self.visit_Name(node.target, named_expr=True)
         return dnode
 
     def is_in_current_scope(self, name):
         return any(name in defs
                    for defs in self._definitions[self._scope_depths[-1]:])
 
-    def visit_Name(self, node, skip_annotation=False):
+    def _first_non_comprehension_scope(self):
+        index = -1
+        enclosing_scope = self._scopes[index]
+        while isinstance(enclosing_scope, (ast.DictComp, ast.ListComp, 
+                                            ast.SetComp, ast.GeneratorExp)):
+            index -= 1
+            enclosing_scope = self._scopes[index]
+        return index, enclosing_scope
+
+    def visit_Name(self, node, skip_annotation=False, named_expr=False):
         if isinstance(node.ctx, (ast.Param, ast.Store)):
             dnode = self.chains.setdefault(node, Def(node))
             if any(node.id in _globals for _globals in self._globals):
                 self.set_or_extend_global(node.id, dnode)
             else:
-                self.set_definition(node.id, dnode)
-                if dnode not in self.locals[self._scopes[-1]]:
-                    self.locals[self._scopes[-1]].append(dnode)
+                # special code for warlus target: should be 
+                # stored in first non comprehension scope
+                index, enclosing_scope = (self._first_non_comprehension_scope() 
+                                          if named_expr else (-1, self._scopes[-1]))
+
+                if index < -1 and isinstance(enclosing_scope, ast.ClassDef):
+                    # invalid named expression, not calling set_definition.
+                    self.warn('assignment expression within a comprehension '
+                              'cannot be used in a class body', node)
+                    return dnode
+            
+                self.set_definition(node.id, dnode, index)
+                if dnode not in self.locals[self._scopes[index]]:
+                    self.locals[self._scopes[index]].append(dnode)
 
             # Name.annotation is a special case because of gast
             if node.annotation is not None and not skip_annotation and not self.future_annotations:
@@ -1187,7 +1234,7 @@ class DefUseChains(ast.NodeVisitor):
                 tmp_store, elt.ctx = elt.ctx, tmp_store
                 self.visit(elt)
                 tmp_store, elt.ctx = elt.ctx, tmp_store
-            elif isinstance(elt, ast.Subscript):
+            elif isinstance(elt, (ast.Subscript, ast.Starred, ast.Attribute)):
                 self.visit(elt)
             elif isinstance(elt, (ast.List, ast.Tuple)):
                 self.visit_Destructured(elt)
@@ -1220,9 +1267,18 @@ class DefUseChains(ast.NodeVisitor):
 
     # misc
 
-    def visit_comprehension(self, node):
+    def visit_comprehension(self, node, is_nested):
         dnode = self.chains.setdefault(node, Def(node))
-        self.visit(node.iter).add_user(dnode)
+        if not is_nested and sys.version_info.major >= 3:
+            # There's one part of a comprehension or generator expression that executes in the surrounding scope, 
+            # it's the expression for the outermost iterable.
+            with self.SwitchScopeContext(self._definitions[:-1], self._scopes[:-1], 
+                                        self._scope_depths[:-1], self._precomputed_locals[:-1]):
+                self.visit(node.iter).add_user(dnode)
+        else:
+            # If a comprehension has multiple for clauses, 
+            # the iterables of the inner for clauses are evaluated in the comprehension's scope:
+            self.visit(node.iter).add_user(dnode)
         self.visit(node.target)
         for if_ in node.ifs:
             self.visit(if_).add_user(dnode)
@@ -1248,6 +1304,25 @@ class DefUseChains(ast.NodeVisitor):
             self.visit(node.optional_vars)
         return dnode
 
+def _validate_comprehension(node):
+    """
+    Raises SyntaxError if:
+     - a named expression is used in a comprehension iterable expression
+     - a named expression rebinds a comprehension iteration variable
+    """
+    iter_names = set() # comprehension iteration variables
+    for gen in node.generators:
+        for namedexpr in (n for n in ast.walk(gen.iter) if isinstance(n, ast.NamedExpr)):
+            raise SyntaxError('assignment expression cannot be used '
+                                'in a comprehension iterable expression')
+        iter_names.update(n.id for n in ast.walk(gen.target) 
+            if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Store))
+    for namedexpr in (n for n in ast.walk(node) if  isinstance(n, ast.NamedExpr)):
+        bound = getattr(namedexpr.target, 'id', None)
+        if bound in iter_names:
+            raise SyntaxError('assignment expression cannot rebind '
+                              "comprehension iteration variable '{}'".format(bound))
+
 def _iter_arguments(args):
     """
     Yields all arguments of the given ast.arguments instance.


Test patch:
diff --git a/tests/test_chains.py b/tests/test_chains.py
index e9d33f8..1e9bdaa 100644
--- a/tests/test_chains.py
+++ b/tests/test_chains.py
@@ -26,10 +26,10 @@ def captured_output():
 class TestDefUseChains(TestCase):
     def checkChains(self, code, ref, strict=True):
         class StrictDefUseChains(beniget.DefUseChains):
-            def unbound_identifier(self, name, node):
+            def warn(self, msg, node):
                 raise RuntimeError(
-                    "W: unbound identifier '{}' at {}:{}".format(
-                        name, node.lineno, node.col_offset
+                    "W: {} at {}:{}".format(
+                        msg, node.lineno, node.col_offset
                     )
                 )
 
@@ -75,6 +75,11 @@ class TestDefUseChains(TestCase):
         code = "for a, b in ((1,2), (3,4)): a"
         self.checkChains(code, ["a -> (a -> ())", "b -> ()"])
 
+    if sys.version_info.major >= 3:
+        def test_type_destructuring_starred(self):
+            code = "a, *b = range(2); b"
+            self.checkChains(code, ['a -> ()', 'b -> (b -> ())'])
+
     def test_assign_in_loop(self):
         code = "a = 2\nwhile 1: a = 1\na"
         self.checkChains(code, ["a -> (a -> ())", "a -> (a -> ())"])
@@ -332,6 +337,22 @@ while done:
     def test_simple_import_as(self):
         code = "import x as y; y()"
         self.checkChains(code, ["y -> (y -> (Call -> ()))"])
+    
+    def test_simple_lambda(self):
+        node, c = self.checkChains( "lambda y: True", [])
+        self.assertEqual(c.dump_chains(node.body[0].value), ['y -> ()'])
+    
+    def test_lambda_defaults(self):
+        node, c = self.checkChains( "x=y=1;(lambda y, x=x: (True, x, y, z)); x=y=z=2", 
+                                   ['x -> (x -> (Lambda -> ()))',
+                                    'y -> ()', 
+                                    'x -> ()', 
+                                    'y -> ()',
+                                    'z -> (z -> (Tuple -> (Lambda -> ())))']) 
+        self.assertEqual(c.dump_chains(node.body[1].value), [
+            'y -> (y -> (Tuple -> (Lambda -> ())))',
+            'x -> (x -> (Tuple -> (Lambda -> ())))',
+        ])
 
     def test_multiple_import_as(self):
         code = "import x as y, z; y"
@@ -540,7 +561,44 @@ class Visitor:
                              ['Attr -> (Attr -> (Attr -> ()))',
                               'Visitor -> ()'])
             self.assertEqual(c.dump_chains(node.body[-1]), ['Attr -> ()'])
+    
+    @skipIf(sys.version_info < (3, 0), 'Python 3 syntax')
+    def test_star_assignment(self):
+        code = '''
+curr, *parts = [1,2,3]
+while curr:
+    print(curr)
+    if parts:
+        curr, *parts = parts
+    else:
+        break
+'''
+        self.checkChains(code, ['curr -> (curr -> (), curr -> (Call -> ()))', 
+                                'parts -> (parts -> (), parts -> ())']*2)
+    
+    @skipIf(sys.version_info < (3, 0), 'Python 3 syntax')
+    def test_star_assignment_nested(self):
+        code = '''
+(curr, *parts),i = [1,2,3],0
+while curr:
+    print(curr)
+    if parts:
+        (curr, *parts),i = parts,i
+    else:
+        break
+'''
+        self.checkChains(code, ['curr -> (curr -> (), curr -> (Call -> ()))',
+                                'parts -> (parts -> (), parts -> (Tuple -> ()))',
+                                'i -> (i -> (Tuple -> ()))']*2)
+    
+    def test_attribute_assignment(self):
+        code = "d=object();d.name,x = 't',1"
+        self.checkChains(code, ['d -> (d -> (Attribute -> ()))',
+                                'x -> ()'])
 
+    def test_call_assignment(self):
+        code = "NameError().name = 't'"
+        self.checkChains(code, [])
 
     @skipIf(sys.version_info < (3, 0), 'Python 3 syntax')
     def test_annotation_uses_class_level_name(self):
@@ -597,6 +655,36 @@ cos = pop()'''
             'pop -> (pop -> (Call -> ()))',
             'cos -> (cos -> (Call -> ()))'
         ])
+    
+    @skipIf(sys.version_info < (3, 0), 'Python 3 semantics')
+    def test_class_scope_comprehension(self):
+        code = '''
+class Cls:
+    foo = b'1',
+    [_ for _ in foo]
+    {_ for _ in foo}
+    (_ for _ in foo)
+    {_:1 for _ in foo}
+'''
+        node, chains = self.checkChains(code, ['Cls -> ()'])
+        self.assertEqual(chains.dump_chains(node.body[0]),
+                         ['foo -> ('
+                          'foo -> (comprehension -> (ListComp -> ())), '
+                          'foo -> (comprehension -> (SetComp -> ())), '
+                          'foo -> (comprehension -> (GeneratorExp -> ())), '
+                          'foo -> (comprehension -> (DictComp -> ())))'])
+    
+    @skipIf(sys.version_info < (3, 0), 'Python 3 semantics')
+    def test_class_scope_comprehension_invalid(self):
+        code = '''
+class Foo:
+    x = 5
+    y = [x for i in range(1)]
+    z = [i for i in range(1) for j in range(x)]
+'''
+        self.check_message(code, ["W: unbound identifier 'x' at test:4:9", 
+                                  "W: unbound identifier 'x' at test:5:44"], 'test')
+
 
     @skipIf(sys.version_info < (3, 8), 'Python 3.8 syntax')
     def test_named_expr_simple(self):
@@ -626,6 +714,59 @@ if (a := a + a):
             code, ['a -> (a -> (BinOp -> (NamedExpr -> ())), a -> (BinOp -> (NamedExpr -> ())))', 'a -> ()']
         )
     
+    @skipIf(sys.version_info < (3, 8), 'Python 3.8 syntax')
+    def test_named_expr_comprehension(self):
+        # Warlus target should be stored in first non comprehension scope
+        code = ('cities = ["Halifax", "Toronto"]\n'
+            'if any((witness := city).startswith("H") for city in cities):'
+            'witness')
+        self.checkChains(
+            code, ['cities -> (cities -> (comprehension -> (GeneratorExp -> (Call -> ()))))', 
+                   'witness -> (witness -> ())']
+        )
+        
+    @skipIf(sys.version_info < (3, 8), 'Python 3.8 syntax')
+    def test_named_expr_comprehension_invalid(self):
+        # an assignment expression target name cannot be the same as a 
+        # for-target name appearing in any comprehension containing the assignment expression.
+        # A further exception applies when an assignment expression occurs in a comprehension whose 
+        # containing scope is a class scope. If the rules above were to result in the target 
+        # being assigned in that class's scope, the assignment expression is expressly invalid.
+        code = '''
+stuff = []
+
+# assignment expression cannot rebind comprehension iteration variable
+[[(a := a) for _ in range(5)] for a in range(5)] # INVALID
+[b := 0 for b, _ in stuff] # INVALID
+[c for c in (c := stuff)] # INVALID
+[False and (d := 0) for d, _ in stuff] # INVALID
+[_ for _, e in stuff if True or (e := 1)] # INVALID
+
+# assignment expression cannot be used in a comprehension iterable expression
+[_ for _ in (f := stuff)] # INVALID
+[_ for _ in range(2) for _ in (g := stuff)] # INVALID
+[_ for _ in [_ for _ in (h := stuff)]] # INVALID
+[_ for _ in (lambda: (i := stuff))()] # INVALID
+
+class Example:
+    # assignment expression within a comprehension cannot be used in a class body
+    [(j := i) for i in range(5)] # INVALID
+'''
+        # None of the invalid assigned name shows up.
+        node, chains = self.checkChains(code, ['stuff -> ()', 'Example -> ()'], strict=False)        
+        self.assertEqual(chains.dump_chains(node.body[-1]), [])
+        # It triggers useful warnings
+        self.check_message(code, ["W: assignment expression cannot rebind comprehension iteration variable 'a' at <unknown>:5:0", 
+                                  "W: assignment expression cannot rebind comprehension iteration variable 'b' at <unknown>:6:0", 
+                                  'W: assignment expression cannot be used in a comprehension iterable expression at <unknown>:7:0', 
+                                  "W: assignment expression cannot rebind comprehension iteration variable 'd' at <unknown>:8:0", 
+                                  "W: assignment expression cannot rebind comprehension iteration variable 'e' at <unknown>:9:0", 
+                                  'W: assignment expression cannot be used in a comprehension iterable expression at <unknown>:12:0', 
+                                  'W: assignment expression cannot be used in a comprehension iterable expression at <unknown>:13:0', 
+                                  'W: assignment expression cannot be used in a comprehension iterable expression at <unknown>:14:0', 
+                                  'W: assignment expression cannot be used in a comprehension iterable expression at <unknown>:15:0',
+                                  'W: assignment expression within a comprehension cannot be used in a class body at <unknown>:19:6'])
+    
     @skipIf(sys.version_info.major < 3, "Python 3 syntax")
     def test_annotation_unbound(self):
         code = '''
diff --git a/tests/test_definitions.py b/tests/test_definitions.py
index 17f80e2..3e432a2 100644
--- a/tests/test_definitions.py
+++ b/tests/test_definitions.py
@@ -89,6 +89,11 @@ class TestGlobals(TestCase):
         code = "x, y = 1, 2"
         self.checkGlobals(code, ["x", "y"])
 
+    if sys.version_info.major >= 3:
+        def testGlobalStarredDestructuring(self):
+            code = "x, *y = 1, [2]"
+            self.checkGlobals(code, ["x", "y"])
+
     def testGlobalAugAssign(self):
         code = "x = 1; x += 2"
         self.checkGlobals(code, ["x"])
@@ -369,6 +374,10 @@ class TestLocals(TestCase):
             )
             self.checkLocals(code, ["a", "bar"])
 
+        def test_LocalDestructuring(self):
+            code = "def foo(x): y, *z = x"
+            self.checkLocals(code, ["x", "y", "z"])
+
     def test_LocalMadeGlobal(self):
         code = "def foo(): global a; a = 1"
         self.checkLocals(code, [])
@@ -468,4 +477,4 @@ class TestDefIsLive(TestCase):
             a = a + 1
             a = a + 1
             """
-        self.checkLiveLocals(code, ["a:4"], ["a:2,3,4"])
\ No newline at end of file
+        self.checkLiveLocals(code, ["a:4"], ["a:2,3,4"])