Merge pull request bpython#368 from mlauter/master

thomasballinger · thomasballinger · commit 0513f1a2cd95 · 2014-09-04T12:56:32.000-04:00
Fix bpython#194
diff --git a/bpython/curtsiesfrontend/interpreter.py b/bpython/curtsiesfrontend/interpreter.py
@@ -0,0 +1,155 @@
+import code
+import traceback
+import sys
+from pygments.style import Style
+from pygments.token import *
+from pygments.formatter import Formatter
+from curtsies.bpythonparse import parse
+from codeop import CommandCompiler, compile_command
+from pygments.lexers import get_lexer_by_name
+from pygments.styles import get_style_by_name
+
+default_colors = {
+        Generic.Error:'R',
+        Keyword:'d',
+        Name:'c',
+        Name.Builtin:'g',
+        Comment:'b',
+        String:'m',
+        Error:'r',
+        Literal:'d',
+        Number:'M',
+        Number.Integer:'d',
+        Operator:'d',
+        Punctuation:'d',
+        Token:'d',
+        Whitespace:'d',
+        Token.Punctuation.Parenthesis:'R',
+        Name.Function:'d',
+        Name.Class:'d',
+    }
+ 
+class BPythonFormatter(Formatter):
+    """This is subclassed from the custom formatter for bpython.
+    Its format() method receives the tokensource
+    and outfile params passed to it from the
+    Pygments highlight() method and slops
+    them into the appropriate format string
+    as defined above, then writes to the outfile
+    object the final formatted string. This does not write real strings. It writes format string (FmtStr) objects.
+ 
+    See the Pygments source for more info; it's pretty
+    straightforward."""
+ 
+    def __init__(self, color_scheme, **options):
+        self.f_strings = {}
+        for k, v in color_scheme.iteritems():
+            self.f_strings[k] = '\x01%s' % (v,)
+        Formatter.__init__(self, **options)
+ 
+    def format(self, tokensource, outfile):
+        o = ''
+
+        for token, text in tokensource:
+            while token not in self.f_strings:
+                token = token.parent
+            o += "%s\x03%s\x04" % (self.f_strings[token], text)
+        outfile.write(parse(o.rstrip()))
+ 
+class Interp(code.InteractiveInterpreter):
+    def __init__(self, locals=None):
+        """Constructor.
+
+        The optional 'locals' argument specifies the dictionary in
+        which code will be executed; it defaults to a newly created
+        dictionary with key "__name__" set to "__console__" and key
+        "__doc__" set to None.
+
+        We include an argument for the outfile to pass to the formatter for it to write to.
+
+        """
+        if locals is None:
+            locals = {"__name__": "__console__", "__doc__": None}
+        self.locals = locals
+        self.compile = CommandCompiler()
+
+        # typically changed after being instantiated
+        self.write = lambda stuff: sys.stderr.write(stuff)
+        self.outfile = self
+
+    def showsyntaxerror(self, filename=None):
+        """Display the syntax error that just occurred.
+
+        This doesn't display a stack trace because there isn't one.
+
+        If a filename is given, it is stuffed in the exception instead
+        of what was there before (because Python's parser always uses
+        "<string>" when reading from a string).
+
+        The output is written by self.write(), below.
+
+        """
+        type, value, sys.last_traceback = sys.exc_info()
+        sys.last_type = type
+        sys.last_value = value
+        if filename and type is SyntaxError:
+            # Work hard to stuff the correct filename in the exception
+            try:
+                msg, (dummy_filename, lineno, offset, line) = value
+            except:
+                # Not the format we expect; leave it alone
+                pass
+            else:
+                # Stuff in the right filename
+                value = SyntaxError(msg, (filename, lineno, offset, line))
+                sys.last_value = value
+        l = traceback.format_exception_only(type, value)
+        tbtext = ''.join(l)
+        lexer = get_lexer_by_name("pytb")
+        self.format(tbtext,lexer)
+
+    def showtraceback(self):
+        """Display the exception that just occurred.
+
+        We remove the first stack item because it is our own code.
+
+
+        """
+        type, value, tb = sys.exc_info()
+        sys.last_type = type
+        sys.last_value = value
+        sys.last_traceback = tb
+        tblist = traceback.extract_tb(tb)
+        del tblist[:1]
+        l = traceback.format_list(tblist)
+        if l:
+            l.insert(0, "Traceback (most recent call last):\n")
+        l[len(l):] = traceback.format_exception_only(type, value)
+        tbtext = ''.join(l)
+        lexer = get_lexer_by_name("pytb", stripall=True)
+
+        self.format(tbtext,lexer)
+        
+
+    def format(self, tbtext, lexer):
+        traceback_informative_formatter = BPythonFormatter(default_colors)
+        traceback_code_formatter = BPythonFormatter({Token: ('d')})
+        tokens= list(lexer.get_tokens(tbtext))
+
+        no_format_mode = False
+        cur_line = []
+        for token, text in tokens:
+            if text.endswith('\n'):
+                cur_line.append((token,text))
+                if no_format_mode:
+                    traceback_code_formatter.format(cur_line, self.outfile)
+                    no_format_mode = False
+                else:
+                    traceback_informative_formatter.format(cur_line, self.outfile)
+                cur_line = []
+            elif text == '    ' and cur_line == []:
+                no_format_mode = True
+                cur_line.append((token,text))
+            else:
+                cur_line.append((token,text))
+        assert cur_line == [], cur_line
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -17,6 +17,7 @@
 from pygments import format
 from pygments.lexers import PythonLexer
 from pygments.formatters import TerminalFormatter
+from interpreter import Interp
 
 import blessings
 
@@ -236,7 +237,8 @@ def __init__(self,
                                                    # would be unsafe because initial
                                                    # state was passed in
         if interp is None:
-            interp = code.InteractiveInterpreter(locals=locals_)
+            interp = Interp(locals=locals_)
+            interp.writetb = self.send_to_stderr
         if banner is None:
             if config.help_key:
                 banner = _('Welcome to bpython!') + ' ' + (_('Press <%s> for help.') % config.help_key)
@@ -867,9 +869,7 @@ def send_to_stderr(self, error):
         lines = error.split('\n')
         if lines[-1]:
             self.current_stdouterr_line += lines[-1]
-        self.display_lines.extend([func_for_letter(self.config.color_scheme['error'])(line)
-                                   for line in sum([paint.display_linize(line, self.width, blank_line=True)
-                                                    for line in lines[:-1]], [])])
+        self.display_lines.extend(sum([paint.display_linize(line, self.width, blank_line=True) for line in lines[:-1]], []))
 
     def send_to_stdin(self, line):
         if line.endswith('\n'):
@@ -1202,6 +1202,7 @@ def reevaluate(self, insert_into_history=False):
 
         if not self.weak_rewind:
             self.interp = self.interp.__class__()
+            self.interp.writetb = self.send_to_stderr
             self.coderunner.interp = self.interp
 
         self.buffer = []
diff --git a/bpython/test/test_interpreter.py b/bpython/test/test_interpreter.py
@@ -0,0 +1,42 @@
+import unittest
+
+from bpython.curtsiesfrontend import interpreter
+from curtsies.fmtfuncs import *
+
+class TestInterpreter(unittest.TestCase):
+    def test_syntaxerror(self):
+        i = interpreter.Interp()
+        a = []
+
+        def append_to_a(message):
+            a.append(message)
+        
+        i.write = append_to_a
+        i.runsource('1.1.1.1')
+
+        expected = ''+u''+u'  File '+green(u'"<input>"')+u', line '+bold(magenta(u'1'))+u'\n'+u'    '+u'1.1'+u'.'+u'1.1'+u'\n'+u'    '+u'    '+u'^'+u'\n'+bold(red(u'SyntaxError'))+u': '+cyan(u'invalid syntax')+u'\n'
+
+        self.assertEquals(str(plain('').join(a)), str(expected))
+        self.assertEquals(plain('').join(a), expected)
+
+    def test_traceback(self):
+        i = interpreter.Interp()
+        a = []
+
+        def append_to_a(message):
+            a.append(message)
+        
+        i.write = append_to_a
+
+        def f():
+            return 1/0
+
+        def g():
+            return f()
+
+        i.runsource('g()')
+
+        expected = u'Traceback (most recent call last):\n'+''+u'  File '+green(u'"<input>"')+u', line '+bold (magenta(u'1'))+u', in '+cyan(u'<module>')+u'\n'+''+bold(red(u'NameError'))+u': '+cyan(u"name 'g' is not defined")+u'\n'
+
+        self.assertEquals(str(plain('').join(a)), str(expected))
+        self.assertEquals(plain('').join(a), expected)
