import dis import new class MissingLabelError(Exception): """'goto' without matching 'label'.""" pass def goto(fn): """ A function decorator to add the goto command for a function. Specify labels like so: label .foo Goto labels like so: goto .foo """ labels = {} gotos = {} globalName = None index = 0 end = len(fn.func_code.co_code) i = 0 # scan through the byte codes to find the labels and gotos while i < end: op = ord(fn.func_code.co_code[i]) i += 1 name = dis.opname[op] if op > dis.HAVE_ARGUMENT: b1 = ord(fn.func_code.co_code[i]) b2 = ord(fn.func_code.co_code[i+1]) num = b2 * 256 + b1 if name == 'LOAD_GLOBAL': globalName = fn.func_code.co_names[num] index = i - 1 i += 2 continue if name == 'LOAD_ATTR': if globalName == 'label': labels[fn.func_code.co_names[num]] = index elif globalName == 'goto': gotos[fn.func_code.co_names[num]] = index name = None i += 2 # no-op the labels ilist = list(fn.func_code.co_code) for label,index in labels.items(): ilist[index:index+7] = [chr(dis.opmap['NOP'])]*7 # change gotos to jumps for label,index in gotos.items(): if label not in labels: raise MissingLabelError("Missing label: %s"%label) target = labels[label] + 7 # skip NOPs ilist[index] = chr(dis.opmap['JUMP_ABSOLUTE']) ilist[index + 1] = chr(target & 255) ilist[index + 2] = chr(target >> 8) # create new function from existing function c = fn.func_code newcode = new.code(c.co_argcount, c.co_nlocals, c.co_stacksize, c.co_flags, ''.join(ilist), c.co_consts, c.co_names, c.co_varnames, c.co_filename, c.co_name, c.co_firstlineno, c.co_lnotab) newfn = new.function(newcode,fn.func_globals) return newfn if __name__ == '__main__': @goto def test1(n): s = 0 label .myLoop if n <= 0: return s s += n n -= 1 goto .myLoop assert(test1(10) == 55)