From dd5c023af56b9a424a3965364521c50b73308379 Mon Sep 17 00:00:00 2001
From: Armin Rigo <arigo@tunes.org>
Date: Sun, 25 Sep 2005 11:45:45 +0000
Subject: [PATCH] some more fixes and tests for inspect.getsource(), triggered
 by crashes from the PyPy project as well as the SF bug #1295909.

---
 Lib/inspect.py              | 47 ++++++++++++++++---------------------
 Lib/test/inspect_fodder2.py | 24 +++++++++++++++++++
 Lib/test/test_inspect.py    | 12 ++++++++++
 3 files changed, 56 insertions(+), 27 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index a801a298eac..57bf18cd29b 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -485,19 +485,6 @@ def getcomments(object):
                 comments[-1:] = []
             return string.join(comments, '')
 
-class ListReader:
-    """Provide a readline() method to return lines from a list of strings."""
-    def __init__(self, lines):
-        self.lines = lines
-        self.index = 0
-
-    def readline(self):
-        i = self.index
-        if i < len(self.lines):
-            self.index = i + 1
-            return self.lines[i]
-        else: return ''
-
 class EndOfBlock(Exception): pass
 
 class BlockFinder:
@@ -507,40 +494,46 @@ class BlockFinder:
         self.islambda = False
         self.started = False
         self.passline = False
-        self.last = 0
+        self.last = 1
 
     def tokeneater(self, type, token, (srow, scol), (erow, ecol), line):
         if not self.started:
+            # look for the first "def", "class" or "lambda"
             if token in ("def", "class", "lambda"):
                 if token == "lambda":
                     self.islambda = True
                 self.started = True
-            self.passline = True
+            self.passline = True    # skip to the end of the line
         elif type == tokenize.NEWLINE:
-            self.passline = False
+            self.passline = False   # stop skipping when a NEWLINE is seen
             self.last = srow
+            if self.islambda:       # lambdas always end at the first NEWLINE
+                raise EndOfBlock
         elif self.passline:
             pass
-        elif self.islambda:
-            raise EndOfBlock, self.last
         elif type == tokenize.INDENT:
             self.indent = self.indent + 1
             self.passline = True
         elif type == tokenize.DEDENT:
             self.indent = self.indent - 1
-            if self.indent == 0:
-                raise EndOfBlock, self.last
-        elif type == tokenize.NAME and scol == 0:
-            raise EndOfBlock, self.last
+            # the end of matching indent/dedent pairs end a block
+            # (note that this only works for "def"/"class" blocks,
+            #  not e.g. for "if: else:" or "try: finally:" blocks)
+            if self.indent <= 0:
+                raise EndOfBlock
+        elif self.indent == 0 and type not in (tokenize.COMMENT, tokenize.NL):
+            # any other token on the same indentation level end the previous
+            # block as well, except the pseudo-tokens COMMENT and NL.
+            raise EndOfBlock
 
 def getblock(lines):
     """Extract the block of code at the top of the given list of lines."""
+    blockfinder = BlockFinder()
     try:
-        tokenize.tokenize(ListReader(lines).readline, BlockFinder().tokeneater)
-    except EndOfBlock, eob:
-        return lines[:eob.args[0]]
-    # Fooling the indent/dedent logic implies a one-line definition
-    return lines[:1]
+        tokenize.tokenize(iter(lines).next, blockfinder.tokeneater)
+    except (EndOfBlock, IndentationError):
+        pass
+    return lines[:blockfinder.last]
 
 def getsourcelines(object):
     """Return a list of source lines and starting line number for an object.
diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index f216c825a90..f150ec6af5f 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -64,3 +64,27 @@ multiline_sig = [
             y): x+y,
     None,
     ]
+
+# line 68
+def func69():
+    class cls70:
+        def func71():
+            pass
+    return cls70
+extra74 = 74
+
+# line 76
+def func77(): pass
+(extra78, stuff78) = 'xy'
+extra79 = 'stop'
+
+# line 81
+class cls82:
+    def func83(): pass
+(extra84, stuff84) = 'xy'
+extra85 = 'stop'
+
+# line 87
+def func88():
+    # comment
+    return 90
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 1fb48c526de..ce346b9b211 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -238,6 +238,18 @@ class TestBuggyCases(GetSourceBase):
     def test_multiline_sig(self):
         self.assertSourceEqual(mod2.multiline_sig[0], 63, 64)
 
+    def test_nested_class(self):
+        self.assertSourceEqual(mod2.func69().func71, 71, 72)
+
+    def test_one_liner_followed_by_non_name(self):
+        self.assertSourceEqual(mod2.func77, 77, 77)
+
+    def test_one_liner_dedent_non_name(self):
+        self.assertSourceEqual(mod2.cls82.func83, 83, 83)
+
+    def test_with_comment_instead_of_docstring(self):
+        self.assertSourceEqual(mod2.func88, 88, 90)
+
 # Helper for testing classify_class_attrs.
 def attrs_wo_objs(cls):
     return [t[:3] for t in inspect.classify_class_attrs(cls)]