Patch #1006219: let inspect.getsource show '@' decorators and add tests for

this (which are rather ugly, but it'll have to do until test_inspect gets a
major overhaul and a conversion to unittest). Thanks Simon Percivall!

2 files changed, 44 insertions, 2 deletions

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 42eda77..3c47cd9 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -433,7 +433,7 @@ def findsource(object):
         if not hasattr(object, 'co_firstlineno'):
             raise IOError('could not find function definition')
         lnum = object.co_firstlineno - 1
-        pat = re.compile(r'^(\s*def\s)|(.*\slambda(:|\s))')
+        pat = re.compile(r'^(\s*def\s)|(.*\slambda(:|\s))|^(\s*@)')
         while lnum > 0:
             if pat.match(lines[lnum]): break
             lnum = lnum - 1
@@ -509,7 +509,8 @@ class BlockFinder:
 
     def tokeneater(self, type, token, (srow, scol), (erow, ecol), line):
         if not self.started:
-            if type == tokenize.NAME: self.started = 1
+            if '@' in line: pass
+            elif type == tokenize.NAME: self.started = 1
         elif type == tokenize.NEWLINE:
             self.last = srow
         elif type == tokenize.INDENT:
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 2f79fc3..e3bf73a 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -202,6 +202,47 @@ for fname in files_to_clean_up:
     except:
         pass
 
+# Test for decorators as well.
+
+source = r"""
+def wrap(foo=None):
+  def wrapper(func):
+    return func
+  return wrapper
+
+def replace(func):
+  def insteadfunc():
+    print 'hello'
+  return insteadfunc
+
+# two decorators, one with argument
+@wrap()
+@wrap(wrap)
+def wrapped():
+  pass
+
+@replace
+def gone():
+  pass"""
+
+file = open(TESTFN + "2", "w")
+file.write(source)
+file.close()
+files_to_clean_up = [TESTFN + "2", TESTFN + '2c', TESTFN + '2o']
+
+mod2 = imp.load_source("testmod3", TESTFN + "2")
+
+test(inspect.getsource(mod2.wrapped) == sourcerange(13, 16),
+     "inspect.getsource(mod.wrapped)")
+test(inspect.getsource(mod2.gone) == sourcerange(8, 9),
+     "inspect.getsource(mod.gone)")
+
+for fname in files_to_clean_up:
+    try:
+        os.unlink(fname)
+    except:
+        pass
+
 # Test classic-class method resolution order.
 class A:    pass
 class B(A): pass