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)
