Make other comprehensions work on py2 and py3

--HG--
branch : ast-branch

1 files changed, 36 insertions, 1 deletions

diff --git a/tests/test_arcs.py b/tests/test_arcs.py
index 0407b560..a371401f 100644
--- a/tests/test_arcs.py
+++ b/tests/test_arcs.py
@@ -260,7 +260,7 @@ class LoopArcTest(CoverageTest):
         if env.PY3:
             arcz = ".1 12 23 34 45 36 63 57 7."
         else:
-            arcz = ".1 12 23 27 34 45 36 62 57 7."
+            arcz = ".1 12 23 34 45 36 62 57 7."
         self.check_coverage("""\
             a, i = 1, 0
             while True:
@@ -341,6 +341,41 @@ class LoopArcTest(CoverageTest):
             """,
             arcz=arcz, arcz_missing="", arcz_unpredicted="")
 
+    def test_other_comprehensions(self):
+        # Generator expression:
+        self.check_coverage("""\
+            o = ((1,2), (3,4))
+            o = (a for a in o)
+            for tup in o:
+                x = tup[0]
+                y = tup[1]
+            """,
+            arcz=".1 .2 2-2 12 23 34 45 53 3.",
+            arcz_missing="", arcz_unpredicted=""
+        )
+        # Set comprehension:
+        self.check_coverage("""\
+            o = ((1,2), (3,4))
+            o = {a for a in o}
+            for tup in o:
+                x = tup[0]
+                y = tup[1]
+            """,
+            arcz=".1 .2 2-2 12 23 34 45 53 3.",
+            arcz_missing="", arcz_unpredicted=""
+        )
+        # Dict comprehension:
+        self.check_coverage("""\
+            o = ((1,2), (3,4))
+            o = {a:1 for a in o}
+            for tup in o:
+                x = tup[0]
+                y = tup[1]
+            """,
+            arcz=".1 .2 2-2 12 23 34 45 53 3.",
+            arcz_missing="", arcz_unpredicted=""
+        )
+
 
 class ExceptionArcTest(CoverageTest):
     """Arc-measuring tests involving exception handling."""