CODEWITHSUNDEEP
pythondecorator
user3758232
user3758232

Reputation: 862

Avoiding double wrapping in a decorator

I have a Python 3 script similar to the following:

def transaction(write=False):
    def _transaction_deco(fn):
        @wraps(fn)
        def _wrapper(*args, **kwargs):
            env.timestamp = arrow.utcnow()
            with TxnManager(write=write) as txn:
                ret = fn(*args, **kwargs)
            delattr(env, 'timestamp')
            return ret
        return _wrapper
    return _transaction_deco

@transaction(True)
def fn1(x=False):
    if x:
        return fn2()
    do_something()

@transaction(True)
def fn2():
    do_something_else()

fn1 and fn2 can be called independently, and both are wrapped in the transaction decorator.

Also, fn1 can call fn2 under certain circumstances. This is where I have a problem, because the decorator code is called twice, transactions are nested, and the timestamp is set twice and deleted twice (actually, the second time delattr(env, 'timestamp') is called, an exception is raised because timestamp is no longer there).

I want whatever is decorated with transaction to only be decorated if not already so.

I tried to modify my code based on hints from a separate Q&A:

def transaction(write=False):
    def _transaction_deco(fn):
        # Avoid double wrap.
        if getattr(fn, 'in_txn', False):
            return fn(*args, **kwargs)

        @wraps(fn)
        def _wrapper(*args, **kwargs):
            # Mark transaction begin timestamp.
            env.timestamp = arrow.utcnow()
            with TxnManager(write=write) as txn:
                ret = fn(*args, **kwargs)
            delattr(env, 'timestamp')
            return ret
        _wrapper.in_txn = True
        return _wrapper

    return _transaction_deco

I.e. if the _in_txn attribute is set for the function, return the bare function because we are already in a transaction.

However, if I run fn1(True) with a debugger, I see that the check for in_txn is never true because the _wrapper variable is set to f1 for the outer call, and to f2 for the inner call, so the functions are being both decorated separately.

I could set the timestamp variables as properties of the context manager, but I still end up with two bested context managers, and if both trnsactions are R/W, Bad Things will happen.

Can someone suggest a better solution?

Edit: Based on the replied received, I believe I omitted a key element: fn1 and fn2 are API metods and must always run within a transaction. The reason for using a decorator is to prevent the API implementer from having to make a decision about the transaction, or having to wrap the methods manually in a context manager or a decorator. However I would be in favor of a non-decorator approach as long as it keeps the implementation of fn1 and fn2 plain.

Upvotes: 1

Views: 698

Answers (3)

chepner
chepner

Reputation: 531345

Rather than a decorator, I'd just make another context manager to use explicitly.

@contextlib.contextmanager
def MyTxnManager(write):
    env.timestamp = arrow.utcnow()
    with TxnManager(write=write) as txn:
        yield txn
    delattr(env, 'timestamp')

with MyTxnManager(True) as txn:
    fn1()

Here, instead of hiding the transaction inside fn1 and fn2, we wrap just the timestamp management inside a new transaction manager. This way, calls to fn1 and fn2 occur in a transaction only if we explicitly cause them to do so.

Upvotes: 4

Mad Physicist
Mad Physicist

Reputation: 114330

You can always just check if the environment has a timestamp already and not create it in that case. It is using the already-existing attribute as a flag instead of making a new one:

def transaction(write=False):
    def _transaction_deco(fn):
        @wraps(fn)
        def _wrapper(*args, **kwargs):
            if hasattr(env, 'timestamp'):
                ret = fn(*args, **kwargs)
            else:
                env.timestamp = arrow.utcnow()
                with TxnManager(write=write) as txn:
                    ret = fn(*args, **kwargs)
                delattr(env, 'timestamp')
            return ret
        return _wrapper
    return _transaction_deco

This will let you nest an arbitrary number of transacted calls, and only use the timestamp and context manager of the outermost one.

Upvotes: 0

jonrsharpe
jonrsharpe

Reputation: 122061

There's no need to handle this in the decorator. It's much simpler to separate out the wrapped and unwrapped versions of fn2, e.g.:

def _fn2():
    do_something_else()

@transaction(True)
def fn1(x=False):
    if x:
        return _fn2()
    do_something()

fn2 = transaction(True)(_fn2)

The decorator syntax is just syntactic sugar, so the result will be exactly the same.

Upvotes: 2

Related Questions