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

2 files changed, 14 insertions, 1 deletions

diff --git a/coverage/parser.py b/coverage/parser.py
index f488367d..91b685ce 100644
--- a/coverage/parser.py
+++ b/coverage/parser.py
@@ -308,7 +308,7 @@ OPS_CODE_END = _opcode_set('RETURN_VALUE')
 # Opcodes that unconditionally end the code chunk.
 OPS_CHUNK_END = _opcode_set(
     'JUMP_ABSOLUTE', 'JUMP_FORWARD', 'RETURN_VALUE', 'RAISE_VARARGS',
-    'BREAK_LOOP', 'CONTINUE_LOOP',
+    'BREAK_LOOP', 'CONTINUE_LOOP', 'YIELD_VALUE',
     )
 
 # Opcodes that unconditionally begin a new code chunk.  By starting new chunks
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 244d4c70..a39820e6 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: