Merge pull request #164 from bcaller/unpacking

Handle assignment unpacking a, b, c = d

Author: KevinHock

pyt/cfg/stmt_visitor.py

@@ -326,11 +326,11 @@ def visit_Try(self, node):
 
         return ControlFlowNode(try_node, last_statements, break_statements=body.break_statements)
 
-    def assign_tuple_target(self, node, right_hand_side_variables):
+    def assign_tuple_target(self, target_nodes, value_nodes, right_hand_side_variables):
         new_assignment_nodes = []
         remaining_variables = list(right_hand_side_variables)
-        remaining_targets = list(node.targets[0].elts)
-        remaining_values = list(node.value.elts)  # May contain duplicates
+        remaining_targets = list(target_nodes)
+        remaining_values = list(value_nodes)  # May contain duplicates
 
         def visit(target, value):
             label = LabelVisitor()
@@ -339,7 +339,7 @@ def visit(target, value):
             rhs_visitor.visit(value)
             if isinstance(value, ast.Call):
                 new_ast_node = ast.Assign(target, value)
-                ast.copy_location(new_ast_node, node)
+                ast.copy_location(new_ast_node, target)
                 new_assignment_nodes.append(self.assignment_call_node(label.result, new_ast_node))
             else:
                 label.result += ' = '
@@ -349,7 +349,7 @@ def visit(target, value):
                     extract_left_hand_side(target),
                     ast.Assign(target, value),
                     rhs_visitor.result,
-                    line_number=node.lineno,
+                    line_number=target.lineno,
                     path=self.filenames[-1]
                 )))
             remaining_targets.remove(target)
@@ -358,7 +358,7 @@ def visit(target, value):
                 remaining_variables.remove(var)
 
         # Pair targets and values until a Starred node is reached
-        for target, value in zip(node.targets[0].elts, node.value.elts):
+        for target, value in zip(target_nodes, value_nodes):
             if isinstance(target, ast.Starred) or isinstance(value, ast.Starred):
                 break
             visit(target, value)
@@ -380,7 +380,7 @@ def visit(target, value):
                     extract_left_hand_side(target),
                     ast.Assign(target, remaining_values[0]),
                     remaining_variables,
-                    line_number=node.lineno,
+                    line_number=target.lineno,
                     path=self.filenames[-1]
                 )))
 
@@ -413,14 +413,18 @@ def visit_Assign(self, node):
         rhs_visitor.visit(node.value)
         if isinstance(node.targets[0], (ast.Tuple, ast.List)):  # x,y = [1,2]
             if isinstance(node.value, (ast.Tuple, ast.List)):
-                return self.assign_tuple_target(node, rhs_visitor.result)
+                return self.assign_tuple_target(node.targets[0].elts, node.value.elts, rhs_visitor.result)
             elif isinstance(node.value, ast.Call):
                 call = None
                 for element in node.targets[0].elts:
                     label = LabelVisitor()
                     label.visit(element)
                     call = self.assignment_call_node(label.result, node)
                 return call
+            elif isinstance(node.value, ast.Name):  # Treat `x, y = z` like `x, y = (*z,)`
+                value_node = ast.Starred(node.value, ast.Load())
+                ast.copy_location(value_node, node)
+                return self.assign_tuple_target(node.targets[0].elts, [value_node], rhs_visitor.result)
             else:
                 label = LabelVisitor()
                 label.visit(node)

tests/cfg/cfg_test.py

@@ -820,6 +820,21 @@ def test_assignment_starred_list(self):
             [('a', ['d']), ('b', ['d']), ('c', ['e'])],
         )
 
+    def test_unpacking_to_tuple(self):
+        self.cfg_create_from_ast(ast.parse('a, b, c = d'))
+
+        middle_nodes = self.cfg.nodes[1:-1]
+        self.assert_length(middle_nodes, expected_length=3)
+
+        self.assertCountEqual(
+            [n.label for n in middle_nodes],
+            ['a, b, c = *d'] * 3,
+        )
+        self.assertCountEqual(
+            [(n.left_hand_side, n.right_hand_side_variables) for n in middle_nodes],
+            [('a', ['d']), ('b', ['d']), ('c', ['d'])],
+        )
+
     def test_augmented_assignment(self):
         self.cfg_create_from_ast(ast.parse('a+=f(b,c)'))