Using Decorators to Add Arguments

I have a number of functions that all do very different things, but they all create and close a “context”. So the pattern is something like this:

def some_func(arg1, arg2):
    ctx = Context()
    try:
        # do something with ctx that could raise an exception
    finally:
        ctx.close()

I’d rather avoid the code duplication, and the danger of forgetting that try/finally block and the close(). The most straight-forward way I could see to achieve this is to use a decorator which will ensure the creation and close of the context. So something like this:

def ensure_ctx(func):
    def decorate(*args, **kw):
        ctx = Context()
        try:
            func(ctx, *args, **kw)
        finally:
            ctx.close()
 
    decorate.__name__ = func.__name__
    decorate.__dict__ = func.__dict__
    decorate.__doc__ = func.__doc__
 
    return decorate

I don’t think I have ever seen a decorator that adds arguments like this. While this solves the immediate problem I outlined, it creates other problems. Of course I will need to change all wrapped functions, but that is to be expected. But the decorator also forces a change to the wrapped function signature which I don’t like. Also, just by reading the function signature you would think you need to pass it a context, which would be incorrect since the decorator does that. This point could be fixed by making the decorator attach the context to the func object, but I don’t like that approach either because it makes writing and reading the wrapped callables weird; you’d have to go read the decorator source to understand what is going on. The best solution so far seems to be to add new optional ctx=None argument to the callables, and modify the decorator to inspect the kw dictionary for presence of ctx and create and close one only if it isn’t there. This makes it possible to call the original callable with a context if you want to reuse the context, while also allowing you to call it without a context thereby letting the decorator take care of it. Here’s what all that that could look like:

def ensure_ctx(func):
    def decorate(*args, **kw):
        close = False
        if not kw.get('ctx'):
            kw['ctx'] = Context()
            close = True
        try:
            func(*args, **kw)
        finally:
            if close:
                kw['ctx'].close()
                del kw['ctx']
 
    decorate.__name__ = func.__name__
    decorate.__dict__ = func.__dict__
    decorate.__doc__ = func.__doc__
 
    return decorate
 
@ensure_ctx
def some_func(arg1, arg2, ctx=None):
    # do something with ctx that could raise an exception

I still don’t especially like it, but I can’t think of anything better right now (besides using the decorator module).

Similar Posts:

    None Found

7 Comments

  1. Brian Zimmer:

    I think the solution you want is the “with” statement available in Python 2.5+. I thought this an interesting post so I responded with more details than I felt appropriate in a comment box.

  2. Ornitorrinco Enmascarado:

    I agree with Brian, the “with” statment is a more natural solution.

  3. Christof:

    Not sure you dismissed this option in your article above but maybe just add ctx to the func in the decorator “func.ctx = Context()“? Usage in the some_func is “some_func.ctx“ but with the decorator before the “def some_func“ it makes this kind of obvious where .ctx may come from.

    Not sure how to format code in this comment so I put the code here: http://pastebin.com/f6fd359ed

    Guess usage of “with“ may even be better but this seems to work and feels logical in usage I think.

  4. Love Encounter Flow:

    i agree with the foregoing posters that this looks a lot like a with statement use case. you may want to consider using contextlib/@contextmanager to simplify the definition of your context handler. as concerns the implicit introduction of local variables (which you are solving by putting the affected names into the method signature)—i’ve tried it before, and i believe it can only be done by recompiling the method source code (i.e. you would have to use inspection to grab the text from the *.py file, then do a superficial parse and add lines that introduce your new variables, exec the source text, and put the result back into the class or module from whence it came). the only other possible solution i’m aware of is fiddling with the bytecodes, which may or may not be easier. this complication is a sad side-effect of python’s strategy to consider a name local when there is a local assignment without a global statement and vice versa.

  5. Marius Gedminas:

    If you write decorators, you really want to set the __module__ of the inner function too. Things like doctests rely on it. Best use the functools.wraps():

    Here goes a massively-blog-comment-distorted-I-expect example:

    def ensure_ctx(func):
        @functools.wraps(func)
        def decorate(*args, **kw):
            ...
        return decorate
    
  6. Heikki Toivonen:

    @Brian: Thanks, I will definitely need to take a more careful look of “with”. I dismissed that approach too early I guess, since I haven’t actually used “with” much before (just for locks that already had the needed double underscore methods).

    @Cristof: I did consider it (“attach the context to the func object”), but I didn’t like how it looked.

    @Love Encounter Flow: I dismissed the inspection approach as being too kludgy.

    @Marius: Ah, thanks, I wasn’t aware of that.

  7. Phillip J. Eby:

    Note, by the way, that you can always design your context objects so that the only way to use them is via “with”: e.g., “with Context() as ctx:”, where Context.__enter__ returns the real context object, and the Context() itself is just a dummy class that raises errors if you try to do anything with it, calls close() on __exit__.