Remember the source of interactively defined functions

michaelmulley · michaelmulley · commit 38f1ed6bc302 · 2015-03-25T16:19:23.000-04:00
- Patches the linecache module to achieve this - inspect.getsource() works for functions defined on the console - tracebacks show code for frames inside interactively defined functions fixes bpython#229
diff --git a/bpython/history.py b/bpython/history.py
@@ -25,6 +25,7 @@
 
 from __future__ import unicode_literals
 import io
+import linecache
 import os
 import stat
 from itertools import islice
@@ -232,3 +233,74 @@ def append_reload_and_write(self, s, filename, encoding):
                 # Make sure that entries contains at least one element. If the
                 # file and s are empty, this can occur.
                 self.entries = ['']
+
+class BPythonLinecache(dict):
+    """Replaces the cache dict in the standard-library linecache module,
+    to also remember (in an unerasable way) bpython console input."""
+
+    def __init__(self, *args, **kwargs):
+        dict.__init__(self, *args, **kwargs)
+        self.bpython_history = []
+
+    def is_bpython_filename(self, fname):
+        try:
+            return fname.startswith('<bpython-input-')
+        except AttributeError:
+            # In case the key isn't a string
+            return False
+
+    def get_bpython_history(self, key):
+        """Given a filename provided by remember_bpython_input,
+        returns the associated source string."""
+        try:
+            idx = int(key.split('-')[2][:-1])
+            return self.bpython_history[idx]
+        except (IndexError, ValueError):
+            raise KeyError
+
+    def remember_bpython_input(self, source):
+        """Remembers a string of source code, and returns
+        a fake filename to use to retrieve it later."""
+        filename = '<bpython-input-%s>' % len(self.bpython_history)
+        self.bpython_history.append((len(source), None,
+            source.splitlines(True), filename))
+        return filename
+
+    def __getitem__(self, key):
+        if self.is_bpython_filename(key):
+            return self.get_bpython_history(key)
+        return dict.__getitem__(self, key)
+
+    def __contains__(self, key):
+        if self.is_bpython_filename(key):
+            try:
+                self.get_bpython_history(key)
+                return True
+            except KeyError:
+                return False
+        return dict.__contains__(self, key)
+
+    def __delitem__(self, key):
+        if not self.is_bpython_filename(key):
+            return dict.__delitem__(self, key)
+
+def _bpython_clear_linecache():
+    try:
+        bpython_history = linecache.cache.bpython_history
+    except AttributeError:
+        bpython_history = []
+    linecache.cache = BPythonLinecache()
+    linecache.cache.bpython_history = bpython_history
+
+# Monkey-patch the linecache module so that we're able
+# to hold our command history there and have it persist
+linecache.cache = BPythonLinecache(linecache.cache)
+linecache.clearcache = _bpython_clear_linecache
+
+def filename_for_console_input(code_string):
+    """Remembers a string of source code, and returns
+    a fake filename to use to retrieve it later."""
+    try:
+        return linecache.cache.remember_bpython_input(code_string)
+    except AttributeError:
+        return '<input>'
diff --git a/bpython/repl.py b/bpython/repl.py
@@ -47,7 +47,7 @@
 from bpython.clipboard import get_clipboard, CopyFailed
 from bpython.config import getpreferredencoding
 from bpython.formatter import Parenthesis
-from bpython.history import History
+from bpython.history import History, filename_for_console_input
 from bpython.paste import PasteHelper, PastePinnwand, PasteFailed
 from bpython.translations import _, ngettext
 
@@ -94,7 +94,7 @@ def __init__(self, locals=None, encoding=None):
     def reset_running_time(self):
         self.running_time = 0
 
-    def runsource(self, source, filename='<input>', symbol='single',
+    def runsource(self, source, filename=None, symbol='single',
                   encode=True):
         """Execute Python code.
 
@@ -104,6 +104,8 @@ def runsource(self, source, filename='<input>', symbol='single',
         if not py3 and encode:
             source = u'# coding: %s\n%s' % (self.encoding, source)
             source = source.encode(self.encoding)
+        if filename is None:
+            filename = filename_for_console_input(source)
         with self.timer:
             return code.InteractiveInterpreter.runsource(self, source,
                                                          filename, symbol)
diff --git a/bpython/test/test_interpreter.py b/bpython/test/test_interpreter.py
@@ -17,6 +17,14 @@
 
 pypy = 'PyPy' in sys.version
 
+def _last_console_filename():
+    """Returns the last 'filename' used for console input
+    (as will be displayed in a traceback)."""
+    import linecache
+    try:
+        return '<bpython-input-%s>' % (len(linecache.cache.bpython_history) - 1)
+    except AttributeError:
+        return '<input>'
 
 class TestInterpreter(unittest.TestCase):
     def test_syntaxerror(self):
@@ -30,11 +38,11 @@ def append_to_a(message):
         i.runsource('1.1.1.1')
 
         if pypy:
-            expected = '  File ' + green('"<input>"') + ', line ' + \
+            expected = '  File ' + green('"%s"' % _last_console_filename()) + ', line ' + \
                 bold(magenta('1')) + '\n    1.1.1.1\n      ^\n' + \
                 bold(red('SyntaxError')) + ': ' + cyan('invalid syntax') + '\n'
         else:
-            expected = '  File ' + green('"<input>"') + ', line ' + \
+            expected = '  File ' + green('"%s"' % _last_console_filename()) + ', line ' + \
                 bold(magenta('1')) + '\n    1.1.1.1\n        ^\n' + \
                 bold(red('SyntaxError')) + ': ' + cyan('invalid syntax') + '\n'
 
@@ -56,16 +64,16 @@ def f():
         def g():
             return f()
 
-        i.runsource('g()')
+        i.runsource('g()', encode=False)
 
         if pypy:
             global_not_found = "global name 'g' is not defined"
         else:
             global_not_found = "name 'g' is not defined"
 
         expected = 'Traceback (most recent call last):\n  File ' + \
-            green('"<input>"') + ', line ' + bold(magenta('1')) + ', in ' + \
-            cyan('<module>') + '\n' + bold(red('NameError')) + ': ' + \
+            green('"%s"' % _last_console_filename()) + ', line ' + bold(magenta('1')) + ', in ' + \
+            cyan('<module>') + '\n    g()\n' + bold(red('NameError')) + ': ' + \
             cyan(global_not_found) + '\n'
 
         self.assertMultiLineEqual(str(plain('').join(a)), str(expected))
@@ -106,3 +114,11 @@ def test_runsource_unicode(self):
         i.runsource("a = u'\xfe'", encode=True)
         self.assertIsInstance(i.locals['a'], type(u''))
         self.assertEqual(i.locals['a'], u"\xfe")
+
+    def test_getsource_works_on_interactively_defined_functions(self):
+        source = 'def foo(x):\n    return x + 1\n'
+        i = interpreter.Interp()
+        i.runsource(source)
+        import inspect
+        inspected_source = inspect.getsource(i.locals['foo'])
+        self.assertEquals(inspected_source, source)
