[Python-ideas] For-loop variable scope: simultaneous possession and ingestion of cake

[Arnaud Delobelle]

On 13 Oct 2008, at 01:24, George Sakkis wrote:

> Since this idea didn't get much steam, a more modest proposal would  
> be to relax the restriction on cells: allow the creation of new  
> cells and the rebinding of func_closure in pure Python. Then one  
> could explicitly create a new scope without any other change in the  
> language through a 'localize' decorator that would create a new cell  
> for every free variable (i.e. global or value from an enclosing  
> scope) of the function:
>
> lst = []
> for i in range(10):
>    @localize
>    def f(): print i
>    lst.append(f)
>    lst.append(localize(lambda: i**2))
>
> I'd love to be proven wrong but I don't think localize() can be  
> implemented in current Python.

Here is a (very) quick and dirty implementation in CPython (requires  
ctypes).  I'm sure it breaks in all sorts of ways but I don't have  
more time to test it :)  Tests follow the implementation.

------------------------- localize.py ---------------------
from itertools import *
import sys
import ctypes
from array import array
from opcode import opmap, HAVE_ARGUMENT

new_cell = ctypes.pythonapi.PyCell_New
new_cell.restype = ctypes.py_object
new_cell.argtypes = [ctypes.py_object]

from types import CodeType, FunctionType

LOAD_GLOBAL = opmap['LOAD_GLOBAL']
LOAD_DEREF = opmap['LOAD_DEREF']
LOAD_FAST = opmap['LOAD_FAST']
STORE_GLOBAL = opmap['STORE_GLOBAL']
STORE_DEREF = opmap['STORE_DEREF']
STORE_FAST = opmap['STORE_FAST']

code_args = (
     'argcount', 'nlocals', 'stacksize', 'flags', 'code',
     'consts', 'names', 'varnames', 'filename', 'name',
     'firstlineno', 'lnotab', 'freevars', 'cellvars'
     )

def copy_code(code_obj, **kwargs):
     "Make a copy of a code object, maybe changing some attributes"
     for arg in code_args:
         if not kwargs.has_key(arg):
             kwargs[arg] = getattr(code_obj, 'co_%s' % arg)
     return CodeType(*map(kwargs.__getitem__, code_args))

def code_walker(code):
     l = len(code)
     code = array('B', code)
     i = 0
     while i < l:
         op = code[i]
         if op >= HAVE_ARGUMENT:
             yield op, code[i+1] + (code[i+2] << 8)
             i += 3
         else:
             yield op, None
             i += 1


class CodeMaker(object):
     def __init__(self):
         self.code = array('B')
     def append(self, opcode, arg=None):
         app = self.code.append
         app(opcode)
         if arg is not None:
             app(arg & 0xFF)
             app(arg >> 8)
     def getcode(self):
         return self.code.tostring()


def localize(f):
     if not isinstance(f, FunctionType):
         return f
     nonlocal_vars = []
     new_cells = []
     frame = sys._getframe(1)
     values = dict(frame.f_globals)
     values.update(frame.f_locals)
     co = f.func_code
     deref = co.co_cellvars + co.co_freevars
     names = co.co_names
     varnames = co.co_varnames
     offset = len(deref)
     varindex = {}
     new_code = CodeMaker()

     # Disable CO_NOFREE in the code object's flags
     flags = co.co_flags & (0xFFFF - 0x40)

     # Count the number of arguments of f, including *args & **kwargs
     argcount = co.co_argcount
     if flags & 0x04:
         argcount += 1
     if flags & 0x08:
         argcount += 1

     # Change the code object so that the non local variables are
     # bound to new cells which are initialised to the current value
     # of the variable with that name in the surrounding frame.
     for opcode, arg in code_walker(co.co_code):
         vname = None
         if opcode in (LOAD_GLOBAL, STORE_GLOBAL):
             vname = names[arg]
         elif opcode in (LOAD_DEREF, STORE_DEREF):
             vname = deref[arg]
         else:
             new_code.append(opcode, arg)
             continue
         try:
             vi = varindex[vname]
         except KeyError:
             nonlocal_vars.append(vname)
             new_cells.append(new_cell(values[vname]))
             vi = varindex[vname] = offset
             offset += 1
         if opcode in (LOAD_GLOBAL, LOAD_DEREF):
             new_code.append(LOAD_DEREF, vi)
         else:
             new_code.append(STORE_DEREF, vi)

     co = copy_code(co, code=new_code.getcode(),
                    freevars=co.co_freevars + tuple(nonlocal_vars),
                    flags=flags)
     return FunctionType(co, f.func_globals, f.func_name,  
f.func_defaults,
                         (f.func_closure or ()) + tuple(new_cells))

------------------------ /localize.py ---------------------

Some examples:

 >>> y = 3
 >>> @localize
... def f(x):
...     return x, y
...
 >>> f(5)
(5, 3)
 >>> y = 1000
 >>> f(2)
(2, 3)
 >>> def test():
...     acc = []
...     for i in range(10):
...         @localize
...         def pr(): print i
...         acc.append(pr)
...     return acc
...
 >>> for f in test(): f()
...
0
1
2
3
4
5
6
7
8
9
 >>> lambdas = [localize(lambda: i) for i in range(10)]
 >>> for f in lambdas: print f()
...
0
1
2
3
4
5
6
7
8
9
 >>> # Lastly, your example
 >>> lst = []
 >>> for i in range(10):
...    @localize
...    def f(): print i
...    lst.append(f)
...    lst.append(localize(lambda: i**2))
...
 >>> [f() for f in lst]
0
1
2
3
4
5
6
7
8
9
[None, 0, None, 1, None, 4, None, 9, None, 16, None, 25, None, 36,  
None, 49, None, 64, None, 81]
 >>>

-- 
Arnaud