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

Python does not support assignment in if and while statements such as "if (x=func()):". This is an attempt to bring similar functionality to python by injecting bytecode to all functions and methods in a module. This recipe is inspired from recipes 66061, 202234 and 277940.

Python, 136 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
from opcode import opmap, HAVE_ARGUMENT
globals().update(opmap)

class DataHolder(object):
    _varname_ = 'Set'
    # This varname `Set` should be treated like a reserved keyword
    # and should be used for other purpose at any scope.
    def __call__(self, **kwargs):
        if len(kwargs) != 1:
            raise TypeError(
                '%s takes exactly 1 keyword argument (%s given)'%(
                 self._varname_, len(kwargs)))
        name, value = kwargs.popitem()
        setattr(self, name, value)
        return value

def _support_testassign(f):
    co       = f.__code__
    code     = list(co.co_code)
    consts   = list(co.co_consts)
    varnames = list(co.co_varnames)
    if consts[-1] is DataHolder: # already applied
        return
    code.insert(0, LOAD_CONST)
    code.insert(1, len(consts) & 0xFF)
    code.insert(2, len(consts) >> 8)
    code.insert(3, CALL_FUNCTION)
    code.insert(4, 0 & 0xFF)
    code.insert(5, 0 >> 8)
    code.insert(6, STORE_FAST)
    code.insert(7, len(varnames) & 0xFF)
    code.insert(8, len(varnames) >> 8)

    consts.append(DataHolder)
    varnames.append(DataHolder._varname_)

    i, pos = 0, len(varnames)-1
    while i < len(code):
        opcode = code[i]
        if opcode == LOAD_GLOBAL:
            oparg = code[i+1] + (code[i+2] << 8)
            name = co.co_names[oparg]
            if name == DataHolder._varname_:
                code[i] = LOAD_FAST
                code[i+1] = pos & 0xFF
                code[i+2] = pos >> 8
        elif (opcode == CONTINUE_LOOP or
              JUMP_IF_FALSE_OR_POP <= opcode <= POP_JUMP_IF_TRUE):
            oparg = code[i+1] + (code[i+2] << 8) + 9
            code[i+1] = oparg & 0xFF
            code[i+2] = oparg >> 8
        i += 1
        if opcode >= HAVE_ARGUMENT:
            i += 2
    codeobj = type(co)(co.co_argcount, co.co_kwonlyargcount,
                       co.co_nlocals+1, co.co_stacksize, co.co_flags,
                       bytes(code), tuple(consts), co.co_names,
                       tuple(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__)

def install_testassign(mc):
    # mc can be a module or globals() dict
    from types import FunctionType
    if isinstance(mc, dict):
        d = mc
        d[DataHolder._varname_] = DataHolder()
    else:
        try:
            d = vars(mc)
        except TypeError:
            return
    for k, v in d.items():
        if v in (_support_testassign, install_testassign, DataHolder):
            continue
        if isinstance(v, FunctionType):
            newv = _support_testassign(v)
            try:
                d[k] = newv
            except TypeError:
                setattr(mc, k, newv)
        elif isinstance(v, type):
            try:
                setattr(v, DataHolder._varname_, DataHolder())
            except Exception:
                pass
            install_testassign(v)


def test_while(file):
    while Set(line=file.readline()):
        print(Set.line.rstrip())

def test_recursion(file):
    if Set(value=file.readline()):
        test_recursion(file)
        print(Set.value.rstrip())

def test_nonlocal():
    Set(x=100)
    def sub():
        Set(y=1000)
        print('inner function:', getattr(Set, 'x', 'Set has no attribute `x` in this scope'))
        print('inner function:', getattr(Set, 'y', 'Set has no attribute `y` in this scope'))
    sub()
    print('outer function:', getattr(Set, 'x', 'Set has no attribute `x` in this scope'))
    print('outer function:', getattr(Set, 'y', 'Set has no attribute `y` in this scope'))


# This should be called after all function and
# class definitions in a module.
install_testassign(globals())

# ---- Begin Test ---------

from io import StringIO
file = StringIO('\n'.join('Line no : %d'%(i+1) for i in range(5)))

print('Testing while statement:')
test_while(file)

print('\nTesting recursion:')
file.seek(0)
test_recursion(file)

print('\nTesting nonlocal scope:')
test_nonlocal()

print('\nTesting module level:')
# Using `Set` in the Module level scope can only
# be done after calling install_testassign.
file.seek(0)
while Set(line=file.readline()):
    print(Set.line.rstrip())

This function install_testassign when called, modifies every function and method in a module such that the line Set=DataHolder() is added to the beginning if their bytecode. The function,

def test_while(file):
    while Set(line=file.readline()):
        print(Set.line.rstrip())

After modifying the function it becomes,

def test_while(file):
    Set=DataHolder()
    while Set(line=file.readline()):
        print(Set.line.rstrip())

This affects lambda and nested functions too. Since this is done at compile time, LOAD_CONST has been used for faster loading of DataHolder class instead of LOAD_GLOBAL. With this recipe we can use multiple 'assign and test' statements at any scope. Also, the Set object is local to the function/method scope and hence works with recursion too.

2 comments

harish anand (author) 9 years, 11 months ago  # | flag

Test Output: Testing while statement: Line no : 1 Line no : 2 Line no : 3 Line no : 4 Line no : 5

Testing recursion: Line no : 5 Line no : 4 Line no : 3 Line no : 2 Line no : 1

Testing nonlocal scope: inner function: Set has no attribute x in this scope inner function: 1000 outer function: 100 outer function: Set has no attribute y in this scope

Testing module level: Line no : 1 Line no : 2 Line no : 3 Line no : 4 Line no : 5

harish anand (author) 9 years, 11 months ago  # | flag

Test Output:

Testing while statement:
Line no : 1
Line no : 2
Line no : 3
Line no : 4
Line no : 5

Testing recursion:
Line no : 5
Line no : 4
Line no : 3
Line no : 2
Line no : 1

Testing nonlocal scope:
inner function: Set has no attribute `x` in this scope
inner function: 1000
outer function: 100
outer function: Set has no attribute `y` in this scope

Testing module level:
Line no : 1
Line no : 2
Line no : 3
Line no : 4
Line no : 5