import new
import types
import sys
sys.setrecursionlimit(10)

try:
    True
except NameError:
    # Do it this way to avoid getting a warning
    import __builtin__
    __builtin__.__dict__['True'] = 1
    __builtin__.__dict__['False'] = 1

try:

    # We need these for modifying bytecode
    from opcode import opmap, HAVE_ARGUMENT, EXTENDED_ARG

    LOAD_GLOBAL = opmap['LOAD_GLOBAL']
    LOAD_CONST = opmap['LOAD_CONST']
    LOAD_FAST = opmap['LOAD_FAST']
    LOAD_ATTR = opmap['LOAD_ATTR']
    STORE_FAST = opmap['STORE_FAST']
    CALL_FUNCTION = opmap['CALL_FUNCTION']

    STORE_GLOBAL = opmap['STORE_GLOBAL']
    JUMP_FORWARD = opmap['JUMP_FORWARD']
    JUMP_ABSOLUTE = opmap['JUMP_ABSOLUTE']
    CONTINUE_LOOP = opmap['CONTINUE_LOOP']

except ImportError:

    # Python 2.2 doesn't have opcode
    LOAD_GLOBAL = 116
    LOAD_CONST = 100
    LOAD_FAST = 124
    LOAD_ATTR = 105
    STORE_FAST = 125
    CALL_FUNCTION = 131
    STORE_GLOBAL = 97
    JUMP_FORWARD = 110
    JUMP_ABSOLUTE = 113
    CONTINUE_LOOP = 119
    HAVE_ARGUMENT = 90
    EXTENDED_ARG = 143

    # Doesn't have enumerate either
    def enumerate (sequence):
        return zip(range(len(sequence)), sequence)

ABSOLUTE_TARGET = (JUMP_ABSOLUTE, CONTINUE_LOOP)

def _bind_autosuper(func, cls, super=super):
    co = func.func_code
    newcode = map(ord, co.co_code)
    codelen = len(newcode)
    newnames = list(co.co_names)
    newconsts = list(co.co_consts)
    newvarnames = list(co.co_varnames)
    c_pos = -1
    s_pos = -1
    sn_pos = -1
    sv_pos = -1

    for pos, o in enumerate(newnames):
        if o == 'super':
            sn_pos = pos
            break

    # 'super' isn't used anywhere in the function, so we don't have anything to do
    if sn_pos == -1:
        return func

    # Check if the 'cls' and 'super' objects are already in the constants
    for pos, o in enumerate(newconsts):
        if o is cls:
            c_pos = pos

        if o is super:
            s_pos = pos

    # Check if the varname 'super' is already present
    for pos, o in enumerate(newvarnames):
        if o == 'super':
            sv_pos = pos
            break

    # Add in any missing objects to constants and varnames
    if c_pos == -1:
        c_pos = len(newconsts)
        newconsts.append(cls)

    if s_pos == -1:
        s_pos = len(newconsts)
        newconsts.append(super)

    if sv_pos == -1:
        sv_pos = len(newvarnames)
        newvarnames.append('super')

    # This goes at the start of the function. It is:
    #
    #   super = super(cls, self)

    setup_code = [
        LOAD_CONST, s_pos & 0xFF, s_pos >> 8,
        LOAD_CONST, c_pos & 0xFF, c_pos >> 8,
        LOAD_FAST, 0, 0,
        CALL_FUNCTION, 2, 0,
        STORE_FAST, sv_pos & 0xFF, sv_pos >> 8,
    ]

    need_setup = False
    i = 0

    while i < codelen:
        opcode = newcode[i]

        if opcode == EXTENDED_ARG:
            raise TypeError("Cannot use 'super' in function with EXTENDED_ARG opcode")

        # If LOAD_GLOBAL(super) then we want to change it into LOAD_FAST(super)
        elif opcode == LOAD_GLOBAL:
            oparg = newcode[i+1] + (newcode[i+2] << 8)

            if oparg == sn_pos:
                need_setup = True
                newcode[i] = LOAD_FAST
                newcode[i+1] = sv_pos & 0xFF
                newcode[i+2] = sv_pos >> 8

        # If the opcode is an absolute target it needs to be adjusted
        # to take into account the setup code.
        elif opcode in ABSOLUTE_TARGET:
            oparg = newcode[i+1] + (newcode[i+2] << 8) + len(setup_code)
            newcode[i+1] = oparg & 0xFF
            newcode[i+2] = oparg >> 8

        i += 1

        if opcode >= HAVE_ARGUMENT:
            i += 2

    # No changes needed - get out.
    if not need_setup:
        return func

    # Our setup code will have 3 things on the stack
    co_stacksize = max(3, co.co_stacksize)

    # Conceptually, our setup code is on the `def` line.
    co_lnotab = map(ord, co.co_lnotab)

    if co_lnotab:
        co_lnotab[0] += len(setup_code)
    
    co_lnotab = ''.join(map(chr, co_lnotab))

    # Our code consists of the setup code and the modified code.
    codestr = ''.join(map(chr, setup_code + newcode))

    codeobj = new.code(co.co_argcount, len(newvarnames), co_stacksize,
                       co.co_flags, codestr, tuple(newconsts), tuple(newnames),
                       tuple(newvarnames), co.co_filename, co.co_name,
                       co.co_firstlineno, co_lnotab, co.co_freevars,
                       co.co_cellvars)

#    import dis
#    dis.dis(codeobj)

    func.func_code = codeobj
    return func

class _autosuper(type):
    def __init__(cls, name, bases, clsdict):
        for v in vars(cls):
            o = getattr(cls, v)

            if isinstance(o, types.UnboundMethodType):
                _bind_autosuper(o.im_func, cls)

class autosuper(object):
    __metaclass__ = _autosuper

if __name__ == '__main__':

    class A(autosuper):
        def f(self):
            print 'A:', super

    class B(A):
        def f(self):
            print 'B:', super
            super.f()

    class C(A):
        def f(self):
            print 'C:', super
            super.f()

    class D(B, C):
        def f(self):
            print 'D:', super
            super.f()

    D().f()
