135 lines
4.1 KiB
Python
Executable File
135 lines
4.1 KiB
Python
Executable File
#! /usr/bin/env python
|
|
import argparse
|
|
import sys
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from rich.console import Console
|
|
from rich.syntax import Syntax
|
|
from datetime import datetime
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
LINE_NUMBER_RE = re.compile(r", line \d+,")
|
|
OUTPUT_ROOT = Path("~/Pictures/pytb").expanduser()
|
|
|
|
|
|
def main() -> int:
|
|
# Handle program arguments
|
|
ap = argparse.ArgumentParser(
|
|
prog="pytb", description="Tool for analyzing Python back traces"
|
|
)
|
|
ap.add_argument("--file", "-f", help="Read from file instead of stdin", type=Path)
|
|
ap.add_argument(
|
|
"--no-strip-referer", help="Strip referer from flask tbs", action="store_true"
|
|
)
|
|
ap.add_argument("--trace-only", help="Only print the trace", action="store_true")
|
|
ap.add_argument(
|
|
"-v", "--verbose", help="Enable verbose logging", action="store_true"
|
|
)
|
|
args = ap.parse_args()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.DEBUG if args.verbose else logging.INFO,
|
|
format="%(levelname)s: %(message)s",
|
|
)
|
|
|
|
# Attempt to read from file
|
|
if args.file:
|
|
tb_lines = args.file.read_text().splitlines()
|
|
else:
|
|
# Check if the shell is interactive
|
|
if sys.stdin.isatty():
|
|
print("Please paste the backtrace and press Ctrl+D:")
|
|
|
|
# Read from stdin until EOF
|
|
try:
|
|
tb_lines = "".join(list(sys.stdin)).splitlines()
|
|
except KeyboardInterrupt:
|
|
print("\nKeyboard interrupt detected, exiting...")
|
|
return 1
|
|
|
|
# Seek to the first line of the backtrace
|
|
for start_idx, line in enumerate(tb_lines):
|
|
if line.startswith("Traceback"):
|
|
break
|
|
else:
|
|
logger.error("No traceback found")
|
|
return 1
|
|
|
|
# Group the traceback lines into frames
|
|
frames = []
|
|
is_in_frame = False
|
|
for line in tb_lines[start_idx:]:
|
|
if line.lstrip().startswith("File "):
|
|
is_in_frame = True
|
|
frames.append([line])
|
|
elif is_in_frame:
|
|
frames[-1].append(line)
|
|
|
|
# Handle the frames
|
|
output_lines = []
|
|
for frame in frames:
|
|
# Figure out the file
|
|
file = Path(frame[0].split('"')[1])
|
|
line_num = int(LINE_NUMBER_RE.search(frame[0]).group(0)[6:-1])
|
|
|
|
# Print the actual code
|
|
for idx, statement in enumerate(frame[1:]):
|
|
# Remove left padding
|
|
statement = statement.lstrip()
|
|
|
|
# Remove referer if needed
|
|
if not args.no_strip_referer:
|
|
statement = statement.split(", referer")[0]
|
|
|
|
# Build a context string if needed
|
|
context = f" # {file}#{line_num}" if idx == 0 else ""
|
|
|
|
# Print the line
|
|
output_lines.append((statement, context))
|
|
|
|
# Figure out the longest statement
|
|
longest_statement = max(len(line[0]) for line in output_lines[:-1])
|
|
|
|
# Build the lines, padding the statements so that the files line up
|
|
output = ""
|
|
for statement, context in output_lines:
|
|
output += f"{statement.ljust(longest_statement)}{context}\n"
|
|
|
|
# remove any trailing newlines
|
|
output = output.rstrip()
|
|
|
|
# Figure out the longest line
|
|
output_trace = "\n".join(output.splitlines()[:-1])
|
|
output_error = output.splitlines()[-1]
|
|
if args.trace_only:
|
|
longest_line = max(len(line) for line in output_trace.splitlines())
|
|
else:
|
|
longest_line = max(len(line) for line in output.splitlines())
|
|
|
|
# Pass over to rich to do the syntax highlighting
|
|
console = Console(record=True, width=longest_line + 1)
|
|
console.print(Syntax(output_trace, "python", background_color="default"))
|
|
if not args.trace_only:
|
|
console.print(
|
|
Syntax(output_error, "python", background_color="default")
|
|
)
|
|
|
|
# Export an image
|
|
file_name = f"Traceback {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
|
|
if args.file:
|
|
file_name += f" ({args.file.stem})"
|
|
file_name += ".svg"
|
|
OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
|
|
console.save_svg(
|
|
OUTPUT_ROOT / file_name, title="Evan's Python Traceback Visualizer"
|
|
)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|