import types, dis def generate_valid_lnopairs(addr, line): """generates the byte code and line number offsets used in code objects this implements the weird rules for dealing with offsets higher than will fit in the bytes and negative this throws a NotImplementedError for line values less than -128 since I am not sure what is suppose to happen in that case for invalid values like a negative address offset will give an error when passing this value to the bytes constructor.""" if line < -128: raise NotImplementedError("not implemented stepping beyond 128 lines backwards") while addr > 255: yield 255 yield 0 addr -= 255 while line > 127: yield addr yield 127 addr = 0 line -= 127 if line == 0 and addr == 0: return # one or both loops ended perfectly at 0 so don't yield 0,0 pair. if line < 0: # storing negative values as 0x100 + (negative value) # no handling here for values below 0x80 line += 0x100 yield addr yield line def generate_new_lnotab(func): """generates lnotab to fix the line numbers for with statement cleanups""" orig_lnotab = func.__code__.co_lnotab bytecode = dis.Bytecode(func.__code__) last_line = func.__code__.co_firstlineno last_offset = 0 enter_with = [] for inst in bytecode: if inst.starts_line is not None: yield from generate_valid_lnopairs(inst.offset - last_offset, inst.starts_line - last_line) last_offset = inst.offset last_line = inst.starts_line elif inst.opname == "SETUP_WITH": enter_with.append(last_line) elif inst.opname == "WITH_CLEANUP_START": # need to walk back to line of enter_with yield from generate_valid_lnopairs(inst.offset - last_offset, (enter_with[-1] - last_line)) last_offset = inst.offset last_line = enter_with.pop(-1) def correct_with_cleanup_traceback(func): """given a function, corrects the line numbers for the with statement cleanup so errors in exits will point the with statement""" c = func.__code__ func.__code__ = types.CodeType(c.co_argcount, c.co_kwonlyargcount, c.co_nlocals, c.co_stacksize, c.co_flags, c.co_code, c.co_consts, c.co_names, c.co_varnames, c.co_filename, c.co_name, c.co_firstlineno, bytes(generate_new_lnotab(func)), c.co_freevars, c.co_cellvars) class Test(): "example context manager for testing errors in exit method." def __init__(self, fail=False): self.fail = fail def __enter__(self): print("aenter used") def __exit__(self, *errors): if self.fail: raise Exception("error in exit") def my_test(): with Test(True) as fail_during_handling: with Test(True) as will_fail_first: print("inside with") return "hi" #dis.dis(my_test) # show line numbers before correct_with_cleanup_traceback(my_test) dis.dis(my_test) # show line numbers after correction my_test()