Capture functions within a closure scope

Question

Python3 added a __prepare__ so that you could substitute the type of dictionary used to gather items from a class declaration (see here.) Using __prepare__ I can set up the class to allow multiple definitions of the same member function.

class MultiType(type):

    @classmethod
    def __prepare__(metacls, name, bases, **kwds):
        return collections.defaultdict(list)


class Foo(metaclass=MultiType):
    def bar(self, x):
        return x - 1

    def bar(self, x):
        return x

At the module level, I could play some tricks using a decorator:

def capture(f):
    module = importlib.import_module(f.__module__)
    if not hasattr(module, 'my_cache'):
        setattr(module, 'my_cache', [])
    module.my_cache.append(f)

    @functools.wraps(f)
    def _wrapper(*args, **kwargs):
        return (func(*args, **kwargs) for func in module.my_cache)
    return _wrapper

@capture
def foo(x):
    return x - 1

@capture
def foo(x):
    return 42

However, if I do this within a closure, I could potentially add things scoped at the module level which should not be.

def foo(x):
    @some_awesome_deco
    def bar(y):
        return y

    @some_awesome_deco
    def bar(y):
        return 24

    return bar(x+1)

Is there a way to identify and capture the functions declared within the scope of a closure so that I can handle these differently than functions declared at module scope without having to reference the closed over function (i.e. @some_awesome_deco(foo)?)

Answer 1

If all you need to support is CPython, then your decorator could look at the sys._getframe(1) frame object, which represents the execution frame for the code executing the decorator. If the frame.f_locals dictionary is the same object as the frame.f_globals dictionary you are at module level. If not, you are in a nested scope.

You'll have to generate some kind of scope key however; you may get away with storing something in f_locals (which won't actually affect the actual locals). Just remember that locals (as well as the frame) are cleared when a function exits. I'd return a special callable instead, one that is mutable, so you can refer to it on subsequent decorator calls. You'd be able to retrieve that object with frame.f_locals[decorated_function.__name__], for example.

See the inspect module documenation for an overview of what attributes you can expect to find on a frame object.

Demo:

>>> import sys
>>> def nested_scope_detector(func):
...     frame = sys._getframe(1)
...     nested_scope = frame.f_globals is not frame.f_locals
...     redefinition = func.__name__ in frame.f_locals
...     if nested_scope: print('{!r} is located in a nested scope.'.format(func))
...     if redefinition: print('Redefining {!r}, previously bound to {!r}'.format(
...         func.__name__, frame.f_locals[func.__name__]))
...     return func
... 
>>> @nested_scope_detector
... def foo(): pass
... 
>>> @nested_scope_detector
... def foo(): pass
... 
Redefining 'foo', previously bound to <function foo at 0x10e931d08>
>>> def bar():
...     @nested_scope_detector
...     def foo(): pass
...     @nested_scope_detector
...     def foo(): pass
... 
>>> bar()
<function bar.<locals>.foo at 0x10eb4ef28> is located in a nested scope.
<function bar.<locals>.foo at 0x10eb4eea0> is located in a nested scope.
Redefining 'foo', previously bound to <function bar.<locals>.foo at 0x10eb4ef28>

As such, you could use a function attribute on the returned wrapper function to store your functions:

def capture(f):
    locals = sys._getframe(1).f_locals
    preexisting = locals.get(f.__name__)
    if preexisting is not None and hasattr(preexisting, 'functions'):
        preexisting.functions.append(f)
        return preexisting

    @functools.wraps(f)
    def _wrapper(*args, **kwargs):
        return (func(*args, **kwargs) for func in _wrapper.functions)
    _wrapper.functions = [f]
    return _wrapper

and it'll work in any scope:

>>> @capture
... def foo(): return 'bar'
... 
>>> @capture
... def foo(): return 'baz'
... 
>>> foo()
<generator object <genexpr> at 0x10eb45ee8>
>>> list(foo())
['bar', 'baz']
>>> def bar():
...     @capture
...     def foo(): return 'bar'
...     @capture
...     def foo(): return 'baz'
...     return foo
... 
>>> list(bar()())
['bar', 'baz']