Parse Python Stack Trace and Open Selected Source References for Editing in OS X

#! /usr/bin/env python
 
import subprocess
import StringIO
import re
import os
import sys
 
from optparse import OptionGroup
from optparse import OptionParser
 
class StackDesc(object):
    def __init__(self, items):
        self.filepath = os.path.abspath(items[0])
        self.line_num = int(items[1])
        if len(items) >= 4 and items[3] is not None:
            self.object_name = items[3]
        else:
            self.object_name = ""
        self.statement = None
 
def scrape_terminal_history():
    script = """\
/usr/bin/osascript -s o -e   '
tell application "Terminal"
activate
    tell front window
        return (the history of selected tab)
    end tell
end tell
'
"""
    p = subprocess.Popen(script, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()
    return stdout, stderr
 
def scrape_traceback_lines(most_recent_only=False):
    stdout, stderr = scrape_terminal_history()
    stdout = stdout.split("\n")[::-1]
    if most_recent_only:
        for idx, f in enumerate(stdout):
            if f.startswith("Traceback (most recent call last)"):
                tb = (stdout[0:idx])[::-1]
                return tb
    else:
        return stdout[::-1]
    return []
 
def parse_traceback(traceback_lines):
    pattern = re.compile("""^\s+File \"(.*)\", line (\d+)(, in (.*))*\s*""")
    if traceback_lines is None or len(traceback_lines) == 0:
        return [], ""
    stack_descs = []
    error_str = None
    prev_line_is_stack_frame = False
    for tb in traceback_lines:
        m = pattern.match(tb)
        if m is not None:
            items = m.groups()
            stack_descs.append(StackDesc(items))
            prev_line_is_stack_frame = True
        elif prev_line_is_stack_frame:
            stack_descs[-1].statement = tb.strip()
            prev_line_is_stack_frame = False
        elif error_str is None and len(stack_descs) > 0:
            tbs = tb.strip()
            if tbs and tbs != "^":
                error_str = tbs
            elif tbs == "^":
                error_str = tb
            prev_line_is_stack_frame = False
        elif (error_str is not None and error_str.strip() == "^"):
            error_str = error_str + "\n" + tb.strip()
            prev_line_is_stack_frame = False
    return stack_descs, error_str
 
def get_traceback_stack(most_recent_only=False):
    traceback_lines = scrape_traceback_lines(most_recent_only)
    stack_descs, error_str = parse_traceback(traceback_lines)
    return stack_descs, error_str
 
def compose_ansi_color(code):
    return chr(27) + '[' + code + 'm'
 
def display_stack_descs(stack_descs, error_str, color):
    if color:
        title_color = compose_ansi_color("1;91")
        error_color = compose_ansi_color("1;91")
        index_color = compose_ansi_color("0;37;40")
        filepath_color = compose_ansi_color("1;94")
        line_num_color = compose_ansi_color("1;94")
        obj_color = compose_ansi_color("1;34")
        statement_color = compose_ansi_color("0;31")
        clear_color = compose_ansi_color("0")
    else:
        title_color = ""
        error_color = ""
        index_color = ""
        filepath_color = ""
        line_num_color = ""
        obj_color = ""
        statement_color = ""
        clear_color = ""
    sys.stdout.write("%sPython stack trace (most recent call last):%s\n" % (title_color, clear_color))
    for sdi, sd in enumerate(stack_descs):
        if sd.object_name:
            obj_name = ": %s%s%s" % (obj_color, sd.object_name, clear_color)
        else:
            obj_name = ""
        sys.stdout.write("%s% 3d %s: %s%s%s [%s%d%s]%s\n" % (
                index_color, sdi, clear_color,
                filepath_color, sd.filepath, clear_color,
                line_num_color, sd.line_num, clear_color,
                obj_name))
        if sd.statement:
            sys.stdout.write("      >>> %s%s%s\n" % (statement_color, sd.statement, clear_color))
    if error_str:
        sys.stdout.write("%s%s%s\n" % (error_color, error_str, clear_color))
 
 
def bbedit_browse(stack_descs, error_str):
    script_template = """
/usr/bin/osascript 2>/dev/null <<EOF
set locationData to {%s}
tell application "BBEdit"
    activate
	set resultItems to {}
    set numLocations to count locationData
	repeat with locationIndex from 1 to numLocations
		set {locationMessage, locationFile, locationLine} to item locationIndex of locationData
        set locationLine to locationLine as number
        if locationIndex equals numLocations then
            set rKind to error_kind
        else
            set rKind to note_kind
        end if
        set resultEntry to {message:locationMessage, result_kind:rKind, result_file:locationFile , result_line:locationLine, message_script:0}
        copy resultEntry to end of resultItems
	end repeat
	make new results browser with data resultItems with properties {name:"Python Traceback"}
end tell
 
on |splittext|(delimiter, someText)
	set prevTIDs to AppleScript's text item delimiters
	set AppleScript's text item delimiters to delimiter
	set output to text items of someText
	set AppleScript's text item delimiters to prevTIDs
	return output
end |splittext|
EOF
"""
    list_items = []
    for sd in stack_descs:
        if sd.statement is None:
            message = ""
        else:
            message = sd.statement.replace('"',r'\"')
        list_items.append('{"%s", "%s", "%s"}' % (message, sd.filepath, sd.line_num))
    list_items_text = ",".join(list_items)
    script = script_template % list_items_text
    p = subprocess.Popen(script, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
_prog_usage = '%prog [options] [ID# [ID# [ID# [...]]]]'
_prog_version = 'Python Stack Trace Parser Version 1.0'
_prog_description = """\
Parses last Python stack trace in current Terminal history and opens most
recently referenced source file for editing at the correct line. If ID#
is given, then the file referenced in the stack frame indexed by ID#
will be opened (the most recent call will be indexed by 0, the previous
stack frame will be indexed by 1, etc. all the way to the first stack
frame). If ID# is preceded by a '^', then indexed starts at the first
stack frame (i.e., the first stack frame will be "^1", the
second stack frame is "^2" etc.).
"""
 
def main():
    """
    Main CLI handler.
    """
 
    parser = OptionParser(usage=_prog_usage,
        add_help_option=True,
        version=_prog_version,
        description=_prog_description)
 
    parser.add_option('-r', '--recent-only',
        action='store_true',
        dest='most_recent_only',
        default=False,
        help='only open most recent stack track in Terminal history (default is all)')
 
    parser.add_option('-e', '--execute',
        action='store',
        dest='execute',
        metavar='<PYTHON>',
        default=None,
        help="Python traceback will be read from results of <PYTHON> (e.g. 'pystack.py -e test.py', 'pystack.py -e '-c print(\"hello\")'")
 
    parser.add_option('--stdin',
        action='store_true',
        dest='from_stdin',
        default=False,
        help="Python traceback will be read from standard input (e.g. 'python test.py 2>&1 | pystack.py --stdin'")
 
    parser.add_option('-l', '--list',
        action='store_true',
        dest='show_list',
        default=False,
        help='list parsed stack track')
 
    parser.add_option('-b', '--browse',
        action='store_true',
        dest='bbedit_browse',
        default=False,
        help='browse list in BBEdit')
 
    parser.add_option('-a', '--open-all',
        action='store_true',
        dest='open_all_files',
        default=False,
        help='open all files referenced in stack trace')
 
    parser.add_option('-i', '--interactive',
        action='store_true',
        dest='interactive',
        default=False,
        help='list parsed stack track and prompt for file to edit')
 
    parser.add_option('-c', '--color',
        action='store_true',
        dest='color',
        default=False,
        help='display in color')
 
    parser.add_option('--editor',
        action='store',
        dest='editor_app',
        default='bbedit',
        help="path to editor application (default = '%default')")
 
    parser.add_option('-o', '--editor-opts',
        dest='editor_opts',
        help='options to be passed to editor')
 
    parser.add_option('-n', '--new-window',
        action='store_const',
        dest='editor_new_window',
        const='--new-window',
        help='open new text editor window (BBEdit only)')
 
    parser.add_option('-f', '--front-window',
        action='store_const',
        dest='editor_front_window',
        const='--front-window',
        help='open in front text editor window (BBEdit only)')
 
    parser.add_option('-q', '--quiet',
        action='store_true',
        dest='quiet',
        help='suppress messages and standard output when executing python')
 
    (opts, args) = parser.parse_args()
 
    if opts.execute:
        pycmd = "python " + opts.execute
        p = subprocess.Popen(pycmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        p.wait()
        stdout_lines = p.stdout.readlines()
        stderr_lines = p.stderr.readlines()
        if not opts.quiet:
            sys.stdout.write("\n".join(stdout_lines))
        stack_descs, error_str = parse_traceback(stderr_lines)
    elif opts.from_stdin:
        traceback_lines = sys.stdin.readlines()
        stack_descs, error_str = parse_traceback(traceback_lines)
    else:
        stack_descs, error_str = get_traceback_stack(most_recent_only=opts.most_recent_only)
 
    if len(stack_descs) == 0:
        sys.exit("Python stack trace not found.")
 
    if opts.show_list:
        display_stack_descs(stack_descs, error_str, opts.color)
        sys.exit(0)
    if opts.bbedit_browse:
        bbedit_browse(stack_descs, error_str)
        sys.exit(0)
 
    if opts.editor_opts is not None:
        editor_opts = opts.editor_opts
    else:
        editor_opts = ""
    if opts.editor_new_window:
        editor_opts = "--new-window " + editor_opts
    if opts.editor_front_window:
        editor_opts = "--front-window " + editor_opts
    editor_invocation = "%s %s" % (opts.editor_app, editor_opts)
    if opts.interactive:
        display_stack_descs(stack_descs, error_str, color=opts.color)
        idx = None
        while idx is None:
            idx = raw_input("Stack frame reference to edit ('q' to quit): ")
            if idx == "" or idx.lower() == "q":
                sys.exit(0)
            try:
                idx = int(idx)
            except ValueError, e:
                sys.stderr.write("'%s' is an invalid selection.\n" % idx)
                idx = None
            if idx >= len(stack_descs):
                sys.stderr.write("Index %d is out of bounds [0, %d]" % (idx, len(stack_descs)-1))
                idx = None
            if idx is not None:
                sd = stack_descs[idx]
                command = "%s +%d %s" % (editor_invocation, sd.line_num, sd.filepath)
                edp = subprocess.Popen(command, shell=True)
                idx = None
    else:
        indexes = []
        if opts.open_all_files:
            indexes = range(0, len(stack_descs)-1)
        elif len(args) == 0:
            indexes = [-1]
        else:
            for i in args:
                if i.startswith("^"):
                    i = -1 * int(i[1:])
                else:
                    i = int(i)
                if i > len(stack_descs):
                    sys.exit("Index %d is out of bounds [0, %d]" % (i, len(stack_descs)-1))
                indexes.append(i)
        editor_args = []
        for i in indexes:
            editor_args.append("+%d %s" % (stack_descs[i].line_num, stack_descs[i].filepath))
        editor_args = " ".join(editor_args)
        command = "%s %s" % (editor_invocation, editor_args)
        edp = subprocess.Popen(command, shell=True)
 
if __name__ == '__main__':
    main()