Support comprehensions better

--HG--
branch : ast-branch

2 files changed, 51 insertions, 12 deletions

diff --git a/coverage/parser.py b/coverage/parser.py
index 39e23d23..c680f63b 100644
--- a/coverage/parser.py
+++ b/coverage/parser.py
@@ -383,10 +383,6 @@ class AstArcAnalyzer(object):
             # TODO: test case for empty module.
             return 1
 
-    def line_comprehension(self, node):
-        # TODO: is this how to get the line number for a comprehension?
-        return node.target.lineno
-
     def line_default(self, node):
         return node.lineno
 
@@ -433,7 +429,6 @@ class AstArcAnalyzer(object):
     # TODO: multi-target assignment with computed targets
     # TODO: listcomps hidden deep in other expressions
     # TODO: listcomps hidden in lists: x = [[i for i in range(10)]]
-    # TODO: multi-line listcomps
     # TODO: nested function definitions
 
     def process_break_exits(self, exits):
@@ -554,7 +549,6 @@ class AstArcAnalyzer(object):
         return set()
 
     def handle_Return(self, node):
-        # TODO: deal with returning through a finally.
         here = self.line_for_node(node)
         self.process_return_exits([here])
         return set()
@@ -711,12 +705,9 @@ class AstArcAnalyzer(object):
                 for xit in exits:
                     self.arcs.add((xit, -start))
             elif node_name in self.CODE_COMPREHENSIONS:
-                # TODO: tests for when generators is more than one?
-                for gen in node.generators:
-                    start = self.line_for_node(gen)
-                    self.arcs.add((-1, start))
-                    self.arcs.add((start, -start))
-                    # TODO: guaranteed this won't work for multi-line comps.
+                start = self.line_for_node(node)
+                self.arcs.add((-1, start))
+                self.arcs.add((start, -start))
             elif node_name == "Lambda":
                 start = self.line_for_node(node)
                 self.arcs.add((-1, start))
diff --git a/tests/test_arcs.py b/tests/test_arcs.py
index 8b524db7..fb4b99eb 100644
--- a/tests/test_arcs.py
+++ b/tests/test_arcs.py
@@ -378,6 +378,54 @@ class LoopArcTest(CoverageTest):
             arcz=".1 .2 2-2 12 23 34 45 53 3.",
         )
 
+    def test_multiline_dict_comp(self):
+        if env.PYVERSION < (2, 7):
+            self.skip("Don't have set or dict comprehensions before 2.7")
+        if env.PY2:
+            arcz = ".2 2B B-4   2-4"
+        else:
+            arcz = ".2 2B B-3   2-3"
+        # Multiline dict comp:
+        self.check_coverage("""\
+            # comment
+            d = \\
+                {
+                i:
+                    str(i)
+                for
+                    i
+                    in
+                    range(9)
+            }
+            x = 11
+            """,
+            arcz=arcz,
+        )
+        # Multi dict comp:
+        if env.PY2:
+            arcz = ".2 2F F-4   2-4"
+        else:
+            arcz = ".2 2F F-3   2-3"
+        self.check_coverage("""\
+            # comment
+            d = \\
+                {
+                (i, j):
+                    str(i+j)
+                for
+                    i
+                    in
+                    range(9)
+                for
+                    j
+                    in
+                    range(13)
+            }
+            x = 15
+            """,
+            arcz=arcz,
+        )
+
 
 class ExceptionArcTest(CoverageTest):
     """Arc-measuring tests involving exception handling."""