1 files changed, 23 insertions, 19 deletions
diff --git a/Lib/inspect.py b/Lib/inspect.py
index eeb54d2..2d88bc1 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -349,28 +349,32 @@ class ListReader:
             return self.lines[i]
         else: return ''
 
-def getblock(lines):
-    """Extract the block of code at the top of the given list of lines."""
-
-    indent = 0
-    started = 0
-    last = 0
-    tokens = tokenize.generate_tokens(ListReader(lines).readline)
-
-    for (type, token, (srow, scol), (erow, ecol), line) in tokens:
-        if not started:
-            if type == tokenize.NAME:
-                started = 1
+class EndOfBlock(Exception): pass
+
+class BlockFinder:
+    """Provide a tokeneater() method to detect the end of a code block."""
+    def __init__(self):
+        self.indent = 0
+        self.started = 0
+        self.last = 0
+
+    def tokeneater(self, type, token, (srow, scol), (erow, ecol), line):
+        if not self.started:
+            if type == tokenize.NAME: self.started = 1
         elif type == tokenize.NEWLINE:
-            last = srow
+            self.last = srow
         elif type == tokenize.INDENT:
-            indent = indent + 1
+            self.indent = self.indent + 1
         elif type == tokenize.DEDENT:
-            indent = indent - 1
-            if indent == 0:
-                return lines[:last]
-    else:
-        raise ValueError, "unable to find block"
+            self.indent = self.indent - 1
+            if self.indent == 0: raise EndOfBlock, self.last
+
+def getblock(lines):
+    """Extract the block of code at the top of the given list of lines."""
+    try:
+        tokenize.tokenize(ListReader(lines).readline, BlockFinder().tokeneater)
+    except EndOfBlock, eob:
+        return lines[:eob.args[0]]
 
 def getsourcelines(object):
     """Return a list of source lines and starting line number for an object.