backtrace.py

importos
importsys
importargparse
importtraceback

importcolorama
fromcoloramaimportFore, Style


DESCRIPTION="""Beautify Tracebacks.

Just pipe stderr into backtrace like so:
 `python bad-program.py 2>&1 | backtrace`
"""

TRACEBACK_IDENTIFIER='Traceback (most recent call last):\n'
STYLES= {
'backtrace': Fore.YELLOW+'{0}',
'error': Fore.RED+Style.BRIGHT+'{0}',
'line': Fore.RED+Style.BRIGHT+'{0}',
'module': '{0}',
'context': Style.BRIGHT+Fore.GREEN+'{0}',
'call': Fore.YELLOW+'--> '+Style.BRIGHT+'{0}',
}

CONVERVATIVE_STYLES= {
'backtrace': Fore.YELLOW+'{0}',
'error': Fore.RED+Style.BRIGHT+'{0}',
'line': 'line '+Fore.RED+Style.BRIGHT+'{0},',
'module': 'File {0},',
'context': 'in '+Style.BRIGHT+Fore.GREEN+'{0}',
'call': Fore.YELLOW+'--> '+Style.BRIGHT+'{0}',
}


def_flush(message):
sys.stderr.write(message+'\n')
sys.stderr.flush()


class_Hook(object):
def__init__(self,
entries,
align=False,
strip_path=False,
conservative=False):
self.entries=entries
self.align=align
self.strip=strip_path
self.conservative=conservative

defreverse(self):
self.entries=self.entries[::-1]

defrebuild_entry(self, entry, styles):
entry=list(entry)
# This is the file path.
entry[0] =os.path.basename(entry[0]) ifself.stripelseentry[0]
# Always an int (entry line number)
entry[1] =str(entry[1])

new_entry= [
styles['line'].format(entry[1]) +Style.RESET_ALL,
styles['module'].format(entry[0]) +Style.RESET_ALL,
styles['context'].format(entry[2]) +Style.RESET_ALL,
styles['call'].format(entry[3]) +Style.RESET_ALL
 ]
ifself.conservative:
new_entry[0], new_entry[1] =new_entry[1], new_entry[0]

returnnew_entry

@staticmethod
defalign_all(entries):
lengths= [0, 0, 0, 0]

forentryinentries:
forindex, fieldinenumerate(entry):
lengths[index] =max(lengths[index], len(str(field)))
returnlengths

@staticmethod
defalign_entry(entry, lengths):
return' '.join(
 ['{0:{1}}'.format(field, lengths[index])
forindex, fieldinenumerate(entry)])

defgenerate_backtrace(self, styles):
"""Return the (potentially) aligned, rebuit traceback

 Yes, we iterate over the entries thrice. We sacrifice
 performance for code readability. I mean.. come on, how long can
 your traceback be that it matters?
 """
backtrace= []
forentryinself.entries:
backtrace.append(self.rebuild_entry(entry, styles))

# Get the lenght of the longest string for each field of an entry
lengths=self.align_all(backtrace) ifself.alignelse [1, 1, 1, 1]

aligned_backtrace= []
forentryinbacktrace:
aligned_backtrace.append(self.align_entry(entry, lengths))
returnaligned_backtrace


defhook(reverse=False,
align=False,
strip_path=False,
enable_on_envvar_only=False,
on_tty=False,
conservative=False,
styles=None,
tb=None,
tpe=None,
value=None):
"""Hook the current excepthook to the backtrace.

 If `align` is True, all parts (line numbers, file names, etc..) will be
 aligned to the left according to the longest entry.

 If `strip_path` is True, only the file name will be shown, not its full
 path.

 If `enable_on_envvar_only` is True, only if the environment variable
 `ENABLE_BACKTRACE` is set, backtrace will be activated.

 If `on_tty` is True, backtrace will be activated only if you're running
 in a readl terminal (i.e. not piped, redirected, etc..).

 If `convervative` is True, the traceback will have more seemingly original
 style (There will be no alignment by default, 'File', 'line' and 'in'
 prefixes and will ignore any styling provided by the user.)

 See https://github.com/nir0s/backtrace/blob/master/README.md for
 information on `styles`.
 """
ifenable_on_envvar_onlyand'ENABLE_BACKTRACE'notinos.environ:
return

isatty=getattr(sys.stderr, 'isatty', lambda: False)
ifon_ttyandnotisatty():
return

ifconservative:
styles=CONVERVATIVE_STYLES
align=alignorFalse
elifstyles:
forkinSTYLES.keys():
styles[k] =styles.get(k, STYLES[k])
else:
styles=STYLES

# For Windows
colorama.init()

defbacktrace_excepthook(tpe, value, tb=None):
# Don't know if we're getting traceback or traceback entries.
# We'll try to parse a traceback object.
try:
traceback_entries=traceback.extract_tb(tb)
exceptAttributeError:
traceback_entries=tb
parser=_Hook(traceback_entries, align, strip_path, conservative)

tpe=tpeifisinstance(tpe, str) elsetpe.__name__
tb_message=styles['backtrace'].format('Traceback ({0}):'.format(
'Most recent call '+ ('first'ifreverseelse'last'))) + \
Style.RESET_ALL
err_message=styles['error'].format(tpe+': '+str(value)) + \
Style.RESET_ALL

ifreverse:
parser.reverse()

_flush(tb_message)
backtrace=parser.generate_backtrace(styles)
backtrace.insert(0ifreverseelselen(backtrace), err_message)
forentryinbacktrace:
_flush(entry)

iftb:
backtrace_excepthook(tpe=tpe, value=value, tb=tb)
else:
sys.excepthook=backtrace_excepthook


defunhook():
"""Restore the default excepthook
 """
sys.excepthook=sys.__excepthook__


def_extract_traceback(text):
"""Receive a list of strings representing the input from stdin and return
 the restructured backtrace.

 This iterates over the output and once it identifies a hopefully genuine
 identifier, it will start parsing output.
 In the case the input includes a reraise (a Python 3 case), the primary
 traceback isn't handled, only the reraise.

 Each of the traceback lines are then handled two lines at a time for each
 stack object.

 Note that all parts of each stack object are stripped from newlines and
 spaces to keep the output clean.
 """
capture=False
entries= []
all_else= []
ignore_trace=False

# In python 3, a traceback may includes output from a reraise.
# e.g, an exception is captured and reraised with another exception.
# This marks that we should ignore
iftext.count(TRACEBACK_IDENTIFIER) ==2:
ignore_trace=True

forindex, lineinenumerate(text):
ifTRACEBACK_IDENTIFIERinline:
ifignore_trace:
ignore_trace=False
continue
capture=True
# We're not capturing and making sure we only read lines
# with spaces since, after the initial identifier, all traceback lines
# contain a prefix spacing.
elifcaptureandline.startswith(' '):
ifindex%2==0:
# Line containing a file, line and module.
line=line.strip().strip('\n')
next_line=text[index+1].strip('\n')
entries.append(line+', '+next_line)
elifcapture:
# Line containing the module call.
entries.append(line)
break
else:
# Add everything else after the traceback.
all_else.append(line)

traceback_entries= []

# Build the traceback structure later passed for formatting.
forindex, lineinenumerate(entries[:-2]):
# TODO: This should be done in a _parse_entry function
element=line.split(',')
element[0] =element[0].strip().lstrip('File').strip(' "')
element[1] =element[1].strip().lstrip('line').strip()
element[2] =element[2].strip().lstrip('in').strip()
traceback_entries.append(tuple(element))
returntraceback_entries, all_else


def_stdin_hook(args):
output=sys.stdin.readlines()

ifTRACEBACK_IDENTIFIERnotinoutput:
sys.exit(
'No Traceback detected. Make sure you pipe stderr to '
'backtrace correctly.')

tb, all_else=_extract_traceback(output)
sys.stdout.write(''.join(all_else))
tpe, value=output[-1].strip('\n').split(': ', 1)
hook(
reverse=args.reverse,
align=args.align,
strip_path=args.strip_path,
conservative=args.conservative,
tpe=tpe,
value=value,
tb=tb
 )


def_add_reverse_argument(parser):
parser.add_argument(
'-r',
'--reverse',
action='store_true',
help='Reverse traceback entry order')
returnparser


def_add_align_argument(parser):
parser.add_argument(
'-a',
'--align',
action='store_true',
help='Right-align the backtrace')
returnparser


def_add_strip_path_argument(parser):
parser.add_argument(
'-s',
'--strip-path',
action='store_true',
help='Strip the path to the module')
returnparser


def_add_conservative_argument(parser):
parser.add_argument(
'-c',
'--conservative',
action='store_true',
help='Activate conservative mode')
returnparser


defparse_args():
parser=argparse.ArgumentParser(
description=DESCRIPTION,
formatter_class=argparse.RawTextHelpFormatter)
parser=_add_reverse_argument(parser)
parser=_add_align_argument(parser)
parser=_add_strip_path_argument(parser)
parser=_add_conservative_argument(parser)
parser.set_defaults(func=_stdin_hook)

returnparser.parse_args()


defmain():
args=parse_args()
args.func(args)


if__name__=='__main__':
main()