Welcome, guest | Sign In | My Account | Store | Cart

Decorator for automatic code optimization. If a global is known at compile time, replace it with a constant. Fold tuples of constants into a single constant. Fold constant attribute lookups into a single constant.

This is only an update of Recipe 277940, making it compatible with Python 3. All credit must go to the original author, Raymond Hettinger.

Python, 200 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
from opcode import opmap, HAVE_ARGUMENT, EXTENDED_ARG
globals().update(opmap)

def _make_constants(f, builtin_only=False, stoplist=[], verbose=False):
    if verbose:
        print('optimizing', f.__name__)
    try:
        co = f.__code__
    except AttributeError:
        return f        # Jython doesn't have a func_code attribute.
    newcode = list(co.co_code)
    newconsts = list(co.co_consts)
    names = co.co_names
    codelen = len(newcode)

    import builtins
    env = vars(builtins).copy()
    if builtin_only:
        stoplist = dict.fromkeys(stoplist)
        stoplist.update(f.__globals__)
    else:
        env.update(f.__globals__)

    # First pass converts global lookups into constants
    changed = False
    i = 0
    while i < codelen:
        opcode = newcode[i]
        if opcode in (EXTENDED_ARG, STORE_GLOBAL):
            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]
                for pos, v in enumerate(newconsts):
                    if v is value:
                        break
                else:
                    pos = len(newconsts)
                    newconsts.append(value)
                newcode[i] = LOAD_CONST
                newcode[i+1] = pos & 0xFF
                newcode[i+2] = pos >> 8
                changed = True
                if verbose:
                    print(name, '-->', value)
        i += 1
        if opcode >= HAVE_ARGUMENT:
            i += 2

    # Second pass folds tuples of constants and constant attribute lookups
    i = 0
    while i < codelen:

        newtuple = []
        while newcode[i] == LOAD_CONST:
            oparg = newcode[i+1] + (newcode[i+2] << 8)
            newtuple.append(newconsts[oparg])
            i += 3

        opcode = newcode[i]
        if not newtuple:
            i += 1
            if opcode >= HAVE_ARGUMENT:
                i += 2
            continue

        if opcode == LOAD_ATTR:
            obj = newtuple[-1]
            oparg = newcode[i+1] + (newcode[i+2] << 8)
            name = names[oparg]
            try:
                value = getattr(obj, name)
            except AttributeError:
                continue
            deletions = 1

        elif opcode == BUILD_TUPLE:
            oparg = newcode[i+1] + (newcode[i+2] << 8)
            if oparg != len(newtuple):
                continue
            deletions = len(newtuple)
            value = tuple(newtuple)

        else:
            continue

        reljump = deletions * 3
        newcode[i-reljump] = JUMP_FORWARD
        newcode[i-reljump+1] = (reljump-3) & 0xFF
        newcode[i-reljump+2] = (reljump-3) >> 8

        n = len(newconsts)
        newconsts.append(value)
        newcode[i] = LOAD_CONST
        newcode[i+1] = n & 0xFF
        newcode[i+2] = n >> 8
        i += 3
        changed = True
        if verbose:
            print("new folded constant:", value)

    if not changed:
        return f

    codestr = bytes(newcode)
    codeobj = type(co)(co.co_argcount, co.co_kwonlyargcount, 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.__globals__, f.__name__, f.__defaults__,
                    f.__closure__)

_make_constants = _make_constants(_make_constants) # optimize thyself!

def bind_all(mc, builtin_only=False, stoplist=[],  verbose=False):
    """Recursively apply constant binding 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.

    """
    from types import FunctionType, ModuleType

    def _bind_all(mc, builtin_only=False, stoplist=[],  verbose=False):

      if isinstance(mc, (ModuleType, type)):
        for k,v in list(vars(mc).items()):
          if type(v) is FunctionType:
              newv = _make_constants(v, builtin_only, stoplist,  verbose)
              setattr(mc, k, newv)
          elif isinstance(v, type):
              _bind_all(v, builtin_only, stoplist, verbose)

    if isinstance(mc, dict): # allow: bind_all(globals())
      for k,v in list(mc.items()):
        if type(v) is FunctionType:
            newv = _make_constants(v, builtin_only, stoplist,  verbose)
            mc[k] = newv
        elif isinstance(v, type):
            _bind_all(v, builtin_only, stoplist, verbose)
    else:
        _bind_all(v, builtin_only, stoplist, verbose)

@_make_constants
def make_constants(builtin_only=False, stoplist=[], verbose=False):
    """ Return a decorator for optimizing 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.
    Also, folds constant attr lookups and tuples of constants.
    If verbose is True, prints each substitution as is occurs

    """
    if type(builtin_only) == type(make_constants):
        raise ValueError("The bind_constants decorator must have arguments.")
    return lambda f: _make_constants(f, builtin_only, stoplist, verbose)

## --------- Example call -----------------------------------------

import random

@make_constants(verbose=True, stoplist=['random'])
def sample(population, k):
    "Choose k unique random elements from a population sequence."

    if not isinstance(population, (list, tuple, str)):
        raise TypeError('Cannot handle type', type(population))
    n = len(population)
    if not 0 <= k <= n:
        raise ValueError("sample larger than population")
    result = [None] * k
    pool = list(population)
    for i in range(k):         # invariant:  non-selected at [0,n-i)
        j = int(random.random() * (n-i))
        result[i] = pool[j]
        pool[j] = pool[n-i-1]   # move non-selected item into vacancy
    return result


""" Output from the example call:

isinstance --> <built-in function isinstance>
list --> <class 'list'>
tuple --> <class 'tuple'>
str --> <class 'str'>
TypeError --> <class 'TypeError'>
type --> <class 'type'>
len --> <built-in function len>
ValueError --> <class 'ValueError'>
list --> <class 'list'>
range --> <class 'range'>
int --> <class 'int'>
new folded constant: (<class 'list'>, <class 'tuple'>, <class 'str'>)
"""

See original discussion at recipe 277940.

Minor changes:

  • typos noted in the comment section are fixed

  • bind_all now accepts a dictionary as well, so you can use bind_all(globals()) at the end of a module to optimize the whole module.

  • comment #7 (optimize closures) is NOT implemented.

2 comments

David Cortesi 8 years, 4 months ago  # | flag

Two things that seem no longer applicable in Python 3.4 at least.

First, I cannot find a way to make Python3 compile the sequence, LOAD_CONST, LOAD_CONST... BUILD_TUPLE. It appears this optimization has already be implemented in the 3 compile phase because code such as z = (1,2) produces a LOAD_CONST with an argument of a tuple, (1,2). So I think the whole exercise of pre-building constant tuples in this way is moot now.

Second, line 69, if opcode == LOAD_ATTR: can only be reached if the preceding opcode was LOAD_CONST. (If it wasn't LOAD_CONST, the test at line 63 will find newtuple null and continue the loop.) However, I can't find a way to get Python3 bytecode to produce the sequence, LOAD_CONST, LOAD_ATTR. Maybe it did in Python2, but in 3, the preceding opcode is always LOAD_FAST. As a result I think probably the code in 70-77 will never be executed in Python3.

I can produce a situation that would benefit from an optimization like this: z = ( 1, an_obj.a_member ) produces the bytecode,

LOAD_CONST  1
LOAD_FAST 0 (an_obj)
LOAD_ATTR 1 ('a_member')
BUILD_TUPLE

The logic of the second pass would have to be heavily revised to handle this.

David Cortesi 8 years, 2 months ago  # | flag

OK I take that back; it IS possible for Python 3.5.1 to generate LOAD_GLOBAL, LOAD_GLOBAL... BUILD_TUPLE. After processing by phase 1, this will be the sequence LOAD_CONST, LOAD_CONST... BUILD_TUPLE, which phase 2 can indeed optimize.

My bad. Sorry for the noise.