1"""
2Some functions for optimzing code, specifically for turning references to
3globals into constants where possible.
4
5See: http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/277940
6"""
7from opcode import opmap, HAVE_ARGUMENT, EXTENDED_ARG
8from types import *
9globals().update(opmap)
10
11def make_constants(f, builtin_only=False, stoplist=[], verbose=False):
12    try:
13        co = f.func_code
14    except AttributeError:
15        return f        # Jython doesn't have a func_code attribute.
16    newcode = map(ord, co.co_code)
17    newconsts = list(co.co_consts)
18    names = co.co_names
19    codelen = len(newcode)
20
21    import __builtin__
22    env = vars(__builtin__).copy()
23    if builtin_only:
24        stoplist = dict.fromkeys(stoplist)
25        stoplist.update(f.func_globals)
26    else:
27        env.update(f.func_globals)
28
29    # First pass converts global lookups into constants
30    i = 0
31    while i < codelen:
32        opcode = newcode[i]
33        if opcode in (EXTENDED_ARG, STORE_GLOBAL):
34            return f    # for simplicity, only optimize common cases
35        if opcode == LOAD_GLOBAL:
36            oparg = newcode[i+1] + (newcode[i+2] << 8)
37            name = co.co_names[oparg]
38            if name in env and name not in stoplist:
39                value = env[name]
40                for pos, v in enumerate(newconsts):
41                    if v is value:
42                        break
43                else:
44                    pos = len(newconsts)
45                    newconsts.append(value)
46                newcode[i] = LOAD_CONST
47                newcode[i+1] = pos & 0xFF
48                newcode[i+2] = pos >> 8
49                if verbose:
50                    print name, '-->', value
51        i += 1
52        if opcode >= HAVE_ARGUMENT:
53            i += 2
54
55    # Second pass folds tuples of constants and constant attribute lookups
56    i = 0
57    while i < codelen:
58
59        newtuple = []
60        while newcode[i] == LOAD_CONST:
61            oparg = newcode[i+1] + (newcode[i+2] << 8)
62            newtuple.append(newconsts[oparg])
63            i += 3
64
65        opcode = newcode[i]
66        if not newtuple:
67            i += 1
68            if opcode >= HAVE_ARGUMENT:
69                i += 2
70            continue
71
72        if opcode == LOAD_ATTR:
73            obj = newtuple[-1]
74            oparg = newcode[i+1] + (newcode[i+2] << 8)
75            name = names[oparg]
76            try:
77                value = getattr(obj, name)
78            except AttributeError:
79                continue
80            deletions = 1
81
82        elif opcode == BUILD_TUPLE:
83            oparg = newcode[i+1] + (newcode[i+2] << 8)
84            if oparg != len(newtuple):
85                continue
86            deletions = len(newtuple)
87            value = tuple(newtuple)
88
89        else:
90            continue
91
92        reljump = deletions * 3
93        newcode[i-reljump] = JUMP_FORWARD
94        newcode[i-reljump+1] = (reljump-3) & 0xFF
95        newcode[i-reljump+2] = (reljump-3) >> 8
96
97        n = len(newconsts)
98        newconsts.append(value)
99        newcode[i] = LOAD_CONST
100        newcode[i+1] = n & 0xFF
101        newcode[i+2] = n >> 8
102        i += 3
103        if verbose:
104            print "new folded constant:", value
105
106    codestr = ''.join(map(chr, newcode))
107    codeobj = type(co)(co.co_argcount, co.co_nlocals, co.co_stacksize,
108                    co.co_flags, codestr, tuple(newconsts), co.co_names,
109                    co.co_varnames, co.co_filename, co.co_name,
110                    co.co_firstlineno, co.co_lnotab, co.co_freevars,
111                    co.co_cellvars)
112    return type(f)(codeobj, f.func_globals, f.func_name, f.func_defaults,
113                    f.func_closure)
114
115make_constants = make_constants(make_constants) # optimize thyself!
116
117def bind_all(mc, builtin_only=False, stoplist=[],  verbose=False):
118    """Recursively apply constant binding to functions in a module or class.
119
120    Use as the last line of the module (after everything is defined, but
121    before test code).  In modules that need modifiable globals, set
122    builtin_only to True.
123
124    """
125    try:
126        d = vars(mc)
127    except TypeError:
128        return
129    for k, v in d.items():
130        if type(v) is FunctionType:
131            newv = make_constants(v, builtin_only, stoplist,  verbose)
132            setattr(mc, k, newv)
133        elif type(v) in (type, ClassType):
134            bind_all(v, builtin_only, stoplist, verbose)
135