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
Download
Copy to clipboard1 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.
Tags: expression, shortcuts
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
xin this scope inner function: 1000 outer function: 100 outer function: Set has no attributeyin this scopeTesting module level: Line no : 1 Line no : 2 Line no : 3 Line no : 4 Line no : 5
Test Output: