/reassembly.py

from dis import disassemble

from opcode import HAVE_ARGUMENT, opname, opmap, hasjabs

import inspect 

from StringIO import StringIO

import new

'''

Doesn't work.



The problem is, I can't capture the local value stack and, more importantly,

block stack. Which means that `for i in range(10)` can't be properly restored,

given that range iterator is stored in one of the them.

'''

class Saved_frame(object):

    def __init__(self, ip, locals, retval):

        self.ip = ip

        self.locals = locals

        self.retval = retval



def save_frame(retval):

    frame = inspect.currentframe(1)

    ldict = frame.f_locals

    print 'Locals: {!r}'.format(ldict)

    locals = []

    for name in frame.f_code.co_varnames:

        print name, ldict[name]

        locals.append(ldict[name])

    return Saved_frame(frame.f_lasti, locals, retval)



class Instruction(object):

    def __init__(self, addr, op, arg = None):

        self.addr = addr

        self.op = op

        self.arg = arg

    def __repr__(self):

        s = '{:>3}: {}'.format(self.addr, opname[self.op])

        if self.arg is not None:

            s += ' ' + str(self.arg)

        return s



def disassemble_to_list(code):

    lst = []

    n = len(code)

    i = 0

    while i < n:

        c = code[i]

        op = ord(c)

        addr = i

        i += 1

        oparg = None

        if op >= HAVE_ARGUMENT:

            oparg = ord(code[i]) + ord(code[i + 1])*256

            i += 2

        lst.append(Instruction(addr, op, oparg))

    return lst

        

def assemble_from_list(lst, offset = 0):

    out = StringIO()

    for it in lst:

        out.write(chr(it.op))

        if it.op >= HAVE_ARGUMENT:

            arg = it.arg

            if it.op in hasjabs:

                arg += offset

            out.write(chr(arg & 0xFF))

            out.write(chr(arg >> 8))

    return out.getvalue()





def patch(f):

    # new.function(code, globals[, name[, argdefs[, closure]]]) 

    # new.code(argcount, nlocals, stacksize, flags, codestring, constants, names, varnames, filename, name, firstlineno, lnotab) 

    code = f.__code__

    lst = disassemble_to_list(code.co_code)

    local_cnt = code.co_nlocals

    

    for it in lst:

        if it.op == opmap['RETURN_VALUE']:

            return_addr = it.addr + 1

            break

    

    patch = [

            Instruction(0, opmap['LOAD_FAST'], 0),

            Instruction(0, opmap['UNPACK_SEQUENCE'], local_cnt),]

    

    for i in range(local_cnt):

        patch.append(Instruction(0, opmap['STORE_FAST'], i))

        

    assembled = assemble_from_list(patch)

    patch_length = len(assembled) + 3

    

    patch = [Instruction(0, opmap['JUMP_ABSOLUTE'], return_addr + patch_length)]

    assembled += assemble_from_list(patch)

    

    assert len(assembled) == patch_length 

    assembled += assemble_from_list(lst, patch_length)



    patched_code = new.code(1, code.co_nlocals, code.co_stacksize, 

            code.co_flags, assembled, code.co_consts, code.co_names, 

            code.co_varnames, code.co_filename, code.co_name, 

            code.co_firstlineno, code.co_lnotab)

    patched_function = new.function(patched_code, f.func_globals, 'generator', None, f.func_closure)

    return patched_function



def facs(step = 1):

    acc = 1

    n = 0

    while True:

#    for _ in range(10):

        n += step

        acc *= n

        return save_frame(acc)

        acc += 0 # to prevent the compiler from removing the jump back 



facs2 = patch(facs)



print disassemble(facs.__code__)

print disassemble(facs2.__code__)

# print disassemble(facs_patched.__code__)

print 'start'

frame = facs()

print frame.__dict__

frame = facs2(frame.locals)

print frame.__dict__

frame = facs2(frame.locals)

print frame.__dict__