Skip to content

Instantly share code, notes, and snippets.

@mscolnick
Created May 1, 2025 20:14
Show Gist options
  • Save mscolnick/cea79c6030f8d25318510b182686545e to your computer and use it in GitHub Desktop.
Save mscolnick/cea79c6030f8d25318510b182686545e to your computer and use it in GitHub Desktop.
import marimo
__generated_with = "0.13.3"
app = marimo.App(width="medium")
@app.cell
def _():
import sys
import traceback
import os
import linecache
import inspect
from types import FrameType
from typing import List, Optional, TextIO, Tuple, Any, Dict, Callable, Union
import re
import colorama
from colorama import Fore, Style, Back
# Initialize colorama for cross-platform color support
colorama.init()
class EnhancedTraceback:
"""Enhanced traceback formatter with improved readability and additional context."""
# Syntax highlighting colors
COLORS = {
"keyword": Fore.BLUE,
"builtin": Fore.CYAN,
"string": Fore.GREEN,
"number": Fore.MAGENTA,
"comment": Fore.LIGHTBLACK_EX,
"exception": Fore.RED + Style.BRIGHT,
"filename": Fore.YELLOW,
"lineno": Fore.LIGHTWHITE_EX,
"name": Fore.WHITE,
"reset": Style.RESET_ALL,
"frame": Fore.LIGHTBLACK_EX,
"error": Fore.RED + Style.BRIGHT,
"highlight": Back.LIGHTBLACK_EX,
}
# Regex patterns for syntax highlighting
PATTERNS = {
"keyword": r"\b(and|as|assert|async|await|break|class|continue|def|del|elif|else|except|finally|for|from|global|if|import|in|is|lambda|nonlocal|not|or|pass|raise|return|try|while|with|yield)\b",
"builtin": r"\b(abs|all|any|bin|bool|bytearray|bytes|callable|chr|classmethod|compile|complex|delattr|dict|dir|divmod|enumerate|eval|exec|filter|float|format|frozenset|getattr|globals|hasattr|hash|help|hex|id|input|int|isinstance|issubclass|iter|len|list|locals|map|max|memoryview|min|next|object|oct|open|ord|pow|print|property|range|repr|reversed|round|set|setattr|slice|sorted|staticmethod|str|sum|super|tuple|type|vars|zip|__import__)\b",
"string": r"(\'[^\']*\'|\"[^\"]*\")",
"number": r"\b\d+\b",
"comment": r"#.*$",
}
def __init__(
self,
show_locals: bool = True,
context_lines: int = 3,
max_str_len: int = 100,
use_color: bool = True,
show_full_paths: bool = False,
include_summary: bool = True,
include_suggestions: bool = True,
):
"""
Initialize the enhanced traceback formatter.
Args:
show_locals: Whether to display local variables at each frame
context_lines: Number of context lines to show before and after the error line
max_str_len: Maximum length for displayed string values
use_color: Whether to use colored output
show_full_paths: Whether to show full file paths
include_summary: Whether to include a summary of the error
include_suggestions: Whether to include suggestions for fixing common errors
"""
self.show_locals = show_locals
self.context_lines = context_lines
self.max_str_len = max_str_len
self.use_color = use_color
self.show_full_paths = show_full_paths
self.include_summary = include_summary
self.include_suggestions = include_suggestions
# Compile regex patterns
self.compiled_patterns = {
k: re.compile(v) for k, v in self.PATTERNS.items()
}
def color(self, text: str, color_name: str) -> str:
"""Apply color to text if colors are enabled."""
if self.use_color:
return f"{self.COLORS.get(color_name, '')}{text}{self.COLORS['reset']}"
return text
def syntax_highlight(self, code: str) -> str:
"""Apply syntax highlighting to code."""
if not self.use_color:
return code
# Make a copy of the code to modify
highlighted = code
# Apply patterns in reverse order to avoid messing up the string
for pattern_name in [
"comment",
"string",
"number",
"builtin",
"keyword",
]:
pattern = self.compiled_patterns[pattern_name]
color = self.COLORS[pattern_name]
reset = self.COLORS["reset"]
# Replace each match with colored version
highlighted = pattern.sub(f"{color}\\g<0>{reset}", highlighted)
return highlighted
def format_value(self, value: Any) -> str:
"""Format a value for display, with truncation for long strings."""
try:
if isinstance(value, str):
if len(value) > self.max_str_len:
half = (self.max_str_len - 3) // 2
return f"'{value[:half]}...{value[-half:]}'"
return repr(value)
else:
result = repr(value)
if len(result) > self.max_str_len:
return result[: self.max_str_len - 3] + "..."
return result
except Exception as e:
return f"<Error displaying value: {e}>"
def format_locals(self, local_vars: Dict[str, Any]) -> str:
"""Format local variables for display."""
if not local_vars:
return " No local variables\n"
result = []
for name, value in sorted(local_vars.items()):
if name.startswith("__") and name.endswith("__"):
continue # Skip magic variables
formatted_value = self.format_value(value)
result.append(
f" {self.color(name, 'name')} = {self.syntax_highlight(formatted_value)}"
)
return "\n".join(result) + "\n"
def get_relevant_code(
self, filename: str, lineno: int
) -> Tuple[List[str], int]:
"""Get context lines around the error line."""
try:
lines = []
start = max(1, lineno - self.context_lines)
end = lineno + self.context_lines + 1
# Get the source lines
for i in range(start, end):
line = linecache.getline(filename, i)
if line:
lines.append((i, line.rstrip()))
# Calculate the index of the error line in our list
error_index = lineno - start
return lines, error_index
except Exception:
return [(lineno, "Unable to retrieve source code")], 0
def format_frame(self, frame: FrameType, lineno: int) -> str:
"""Format a single stack frame."""
filename = frame.f_code.co_filename
function = frame.f_code.co_name
# Format the file path
if not self.show_full_paths and not filename.startswith("<"):
display_filename = os.path.basename(filename)
else:
display_filename = filename
# Get the code context
code_lines, error_index = self.get_relevant_code(filename, lineno)
# Format the frame header
header = f" File {self.color(display_filename, 'filename')}, line {self.color(str(lineno), 'lineno')}, in {self.color(function, 'name')}\n"
# Format the code context
code_context = ""
for i, (line_num, line) in enumerate(code_lines):
prefix = "→ " if i == error_index else " "
highlighted_line = self.syntax_highlight(line)
if i == error_index and self.use_color:
code_context += f" {prefix}{self.color(str(line_num).rjust(4), 'lineno')} {self.COLORS['highlight']}{highlighted_line}{self.COLORS['reset']}\n"
else:
code_context += f" {prefix}{self.color(str(line_num).rjust(4), 'lineno')} {highlighted_line}\n"
# Format local variables if requested
locals_str = ""
if self.show_locals:
locals_str = f"\n {self.color('Local variables:', 'frame')}\n{self.format_locals(frame.f_locals)}"
return f"{header}{code_context}{locals_str}"
def get_suggestion(
self, exc_type: type, exc_value: Exception, tb: traceback
) -> str:
"""Generate suggestions for common errors."""
if not self.include_suggestions:
return ""
suggestion = ""
# Handle specific exception types
if exc_type is NameError:
match = re.search(r"name '(\w+)' is not defined", str(exc_value))
if match:
var_name = match.group(1)
suggestion = f"Make sure '{var_name}' is defined before using it. Check for typos in variable names."
elif exc_type is TypeError:
if "takes" in str(exc_value) and "arguments" in str(exc_value):
suggestion = "Check the number of arguments you're passing to the function."
elif "object is not subscriptable" in str(exc_value):
suggestion = "You're trying to use [] on an object that doesn't support indexing."
elif "object is not callable" in str(exc_value):
suggestion = (
"You're trying to call an object that is not a function."
)
elif exc_type is AttributeError:
match = re.search(
r"'(\w+)' object has no attribute '(\w+)'", str(exc_value)
)
if match:
obj_type, attr = match.group(1), match.group(2)
suggestion = f"The '{obj_type}' type doesn't have a '{attr}' attribute. Check for typos or make sure you're using the right object."
elif exc_type is KeyError:
suggestion = "The key you're trying to access doesn't exist in the dictionary."
elif exc_type is IndexError:
suggestion = "You're trying to access an index that is out of range. Check the length of your sequence."
elif exc_type is ModuleNotFoundError:
match = re.search(r"No module named '(\w+)'", str(exc_value))
if match:
module = match.group(1)
suggestion = f"The module '{module}' is not installed. Try installing it with 'pip install {module}'."
elif exc_type is IndentationError:
suggestion = "Check your code indentation. Python uses indentation to determine code blocks."
elif exc_type is SyntaxError:
suggestion = "There's a syntax error in your code. Check for missing parentheses, quotes, or colons."
elif exc_type is ZeroDivisionError:
suggestion = "You're trying to divide by zero. Check your divisor and add a condition to handle zero values."
if suggestion:
return f"\n{self.color('Suggestion:', 'name')} {suggestion}\n"
return ""
def format_exception(
self, exc_type: type, exc_value: Exception, tb: traceback
) -> str:
"""Format an exception with traceback."""
frames = []
# Collect all frames
current_tb = tb
while current_tb:
frames.append((current_tb.tb_frame, current_tb.tb_lineno))
current_tb = current_tb.tb_next
# Format each frame
formatted_frames = []
for i, (frame, lineno) in enumerate(frames):
formatted_frames.append(self.format_frame(frame, lineno))
# Format the exception
exc_name = exc_type.__name__
exc_msg = str(exc_value)
# Create the traceback header
header = (
f"{self.color('Traceback (most recent call last):', 'frame')}\n"
)
# Format the exception line
exception_line = f"{self.color(exc_name, 'exception')}: {exc_msg}\n"
# Add suggestions if enabled
suggestion = self.get_suggestion(exc_type, exc_value, tb)
# Add summary if enabled
summary = ""
if self.include_summary:
if frames:
last_frame, last_lineno = frames[-1]
filename = last_frame.f_code.co_filename
if not self.show_full_paths:
filename = os.path.basename(filename)
function = last_frame.f_code.co_name
summary = f"\n{self.color('Error Summary:', 'error')} {self.color(exc_name, 'exception')} occurred in {self.color(function, 'name')} at {self.color(filename, 'filename')}:{self.color(str(last_lineno), 'lineno')}\n"
# Combine all parts
return f"{header}{''.join(formatted_frames)}{exception_line}{summary}{suggestion}"
def print_exception(
self,
exc_type: type,
exc_value: Exception,
tb: traceback,
file: TextIO = None,
) -> None:
"""Print an exception with traceback to a file."""
if file is None:
file = sys.stderr
file.write(self.format_exception(exc_type, exc_value, tb))
def print_exc(self, file: TextIO = None) -> None:
"""Print the current exception with traceback."""
exc_type, exc_value, tb = sys.exc_info()
self.print_exception(exc_type, exc_value, tb, file)
def format_exc(self) -> str:
"""Format the current exception with traceback."""
exc_type, exc_value, tb = sys.exc_info()
return self.format_exception(exc_type, exc_value, tb)
# Function to install as the default exception hook
def install(
show_locals=True,
context_lines=3,
use_color=True,
show_full_paths=False,
include_summary=True,
include_suggestions=True,
max_str_len=100,
):
"""
Install the enhanced traceback formatter as the default exception hook.
Args:
show_locals: Whether to display local variables at each frame
context_lines: Number of context lines to show before and after the error line
use_color: Whether to use colored output
show_full_paths: Whether to show full file paths
include_summary: Whether to include a summary of the error
include_suggestions: Whether to include suggestions for fixing common errors
max_str_len: Maximum length for displayed string values
"""
formatter = EnhancedTraceback(
show_locals=show_locals,
context_lines=context_lines,
use_color=use_color,
show_full_paths=show_full_paths,
include_summary=include_summary,
include_suggestions=include_suggestions,
max_str_len=max_str_len,
)
def excepthook(exc_type, exc_value, tb):
formatter.print_exception(exc_type, exc_value, tb)
sys.excepthook = excepthook
# Also patch IPython if it's being used
try:
import IPython
ip = IPython.get_ipython()
if ip:
def ipython_excepthook(
self, exc_type, exc_value, tb, tb_offset=None
):
formatter.print_exception(exc_type, exc_value, tb)
ip.set_custom_exc((Exception,), ipython_excepthook)
except (ImportError, AttributeError):
pass
return formatter
return (EnhancedTraceback,)
@app.cell
def _(EnhancedTraceback):
formatter = EnhancedTraceback(
show_locals=False,
context_lines=3,
use_color=False,
show_full_paths=False,
include_summary=True,
include_suggestions=True,
max_str_len=100,
)
try:
# Your code that might raise an exception
result = 1 / 0
except Exception:
print(formatter.format_exc())
return
if __name__ == "__main__":
app.run()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment