Merge pull request #150 from bcaller/starred-tuple-assign

KevinHock · web-flow · commit 6d81d9185956 · 2018-07-24T11:38:41.000-07:00
Starred tuple assignment
diff --git a/examples/example_inputs/assignment_starred.py b/examples/example_inputs/assignment_starred.py
@@ -0,0 +1 @@
+a, *b, c, d, e = f, *g, *h, f + i, j
diff --git a/pyt/cfg/stmt_visitor.py b/pyt/cfg/stmt_visitor.py
@@ -327,28 +327,59 @@ 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):
-        new_assignment_nodes = list()
-        for i, target in enumerate(node.targets[0].elts):
-            value = node.value.elts[i]
+        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
 
+        def visit(target, value):
             label = LabelVisitor()
             label.visit(target)
-
+            rhs_visitor = RHSVisitor()
+            rhs_visitor.visit(value)
             if isinstance(value, ast.Call):
                 new_ast_node = ast.Assign(target, value)
-                new_ast_node.lineno = node.lineno
-
+                ast.copy_location(new_ast_node, node)
                 new_assignment_nodes.append(self.assignment_call_node(label.result, new_ast_node))
-
             else:
                 label.result += ' = '
                 label.visit(value)
-
                 new_assignment_nodes.append(self.append_node(AssignmentNode(
                     label.result,
                     extract_left_hand_side(target),
                     ast.Assign(target, value),
-                    right_hand_side_variables,
+                    rhs_visitor.result,
+                    line_number=node.lineno,
+                    path=self.filenames[-1]
+                )))
+            remaining_targets.remove(target)
+            remaining_values.remove(value)
+            for var in rhs_visitor.result:
+                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):
+            if isinstance(target, ast.Starred) or isinstance(value, ast.Starred):
+                break
+            visit(target, value)
+
+        # If there was a Starred node, pair remaining targets and values from the end
+        for target, value in zip(reversed(list(remaining_targets)), reversed(list(remaining_values))):
+            if isinstance(target, ast.Starred) or isinstance(value, ast.Starred):
+                break
+            visit(target, value)
+
+        if remaining_targets:
+            label = LabelVisitor()
+            label.handle_comma_separated(remaining_targets)
+            label.result += ' = '
+            label.handle_comma_separated(remaining_values)
+            for target in remaining_targets:
+                new_assignment_nodes.append(self.append_node(AssignmentNode(
+                    label.result,
+                    extract_left_hand_side(target),
+                    ast.Assign(target, remaining_values[0]),
+                    remaining_variables,
                     line_number=node.lineno,
                     path=self.filenames[-1]
                 )))
@@ -380,8 +411,8 @@ def assign_multi_target(self, node, right_hand_side_variables):
     def visit_Assign(self, node):
         rhs_visitor = RHSVisitor()
         rhs_visitor.visit(node.value)
-        if isinstance(node.targets[0], ast.Tuple):  # x,y = [1,2]
-            if isinstance(node.value, ast.Tuple):
+        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)
             elif isinstance(node.value, ast.Call):
                 call = None
diff --git a/pyt/cfg/stmt_visitor_helper.py b/pyt/cfg/stmt_visitor_helper.py
@@ -79,6 +79,8 @@ def _get_names(node, result):
         return node.id + result
     elif isinstance(node, ast.Subscript):
         return result
+    elif isinstance(node, ast.Starred):
+        return _get_names(node.value, result)
     else:
         return _get_names(node.value, result + '.' + node.attr)
 
diff --git a/pyt/helper_visitors/label_visitor.py b/pyt/helper_visitors/label_visitor.py
@@ -320,3 +320,7 @@ def visit_FormattedValue(self, node):
             self.result += ':'
             self.visit_joined_str(node.format_spec)
         self.result += '}'
+
+    def visit_Starred(self, node):
+        self.result += '*'
+        self.visit(node.value)
diff --git a/tests/cfg/cfg_test.py b/tests/cfg/cfg_test.py
@@ -1,3 +1,5 @@
+import ast
+
 from .cfg_base_test_case import CFGBaseTestCase
 
 from pyt.core.node_types import (
@@ -779,6 +781,45 @@ def test_assignment_tuple_value(self):
 
         self.assertEqual(self.cfg.nodes[node].label, 'a = (x, y)')
 
+    def test_assignment_starred(self):
+        self.cfg_create_from_file('examples/example_inputs/assignment_starred.py')
+
+        middle_nodes = self.cfg.nodes[1:-1]
+        self.assert_length(middle_nodes, expected_length=5)
+
+        visited = [self.cfg.nodes[0]]
+        while True:
+            current_node = visited[-1]
+            if len(current_node.outgoing) != 1:
+                break
+            visited.append(current_node.outgoing[0])
+        self.assertCountEqual(self.cfg.nodes, visited, msg="Did not complete a path from Entry to Exit")
+
+        self.assertEqual(middle_nodes[0].label, 'a = f')
+        self.assertCountEqual(  # We don't assert a specific order for the assignment nodes
+            [n.label for n in middle_nodes],
+            ['a = f', 'd = f + i', 'e = j'] + ['*b, c = *g, *h'] * 2,
+        )
+        self.assertCountEqual(
+            [(n.left_hand_side, n.right_hand_side_variables) for n in middle_nodes],
+            [('a', ['f']), ('b', ['g', 'h']), ('c', ['g', 'h']), ('d', ['f', 'i']), ('e', ['j'])],
+        )
+
+    def test_assignment_starred_list(self):
+        self.cfg_create_from_ast(ast.parse('[a, b, c] = *d, e'))
+
+        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 = *d', 'a, b = *d', 'c = e'],
+        )
+        self.assertCountEqual(
+            [(n.left_hand_side, n.right_hand_side_variables) for n in middle_nodes],
+            [('a', ['d']), ('b', ['d']), ('c', ['e'])],
+        )
+
 
 class CFGComprehensionTest(CFGBaseTestCase):
     def test_nodes(self):
diff --git a/tests/helper_visitors/label_visitor_test.py b/tests/helper_visitors/label_visitor_test.py
@@ -79,3 +79,7 @@ def test_joined_str(self):
     def test_joined_str_with_format_spec(self):
         label = self.perform_labeling_on_expression('f"a{b!s:.{length}}"')
         self.assertEqual(label.result, 'f\'a{b!s:.{length}}\'')
+
+    def test_starred(self):
+        label = self.perform_labeling_on_expression('[a, *b] = *c, d')
+        self.assertEqual(label.result, '[a, *b] = (*c, d)')
