Commits

Jendrik Seipp  committed bd68158

Find tuple assignments in comprehensions

  • Participants
  • Parent commits 394e15e

Comments (0)

Files changed (2)

File tests/test_scavenging.py

     assert v.unused_vars == []
 
 
+def test_variable6():
+    v = Vulture(verbose=True)
+    v.scan("[a for a, b in func()]")
+    assert v.defined_vars == ['a', 'b']
+    assert sorted(v.used_vars) == ['a', 'func']
+    assert v.tuple_assign_vars == ['a', 'b']
+    assert v.unused_vars == []
+
+
 def test_prop1():
     v = Vulture(verbose=True)
     v.scan("""\
     def unused_attrs(self):
         return self.get_unused(self.defined_attrs, self.used_attrs)
 
+    def _get_lineno(self, node):
+        return getattr(node, 'lineno', 1)
+
     def _get_line(self, node):
-        lineno = getattr(node, 'lineno', 1)
-        return self.code[lineno - 1] if self.code else ""
+        return self.code[self._get_lineno(node) - 1] if self.code else ""
 
     def _get_item(self, node, typ):
         name = getattr(node, 'name', None)
         print
 
     def print_node(self, node):
-        self.log(node.lineno, ast.dump(node), self._get_line(node))
+        self.log(self._get_lineno(node), ast.dump(node), self._get_line(node))
 
     def get_modules(self, paths):
         modules = []
 
     def _find_tuple_assigns(self, node):
         # Find all tuple assignments. Those have the form
-        # Assign->Tuple->Name or For->Tuple->Name
+        # Assign->Tuple->Name or For->Tuple->Name or comprehension->Tuple->Name
         for child in ast.iter_child_nodes(node):
             if not isinstance(child, ast.Tuple):
                 continue
     def visit_For(self, node):
         self._find_tuple_assigns(node)
 
+    def visit_comprehension(self, node):
+        self._find_tuple_assigns(node)
+
     def visit_ClassDef(self, node):
         self.defined_funcs.append(self._get_item(node, 'class'))