Preventing function (or decorator) from being nested

Question

I've got some code in a decorator that I only want run once. Many other functions (utility and otherwise) will be called later down the line, and I want to ensure that other functions that may have this decorator aren't accidentally used way down in the nest of function calls.

I also want to be able to check, at any point, whether or not the current code has been wrapped in the decorator or not.

I've written this, but I just wanted to see if anyone else can think of a better/more elegant solution than checking for the (hopefully!) unique function name in the stack.

import inspect

def my_special_wrapper(fn):
    def my_special_wrapper(*args, **kwargs):
        """ Do some magic, only once! """
        # Check we've not done this before
        for frame in inspect.stack()[1:]:  # get stack, ignoring current!
            if frame[3] == 'my_special_wrapper':
                raise StandardError('Special wrapper cannot be nested')
        # Do magic then call fn
        # ...
        fn(*args, **kwargs)
    return my_special_wrapper

def within_special_wrapper():
    """ Helper to check that the function has been specially wrapped """
    for frame in inspect.stack():
        if frame[3] == 'my_special_wrapper':
            return True
    return False

@my_special_wrapper
def foo():
    print within_special_wrapper()
    bar()
    print 'Success!'

@my_special_wrapper    
def bar():
    pass

foo()

Answer 1

Here is an example of using a global for this task - in what I believe is a relatively safe way:

from contextlib import contextmanager
from functools import wraps

_within_special_context = False

@contextmanager
def flag():
    global _within_special_context
    _within_special_context = True
    try:
        yield
    finally:
        _within_special_context = False


#I'd argue this would be best replaced by just checking the variable, but
#included for completeness.
def within_special_wrapper():
    return _within_special_context


def my_special_wrapper(f):
    @wraps(f)
    def internal(*args, **kwargs):
        if not _within_special_context:
            with flag():
                ...
                f(*args, **kwargs)
        else:
            raise Exception("No nested calls!")
    return internal

@my_special_wrapper
def foo():
    print(within_special_wrapper())
    bar()
    print('Success!')

@my_special_wrapper
def bar():
    pass

foo()

Which results in:

True
Traceback (most recent call last):
  File "/Users/gareth/Development/so/test.py", line 39, in <module>
    foo()
  File "/Users/gareth/Development/so/test.py", line 24, in internal
    f(*args, **kwargs)
  File "/Users/gareth/Development/so/test.py", line 32, in foo
    bar()
  File "/Users/gareth/Development/so/test.py", line 26, in internal
    raise Exception("No nested calls!")
Exception: No nested calls!

Using a context manager ensures that the variable is unset. You could just use try/finally, but if you want to modify the behaviour for different situations, the context manager can be made to be flexible and reusable.

Answer 2

The obvious solution is to have special_wrapper set a global flag, and just skip its magic if the flag is set.

This is about the only good use of a global variable - to allow a single piece of code to store information that is only used within that code, but which needs to survive the life of execution in that code.

It doesn't need to be set in global scope. The function could set the flag on itself, for example, or on any object or class, as long as nothing else will touch it.

As noted by Lattyware in comments, you'll want to use either a try/except, or perhaps even better, a context manager to ensure the variable is unset.

Update: If you need the wrapped code to be able to check if it is wrapped, then provide a function which returns the value of the flag. You might want to wrap it all up with a class for neatness.

Update 2: I see you're doing this for transaction management. There are probably already libraries which do this. I strongly recommend that you at least look at their code.

Answer 3

While my solution technically works, it requires a manual reset of the decorator, but you could very well modify things such that the outermost function is instead a class (with the instances being the wrappers of the decorated functions passed to it in __init__), and have reset() being called in __exit__(), which would then allow you to use the with statement to create the decorator to be usable only once within the context. Also note that it requires Python 3 due to the nonlocal keyword, but that can easily be adapted to 2.7 with a dict in place of the flag variable.

def once_usable(decorator):
    "Apply this decorator function to the decorator you want to be usable only once until it is reset."

    def outer_wrapper():
        flag = False

        def inner_wrapper(*args, **kwargs):
            nonlocal flag
            if not flag:
                flag = True
                return decorator(*args, **kwargs)
            else:
                print("Decorator currently unusable.") # raising an Error also works

        def decorator_reset():
            nonlocal flag
            flag = False

        return (inner_wrapper, decorator_reset)

    return outer_wrapper()

Testing:

>>> def a(aa):
    return aa*2

>>> def b(bb):
    def wrapper(*args, **kwargs):
        print("Decorated.")
        return bb(*args, **kwargs)

    return wrapper

>>> dec, reset = once_usable(b)
>>> aa = dec(a)
>>> aa(22)
Decorated.
44
>>> aaa = dec(a)
Decorator currently unusable.
>>> reset()
>>> aaa = dec(a)
>>> aaa(11)
Decorated.
22