Use generators in the byte code parsing. __next__? yikes!

1 files changed, 19 insertions, 33 deletions

diff --git a/coverage/bytecode.py b/coverage/bytecode.py
index fd5c7da..06bc1df 100644
--- a/coverage/bytecode.py
+++ b/coverage/bytecode.py
@@ -30,7 +30,6 @@ class ByteCodes(object):
     # pylint: disable=R0924
     def __init__(self, code):
         self.code = code
-        self.offset = 0
 
     if sys.version_info >= (3, 0):
         def __getitem__(self, i):
@@ -40,32 +39,26 @@ class ByteCodes(object):
             return ord(self.code[i])
 
     def __iter__(self):
-        return self
+        offset = 0
+        while offset < len(self.code):
+            bc = ByteCode()
+            bc.op = self[offset]
+            bc.offset = offset
 
-    def __next__(self):
-        if self.offset >= len(self.code):
-            raise StopIteration
+            next_offset = offset+1
+            if bc.op >= opcode.HAVE_ARGUMENT:
+                bc.arg = self[offset+1] + 256*self[offset+2]
+                next_offset += 2
 
-        bc = ByteCode()
-        bc.op = self[self.offset]
-        bc.offset = self.offset
+                label = -1
+                if bc.op in opcode.hasjrel:
+                    label = next_offset + bc.arg
+                elif bc.op in opcode.hasjabs:
+                    label = bc.arg
+                bc.jump_to = label
 
-        next_offset = self.offset+1
-        if bc.op >= opcode.HAVE_ARGUMENT:
-            bc.arg = self[self.offset+1] + 256*self[self.offset+2]
-            next_offset += 2
-
-            label = -1
-            if bc.op in opcode.hasjrel:
-                label = next_offset + bc.arg
-            elif bc.op in opcode.hasjabs:
-                label = bc.arg
-            bc.jump_to = label
-
-        bc.next_offset = self.offset = next_offset
-        return bc
-
-    next = __next__     # Py2k uses an old-style non-dunder name.
+            bc.next_offset = offset = next_offset
+            yield bc
 
 
 class CodeObjects(object):
@@ -74,18 +67,11 @@ class CodeObjects(object):
         self.stack = [code]
 
     def __iter__(self):
-        return self
-
-    def __next__(self):
-        if self.stack:
+        while self.stack:
             # We're going to return the code object on the stack, but first
             # push its children for later returning.
             code = self.stack.pop()
             for c in code.co_consts:
                 if isinstance(c, types.CodeType):
                     self.stack.append(c)
-            return code
-
-        raise StopIteration
-
-    next = __next__
+            yield code