Handle Starred assignments where possible

Try to match the targets with the values so we reduce the number of
false positives. Before, all right hand side variables were tainting all
of the left hand side variables.

a, *b = _, _, TAINT

a clean, b tainted

a, *b, c = _, _, TAINT, TAINT, _

a clean, b tainted, c clean

a, *b, c = _, *_, *TAINT, *_

a clean, b tainted, c tainted

Author: bcaller

examples/example_inputs/assignment_starred.py

@@ -0,0 +1 @@
+a, *b, c, d, e = f, *g, *h, f + i, j

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]
                 )))

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)
 

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):