Syntax highlighting for Vyper exceptions

brownie/_cli/console.py

@@ -14,8 +14,6 @@
 from prompt_toolkit.keys import Keys
 from prompt_toolkit.lexers import PygmentsLexer
 from prompt_toolkit.styles.pygments import style_from_pygments_cls
-from pygments import highlight
-from pygments.formatters import get_formatter_by_name
 from pygments.lexers import PythonLexer
 from pygments.styles import get_style_by_name
 
@@ -109,20 +107,6 @@ def __init__(self, project=None, extra_locals=None):
         if extra_locals:
             locals_dict.update(extra_locals)
 
-        # prepare lexer and formatter
-        self.lexer = PythonLexer()
-        fmt_name = "terminal"
-        try:
-            import curses
-
-            curses.setupterm()
-            if curses.tigetnum("colors") == 256:
-                fmt_name = "terminal256"
-        except Exception:
-            # if curses won't import we are probably using Windows
-            pass
-        self.formatter = get_formatter_by_name(fmt_name, style=console_settings["color_style"])
-
         # create prompt session object
         history_file = str(_get_data_folder().joinpath(".history").absolute())
         kwargs = {}
@@ -184,7 +168,7 @@ def _console_write(self, obj):
         except (SyntaxError, NameError):
             pass
         if CONFIG.settings["console"]["show_colors"]:
-            text = highlight(text, self.lexer, self.formatter)
+            text = color.highlight(text)
         self.write(text)
 
     def raw_input(self, prompt=""):

brownie/utils/color.py

@@ -5,13 +5,31 @@
 from pathlib import Path
 from typing import Dict, Optional, Sequence
 
+import pygments
+from pygments.formatters import get_formatter_by_name
+from pygments.lexers import PythonLexer
+from vyper.exceptions import VyperException
+
 from brownie._config import CONFIG
 
 if sys.platform == "win32":
     import colorama
 
     colorama.init()
 
+fmt_name = "terminal"
+try:
+    import curses
+
+    curses.setupterm()
+    if curses.tigetnum("colors") == 256:
+        fmt_name = "terminal256"
+except Exception:
+    # if curses won't import we are probably using Windows
+    pass
+
+formatter = get_formatter_by_name(fmt_name, style=CONFIG.settings["console"]["color_style"])
+
 
 BASE = "\x1b[0;"
 
@@ -106,13 +124,15 @@ def format_tb(
     ) -> str:
         if isinstance(exc, SyntaxError) and exc.text is not None:
             return self.format_syntaxerror(exc)
+
         tb = [i.replace("./", "") for i in traceback.format_tb(exc.__traceback__)]
         if filename and not CONFIG.argv["tb"]:
             try:
                 start = tb.index(next(i for i in tb if filename in i))
                 stop = tb.index(next(i for i in tb[::-1] if filename in i)) + 1
             except Exception:
                 pass
+
         tb = tb[start:stop]
         for i in range(len(tb)):
             info, code = tb[i].split("\n")[:2]
@@ -127,7 +147,15 @@ def format_tb(
             )
             if code:
                 tb[i] += f"\n{code}"
-        tb.append(f"{self('bright red')}{type(exc).__name__}{self}: {exc}")
+
+        msg = str(exc)
+        if isinstance(exc, VyperException):
+            # apply syntax highlight and remove traceback on vyper exceptions
+            msg = self.highlight(msg)
+            if not CONFIG.argv["tb"]:
+                tb.clear()
+
+        tb.append(f"{self('bright red')}{type(exc).__name__}{self}: {msg}")
         return "\n".join(tb)
 
     def format_syntaxerror(self, exc: SyntaxError) -> str:
@@ -140,6 +168,12 @@ def format_syntaxerror(self, exc: SyntaxError) -> str:
             f"{' '*offset}^\n{self('bright red')}SyntaxError{self}: {exc.msg}"
         )
 
+    def highlight(self, text):
+        """
+        Apply python syntax highlighting to a string.
+        """
+        return pygments.highlight(text, PythonLexer(), formatter)
+
 
 def notify(type_, msg):
     """Prepends a message with a colored tag and outputs it to the console."""