Colorize last_error and display variables along the stack trace

bliss/shell/cli/repl.py

@@ -24,9 +24,12 @@ import gevent
 import signal
 import logging
 import platform
+import importlib
 from collections import deque
+from colorama import Fore, Style
 from datetime import datetime
 
+
 import ptpython.layout
 from prompt_toolkit.patch_stdout import patch_stdout as patch_stdout_context
 from prompt_toolkit.output import DummyOutput
@@ -48,7 +51,7 @@ from bliss.shell.cli.no_thread_repl import NoThreadPythonRepl
 from bliss.shell import standard
 
 from bliss import set_bliss_shell_mode
-from bliss.common.utils import ShellStr, Singleton
+from bliss.common.utils import Singleton
 from bliss.common import constants
 from bliss.common import session as session_mdl
 from bliss.common.session import DefaultSession
@@ -69,6 +72,7 @@ if is_windows():
                 raise AttributeError(prop)
             return ""
 
+
 else:
     from blessings import Terminal
 
@@ -81,28 +85,70 @@ session_mdl.set_current_session = functools.partial(
 # =================== ERROR REPORTING ============================
 
 
-class LastError:
-    def __init__(self):
-        self.errors = deque()
+class PrettyTraceback:
+    _blacklist = [
+        # gevent & aiogevent module paths
+        os.path.dirname(importlib.util.find_spec("gevent").origin),
+        os.path.dirname(importlib.util.find_spec("aiogevent").origin),
+        # gevent compiled functions root path
+        "src/gevent",
+    ]
 
-    def __getitem__(self, index):
-        try:
-            return ShellStr(self.errors[index])
-        except IndexError:
-            return ShellStr(
-                f"No exception with index {index} found, size is {len(self.errors)}"
-            )
+    def __init__(self, exc_type, exc_value, tb):
+        self._timestamp = datetime.now()
+        self._exc_type = exc_type
+        self._exc_value = exc_value
 
-    def __repr__(self):
-        try:
-            return ShellStr(self.errors[-1])
-        except IndexError:
-            return "None"
+        # convert traceback to StackSummary to stringify references and avoid memory leak
+        self._stack_summary = traceback.StackSummary.extract(
+            traceback.walk_tb(tb), lookup_lines=True, capture_locals=True
+        )
 
-    def append(self, item):
-        self.errors.append(item)
-        while len(self.errors) > 100:
-            self.errors.popleft()
+    def _is_blacklisted(self, filename):
+        for black_path in PrettyTraceback._blacklist:
+            if filename.startswith(black_path):
+                return True
+        return False
+
+    def format_color(self, raw=True):
+        fmt_error = f"{Fore.YELLOW}--- {self._timestamp.strftime('%d/%m/%Y %H:%M:%S')} ---{Style.RESET_ALL}\n"
+
+        # Stack trace formatting
+        fmt_stack = ""
+        for f in self._stack_summary:
+            if not raw and self._is_blacklisted(f.filename):
+                continue
+
+            # skip bottom calls from ptpython by resetting fmt_stack
+            if f.filename == "<stdin>":
+                fmt_stack = ""
+
+            fmt_stack += f"{Fore.CYAN}{f.filename}{Fore.GREEN}:{f.lineno} "
+            fmt_stack += f"{Fore.BLUE}{f.name}{Style.RESET_ALL}\n"
+
+            if f._line:
+                fmt_stack += f"{Fore.RED}  > {Style.RESET_ALL}{f._line}\n\n"
+                # locals are not printed when f._line is empty, this happens with traces coming from
+                # the interpreter or compiled files and locals dict is full of obscure symbols.
+                if f.locals is not None:
+                    for key, val in f.locals.items():
+                        fmt_stack += f"{Fore.MAGENTA}{key}: "
+                        fmt_stack += (
+                            f"{Style.BRIGHT}{Fore.WHITE}{val}{Style.RESET_ALL}\n"
+                        )
+                    fmt_stack += "\n"
+            else:
+                fmt_stack += "\n"
+        fmt_error += fmt_stack
+
+        # Exception formatting
+        fmt_error += f"{Style.BRIGHT}{Fore.RED}{self._exc_type.__name__}:\n"
+        fmt_error += Style.BRIGHT + Fore.YELLOW
+        for arg in self._exc_value.args:
+            fmt_error += f"{arg}\n"
+        fmt_error += Style.RESET_ALL
+
+        return fmt_error
 
 
 class ErrorReport(ErrorReportInterface):
@@ -121,11 +167,11 @@ class ErrorReport(ErrorReportInterface):
 
     def __init__(self):
         self._expert_mode = False
-        self._last_error = LastError()
+        self._history = deque(maxlen=100)
 
     @property
-    def last_error(self):
-        return self._last_error
+    def history(self):
+        return self._history
 
     @property
     def expert_mode(self):
@@ -150,33 +196,21 @@ def install_excepthook():
             return
         err_file = sys.stderr
 
-        # Store latest traceback (as a string to avoid memory leaks)
-        # next lines are inspired from "_handle_exception()" (ptpython/repl.py)
-        # skip bottom calls from ptpython
-        tblist = list(traceback.extract_tb(tb))
-        to_remove = 0
-        for line_nr, tb_tuple in enumerate(tblist):
-            if tb_tuple.filename == "<stdin>":
-                to_remove = line_nr
-        for i in range(to_remove):
-            tb = tb.tb_next
+        error_report.history.append(PrettyTraceback(exc_type, exc_value, tb))
 
+        # python generic formatting for the exception logger
         exc_text = "".join(traceback.format_exception(exc_type, exc_value, tb))
-        error_report._last_error.append(
-            datetime.now().strftime("%d/%m/%Y %H:%M:%S ") + exc_text
-        )
-
         exc_logger.error(exc_text)
 
         # Adapt the error message depending on the expert_mode
         if error_report._expert_mode:
-            print(error_report._last_error, file=err_file)
+            print(error_report.history[-1].format_color(True), file=err_file)
         elif current_session:
             if current_session.is_loading_config:
                 print(f"{exc_type.__name__}: {exc_value}", file=err_file)
             else:
                 print(
-                    f"!!! === {exc_type.__name__}: {exc_value} === !!! ( for more details type cmd 'last_error' )",
+                    f"!!! === {exc_type.__name__}: {exc_value} === !!! ( for more details type cmd 'last_error()' )",
                     file=err_file,
                 )
 
@@ -220,7 +254,10 @@ class Info:
         self.info_repr = info(obj_with_info)
 
     def __repr__(self):
-        return self.info_repr
+        try:
+            return self.info_repr
+        except AttributeError:
+            return super().__repr__()
 
 
 class WrappedStdout:
@@ -617,7 +654,7 @@ def initialize(
 
 
 def _archive_history(
-    history_filename, file_size_thresh=10**6, keep_active_entries=1000
+    history_filename, file_size_thresh=10 ** 6, keep_active_entries=1000
 ):
     if (
         os.path.exists(history_filename)
@@ -711,9 +748,20 @@ def cli(
                 if alias in protected_user_ns:
                     protected_user_ns._protect(alias)
 
+    def last_error(index=None):
+        hist = user_ns["ERROR_REPORT"].history
+        try:
+            idx = -1 if index is None else index
+            print(hist[idx].format_color(user_ns["ERROR_REPORT"].expert_mode))
+        except IndexError:
+            if index is None:
+                print("None")
+            else:
+                print(f"No exception with index {index} found, size is {len(hist)}")
+
     # handle the last error report
     # (in the shell env only)
-    user_ns["last_error"] = user_ns["ERROR_REPORT"].last_error
+    user_ns["last_error"] = last_error
 
     # protect certain imports and Globals
     to_protect = [