python - Passing function arguments through stacked decorators

Question

I am trying to write two separate but stackable decorators, one for printing the state of an object before and after the method, and one for running some set of internal class tests after the method (which also has arguments).

Here an example of my current attempt:

import functools

class Dog:

    def __init__(self):
        self.happy = False
        self.has_stick = False

    def __str__(self):
        n = ' ' if self.happy else ' not '
        return "I'm%sa happy dog" % n

    def _verbose(func):
        fname = func.func_name
        argnames = func.func_code.co_varnames[:func.func_code.co_argcount]
        @functools.wraps(func)
        def decorator(*args, **kwargs):
            print "Before %s(%s):" % (fname, ', '.join(
                '%s=%r' % entry
                for entry in zip(argnames, args)[1:] + kwargs.items()))
            print args[0]
            result = func(*args, **kwargs)
            print "After %s(%s):" % (fname, ', '.join(
                '%s=%r' % entry
                for entry in zip(argnames, args)[1:] + kwargs.items()))
            print args[0]
            return result
        return decorator

    def _test(printout):
        def actual_decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                self = args[0]
                output = func(*args, **kwargs)
                self._test_not_happy_without_stick(printout)
                return output
            return wrapper
        return actual_decorator

    def _test_not_happy_without_stick(self, printout):
        if printout:
            print "Is happy:", self.happy
            print "Has stick:", self.has_stick
        if self.happy and not self.has_stick:
            print "ERROR"

    @_test(True)
    @_verbose
    def finds_stick(self, good_stick):
        print "I found a stick!"
        self.happy = good_stick
        self.has_stick = True

if __name__ == '__main__':
    fido = Dog()
    fido.finds_stick(False)

If the ordering of the decorators is applied as above, the output is:

Before finds_stick(good_stick=True):
I'm not a happy dog
I found a stick!
After finds_stick(good_stick=True):
I'm a happy dog
Is happy: True
Has stick: True

However, if it is reversed (as I would like to do), the name and value of the argument passed to the decorated functions is lost, as seen below:

Before finds_stick():
I'm not a happy dog
I found a stick!
Is happy: True
Has stick: True
After finds_stick():
I'm a happy dog

How can I ensure that any such stacking of decorators doesn't prevent arguments being passed through the decorators?

Alternatively, I would be happy for suggestions of a more pythonic way to address the problem.

score 0 · Accepted Answer · answered Oct 11 '17 at 14:08

Turns out the problem was due to the backwards compatibility of @functools.wrap(), which doesn't preserve the signatures to achieve what I needed above (as it does in Python 3.4+).

However, you can get the desired functionality by use of the decorator package, as outlined here.

The correct code for the above is then:

import decorator

class Dog:

    def __init__(self):
        self.happy = False
        self.has_stick = False

    def __str__(self):
        n = ' ' if self.happy else ' not '
        return "I'm%sa happy dog" % n

    @decorator.decorator
    def _verbose(func, *args, **kwargs):
        fname = func.func_name
        argnames = func.func_code.co_varnames[:func.func_code.co_argcount]
        print "Before %s(%s):" % (fname, ', '.join(
            '%s=%r' % entry
            for entry in zip(argnames, args)[1:] + kwargs.items()))
        print args[0]
        result = func(*args, **kwargs)
        print "After %s(%s):" % (fname, ', '.join(
            '%s=%r' % entry
            for entry in zip(argnames, args)[1:] + kwargs.items()))
        print args[0]
        return result

    def _test(printout):
        @decorator.decorator
        def wrapper(func, *args, **kwargs):
            self = args[0]
            output = func(*args, **kwargs)
            self._test_not_happy_without_stick(printout)
            return output
        return wrapper

    def _test_not_happy_without_stick(self, printout):
        if printout:
            print "Is happy:", self.happy
            print "Has stick:", self.has_stick
        if self.happy and not self.has_stick:
            print "ERROR"

    @_verbose
    @_test(True)
    def finds_stick(self, good_stick):
        print "I found a stick!"
        self.happy = good_stick
        self.has_stick = True

if __name__ == '__main__':
    fido = Dog()
    fido.finds_stick(True)

With output:

Before finds_stick(good_stick=True):
I'm not a happy dog
I found a stick!
Is happy: True
Has stick: True
After finds_stick(good_stick=True):
I'm a happy dog

python - Passing function arguments through stacked decorators

1 Answers1