import dis
import new
import types

_SUPER_KEYWORD='super'

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

try:

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

    LOAD_GLOBAL = opmap['LOAD_GLOBAL']
    LOAD_NAME = opmap['LOAD_NAME']
    LOAD_CONST = opmap['LOAD_CONST']
    LOAD_FAST = opmap['LOAD_FAST']
    LOAD_ATTR = opmap['LOAD_ATTR']
    STORE_FAST = opmap['STORE_FAST']
    LOAD_DEREF = opmap['LOAD_DEREF']
    STORE_DEREF = opmap['STORE_DEREF']
    CALL_FUNCTION = opmap['CALL_FUNCTION']
    STORE_GLOBAL = opmap['STORE_GLOBAL']
    DUP_TOP = opmap['DUP_TOP']
    POP_TOP = opmap['POP_TOP']
    NOP = opmap['NOP']
    JUMP_FORWARD = opmap['JUMP_FORWARD']

except ImportError:

    # Python 2.2 doesn't have opcode
    HAVE_ARGUMENT = 90
    EXTENDED_ARG = 143

    LOAD_GLOBAL = 116
    LOAD_NAME = 101
    LOAD_CONST = 100
    LOAD_FAST = 124
    LOAD_ATTR = 105
    LOAD_DEREF = 136
    STORE_DEREF = 137
    STORE_FAST = 125
    CALL_FUNCTION = 131
    STORE_GLOBAL = 97
    DUP_TOP = 4
    POP_TOP = 1
    NOP = 9
    JUMP_FORWARD = 110

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

try:
    BaseException
except NameError:
    BaseException = Exception

ABSOLUTE_TARGET = dis.hasjabs

try:
#    raise ImportError

    # First we try to import the compiled pyrex dll
    from _super import _super

except ImportError:
#    import traceback
#    traceback.print_exc()

    _object_getattribute = object.__getattribute__
    _object_setattr = object.__setattr__
    _object_delattr = object.__delattr__
    _builtin_super=super

    class _super (object):
        class _super_exception(object):
            __slots__ = ('exc_type', 'msg')

            def __init__(self, exc):
                self.exc_type = exc.__class__
                self.msg = str(exc)

        """
        Wrapper for the super object.

        If called, a base class method of the same name as the current method
        will be called. Otherwise, attributes will be accessed.
        """

        # We want the lowest overhead possible - both speed and memory.
        # We need to put the mangled name in as Python 2.2.x didn't work
        # with private names specified in __slots__.
        __slots__ = ('_super__super', '_super__method')

        def __init__(self, cls, obj=None, name=None):
            super = _builtin_super(cls, obj)
            _object_setattr(self, '_super__super', super)

            if name is not None:
                try:
                    method = getattr(super, name)
                except AttributeError, e:
                    method = _super._super_exception(e)
            else:
                method = None

            _object_setattr(self, '_super__method', method)

        def __call__(self, *p, **kw):
            """
            Calls the base class method with the passed parameters.
            """
            method = _object_getattribute(self, '_super__method')

            if method is None:
                raise AttributeError("'super' object has no default attribute")
            elif type(method) is _super._super_exception:
                raise method.exc_type(method.msg)

            return method(*p, **kw)

        def __getattribute__ (self, name):
            """
            Gets a base class attribute.
            """
            if name == 'super':
                raise AttributeError("'super' object has no attribute 'super'")

            super = _object_getattribute(self, '_super__super')
            return getattr(super, name)

        def __setattr__(self, name, value):
            """
            All we want to do here is make it look the same as if we called
            setattr() on a real `super` object.
            """
            super = _object_getattribute(self, '_super__super')
            _object_setattr(super, name, value)

        def __delattr__(self, name):
            """
            All we want to do here is make it look the same as if we called
            delattr() on a real `super` object.
            """
            super = _object_getattribute(self, '_super__super')
            _object_delattr(super, name, value)

        def __repr__(self):
            super = _object_getattribute(self, '_super__super')
            method = _object_getattribute(self, '_super__method')

            if method is None:
                method = 'NULL'
            elif type(method) is _super._super_exception:
                method = '<%s: %s>' % (method.exc_type.__name__, method.msg)
            else:
                # Try to display the real method - will only work properly for classes
                # that have _autosuper as the metaclass (because we need func.func_class)
                ms = str(method)

                try:
                    try:
                        func_class_name = method.im_func.func_class.__name__
                    except AttributeError:
                        # Not everything uses _autosuper
                        func_class_name = method.im_class.__name__ + '?'

                    ms = ms.replace('bound method %s.' % (method.im_class.__name__,),
                                    'bound method %s.' % (func_class_name))

                except AttributeError:
                    # Could be a method-wrapper
                    pass

                method = ms

            s = repr(super)
            return '%s, %s%s' % (s[:-1], method, s[-1:])

def _oparg(code, opcode_pos):
    return code[opcode_pos+1] + (code[opcode_pos+2] << 8)

def _bind_autosuper(func, cls):
    co = func.func_code
    name = func.func_name
    newcode = map(ord, co.co_code)
    codelen = len(newcode)
    newconsts = list(co.co_consts)
    newvarnames = list(co.co_varnames)

    # Check if the global super keyword is already present
    try:
        sn_pos = list(co.co_names).index(_SUPER_KEYWORD)
    except ValueError:
        sn_pos = None

    # Check if the varname super keyword is already present
    try:
        sv_pos = newvarnames.index(_SUPER_KEYWORD)
    except ValueError:
        sv_pos = None

    # Check if the callvar super keyword is already present
    try:
        sc_pos = list(co.co_cellvars).index(_SUPER_KEYWORD)
    except ValueError:
        sc_pos = None

    # If 'super' isn't used anywhere in the function, we don't have anything to do
    if sn_pos is None and sv_pos is None and sc_pos is None:
        return func

    c_pos = None
    s_pos = None
    n_pos = None

    # Check if the 'cls_name' 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

        if o == name:
            n_pos = pos

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

    if n_pos is None:
        n_pos = len(newconsts)
        newconsts.append(name)

    if s_pos is None:
        s_pos = len(newconsts)
        newconsts.append(_super)

    if sv_pos is None:
        sv_pos = len(newvarnames)
        newvarnames.append(_SUPER_KEYWORD)

    # This goes at the start of the function. It is:
    #
    #   super = super(cls, self, name)
    #
    # If 'super' is a cell variable, we store to both the
    # local and cell variables (i.e. STORE_FAST and STORE_DEREF).
    #
    setup_code = [
        LOAD_CONST, s_pos & 0xFF, s_pos >> 8,
        LOAD_CONST, c_pos & 0xFF, c_pos >> 8,
        LOAD_FAST, 0, 0,
        LOAD_CONST, n_pos & 0xFF, n_pos >> 8,
        CALL_FUNCTION, 3, 0,
    ]

    if sc_pos is None:
        # 'super' is not a cell variable - we can just use the local variable
        setup_code_noderef = [
            STORE_FAST, sv_pos & 0xFF, sv_pos >> 8,
        ]
    else:
        # If 'super' is a cell variable, we need to handle LOAD_DEREF -
        # we default to just STORE_FAST, but if we actually have a LOAD_DEREF
        # we'll replace it with both STORE_FAST and STORE_DEREF.
        setup_code_noderef = [
            STORE_FAST, sv_pos & 0xFF, sv_pos >> 8,
            JUMP_FORWARD, 1 & 0xFF, 1 >> 8,
            NOP,
        ]

        setup_code_deref = [
            DUP_TOP,
            STORE_FAST, sv_pos & 0xFF, sv_pos >> 8,
            STORE_DEREF, sc_pos & 0xFF, sc_pos >> 8,
        ]

    setup_code += setup_code_noderef

    # Bytecode for loading the local 'super' variable.
    load_super = [
        LOAD_FAST, sv_pos & 0xFF, sv_pos >> 8,
    ]

    # This is used to prevent anyone changing the local or cell variable
    nostore_super = [
        POP_TOP, NOP, NOP
    ]

    setup_code_len = len(setup_code)
    need_setup = False
    need_deref = False
    i = 0

    while i < codelen:
        opcode = newcode[i]
        need_load = False
        remove_store = False

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

        # 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 = _oparg(newcode, i) + setup_code_len
            newcode[i+1] = oparg & 0xFF
            newcode[i+2] = oparg >> 8

        # If LOAD_GLOBAL(super) or LOAD_NAME(super) then we want to change it into
        # LOAD_FAST(super)
        elif (opcode == LOAD_GLOBAL or opcode == LOAD_NAME) and _oparg(newcode, i) == sn_pos:
            need_setup = need_load = True

        # If LOAD_FAST(super) then we just need to add the setup code
        elif opcode == LOAD_FAST and _oparg(newcode, i) == sv_pos:
            need_setup = need_load = True

        # If LOAD_DEREF(super) then we change it into LOAD_FAST(super) because
        # it's slightly faster.
        elif opcode == LOAD_DEREF and _oparg(newcode, i) == sc_pos:
            need_setup = need_deref = need_load = True

        # If it's STORE_FAST(super) then we prevent the store.
        elif opcode == STORE_FAST and _oparg(newcode, i) == sv_pos:
            remove_store = True

        # If it's STORE_DEREF(super) then we prevent the store.
        elif opcode == STORE_DEREF and _oparg(newcode, i) == sc_pos:
            remove_store = True

        if need_load:
            newcode[i:i+3] = load_super
        elif remove_store:
            newcode[i:i+3] = nostore_super

        i += 1

        if opcode >= HAVE_ARGUMENT:
            i += 2

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

    if need_deref:
        setup_code[-len(setup_code_deref):] = setup_code_deref

    # Our setup code will have 4 things on the stack
    co_stacksize = max(4, 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] += setup_code_len
    
    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), co.co_names,
                       tuple(newvarnames), co.co_filename, co.co_name,
                       co.co_firstlineno, co_lnotab, co.co_freevars,
                       co.co_cellvars)

#    dis.dis(codeobj)
#    print

    func.func_code = codeobj
    func.func_class = cls
    return func

class _autosuper(type):
    def __init__(cls, name, bases, clsdict):
        UnboundMethodType = types.UnboundMethodType

        for v in vars(cls):
            o = getattr(cls, v)
            if isinstance(o, UnboundMethodType):
                _bind_autosuper(o.im_func, cls)

class autosuper(object):
    __metaclass__ = _autosuper

if __name__ == '__main__':

    import autosuper_object

    class A(object):
        def f(self):
            # Make super a cell variable. Normally this would throw an
            # UnboundLocalError, but our bytecode hack makes it work.
            super = super

            def a():
                print 'a:', super

            print 'A:', super
            a()

    class B(A):
        def f(self):
            # Make super a cell variable. Normally this would throw an
            # UnboundLocalError, but our bytecode hack makes it work.
            super = super

            def b1():
                def b():
                    print 'b:', super
                    super()

                b()

            print 'B:', super
            b1()

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

    class D(B, C):
        pass

    class E(D):
        def f(self):
            print 'E:', super
            super()

    class F(E, A):
        pass

    F().f()
