Including generator yield statements when doing end of code calculations issue 324

1 files changed, 13 insertions, 0 deletions

diff --git a/tests/test_parser.py b/tests/test_parser.py
index 244d4c7..a39820e 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -34,6 +34,19 @@ class PythonParserTest(CoverageTest):
             2:1, 3:1, 4:2, 5:1, 7:1, 9:1, 10:1
             })
 
+    def test_generator_exit_counts(self):
+        parser = self.parse_source("""\
+            # generators yield lines should only have one exit count
+            def gen(input):
+                for n in inp:
+                    yield (i * 2 for i in range(n))
+
+            list(gen([1,2,3]))
+            """)
+        self.assertEqual(parser.exit_counts(), {
+           2:1, 3:2, 4:1, 6:1
+        })
+
     def test_try_except(self):
         parser = self.parse_source("""\
             try: