import codetweaks, sys
codetweaks.bind_all(sys.modules[__name__])

from types import ClassType, FunctionType
from opcode import opmap, HAVE_ARGUMENT, EXTENDED_ARG
LOAD_GLOBAL, LOAD_CONST = opmap['LOAD_GLOBAL'], opmap['LOAD_CONST']
ABORT_CODES = (EXTENDED_ARG, opmap['STORE_GLOBAL'])

def bind_constants(f, builtin_only=False, stoplist=[], verbose=False):
    """ Return a new function with optimized global references.

    Replaces global references with their currently defined values.
    If not defined, the dynamic (runtime) global lookup is left undisturbed.
    If builtin_only is True, then only builtins are optimized.
    Variable names in the stoplist are also left undisturbed.
    If verbose is True, prints each substitution as is occurs.

    """
    import __builtin__
    env = vars(__builtin__).copy()
    stoplist = dict.fromkeys(stoplist)
    if builtin_only:
        stoplist.update(f.func_globals)
    else:
        env.update(f.func_globals)

    co = f.func_code
    newcode = map(ord, co.co_code)
    newconsts = list(co.co_consts)
    codelen = len(newcode)

    i = 0
    while i < codelen:
        opcode = newcode[i]
        if opcode in ABORT_CODES:
            return f    # for simplicity, only optimize common cases
        if opcode == LOAD_GLOBAL:
            oparg = newcode[i+1] + (newcode[i+2] << 8)
            name = co.co_names[oparg]
            if name in env and name not in stoplist:
                value = env[name]
                try:
                    pos = newconsts.index(value)
                except ValueError:
                    pos = len(newconsts)
                    newconsts.append(value)
                newcode[i] = LOAD_CONST
                newcode[i+1] = pos & 0xFF
                newcode[i+2] = pos >> 8
                if verbose:
                    print name, '-->', value
        i += 1
        if opcode >= HAVE_ARGUMENT:
            i += 2

    codestr = ''.join(map(chr, newcode))
    codeobj = type(co)(co.co_argcount, co.co_nlocals, co.co_stacksize,
                    co.co_flags, codestr, tuple(newconsts), co.co_names,
                    co.co_varnames, co.co_filename, co.co_name,
                    co.co_firstlineno, co.co_lnotab, co.co_freevars,
                    co.co_cellvars)
    return type(f)(codeobj, f.func_globals, f.func_name, f.func_defaults,
                    f.func_closure)


def bind_all(mc, builtin_only=False, stoplist=[], verbose=False):
    """Recursively apply bind_constants() to functions in a module or class.

    Use as the last line of the module (after everything is defined, but
    before test code).

    In modules that need modifiable globals, set builtin_only to True.

    """
    for k, v in vars(mc).items():
        if type(v) is FunctionType:
            newv = bind_constants(v, builtin_only, stoplist, verbose)
            setattr(mc, k, newv)
        elif type(v) in (type, ClassType):
            bind_all(v, builtin_only, stoplist, verbose)


def f(): pass
try:
    f.func_code.code
except AttributeError:                  # detect non-CPython environments
    bind_all = lambda *args, **kwds: 0
del f

import sys
bind_all(sys.modules[__name__])         # Optimizer, optimize thyself!