How to implement Python decorator with arguments as a class?

Question

I'm trying to implement a decorator which accepts some arguments. Usually decorators with arguments are implemented as double-nested closures, like this:

def mydecorator(param1, param2):
    # do something with params
    def wrapper(fn):
        def actual_decorator(actual_func_arg1, actual_func_arg2):
            print("I'm decorated!")

            return fn(actual_func_arg1, actual_func_arg2)

        return actual_decorator

    return wrapper

But personally I don't like such approach because it is very unreadable and difficult to understand.

So I ended up with this:

class jsonschema_validate(object):
    def __init__(self, schema):
        self._schema = schema

    def __call__(self, fn):
        self._fn = fn

        return self._decorator

    def _decorator(self, req, resp, *args, **kwargs):
        try:
            jsonschema.validate(req.media, self._schema, format_checker=jsonschema.FormatChecker())
        except jsonschema.ValidationError as e:
            _log.exception('Validation failed: %r', e)

            raise errors.HTTPBadRequest('Bad request')

        return self._fn(req, resp, *args, **kwargs)

The idea is very simple: at instantiation time we just captures decorator args, and at call time we capture decorated function and return decorator instance's method, which is bound. It is important it to be bound because at decorator's invocation time we want to access self with all information stored in it.

Then we use it on some class:

class MyResource(object):
    @jsonschema_validate(my_resource_schema)
    def on_post(self, req, resp):
        pass

Unfortunately, this approach doesn't work. The problem is that at decorator invocation time we looses context of decorated instance because at decoration time (when defining class) decorated method is not bound. Binding occurs later at attribute access time. But at this moment we already have decorator's bound method (jsonschema_validate._decorator) and self is passed implicitly, and it's value isn't MyResource instance, rather jsonschema_validate instance. And we don't want to loose this self value because we want to access it's attributes at decorator invocation time. In the end it results in TypeError when calling self._fn(req, resp, *args, **kwargs) with complains that "required positional argument 'resp' is missing" because passed in req arg becomes MyResource.on_post "self" and all arguments effectively "shifts".

So, is there a way implement decorator as a class rather than as a bunch of nested functions?

Note

As my first attempt of implementing decorator as simple class was failed rather quickly, I immediately reverted to nested functions. It seems like properly implemented class approach is even more unreadable and tangled, but I want to find solution anyway for the fun of the thing.

UPDATE

Finally found solution, see my own answer.

Jacinator · Accepted Answer · 2019-09-19 19:28:55Z

This is fun! Thanks for posting this question.

Writing a simple decorator that doesn't take arguments is pretty easy, but extending that to a class that then gets called three times is a bit more challenging. I opted to use a functools.partial to solve this problem.

from functools import partial, update_wrapper
from unittest import TestCase, main


class SimpleDecorator(object):

    def __new__(cls, func, **params):
        self = super(SimpleDecorator, cls).__new__(cls)
        self.func = func
        self.params = params
        return update_wrapper(self, func)

    def __call__(self, *args, **kwargs):
        args, kwargs = self.before(*args, **kwargs)
        return self.after(self.func(*args, **kwargs))

    def after(self, value):
        return value

    def before(self, *args, **kwargs):
        return args, kwargs


class ParamsDecorator(SimpleDecorator):

    def __new__(cls, **params):
        return partial(super(ParamsDecorator, cls).__new__, cls, **params)


class DecoratorTestCase(TestCase):

    def test_simple_decorator(self):
        class TestSimpleDecorator(SimpleDecorator):

            def after(self, value):
                value *= 2
                return super().after(value)

        @TestSimpleDecorator
        def _test_simple_decorator(value):
            """Test simple decorator"""
            return value + 1

        self.assertEqual(_test_simple_decorator.__name__, '_test_simple_decorator')
        self.assertEqual(_test_simple_decorator.__doc__, 'Test simple decorator')
        self.assertEqual(_test_simple_decorator(1), 4)

    def test_params_decorator(self):
        class TestParamsDecorator(ParamsDecorator):

            def before(self, value, **kwargs):
                value *= self.params['factor']
                return super().before(value, **kwargs)

        @TestParamsDecorator(factor=3)
        def _test_params_decorator(value):
            """Test params decorator"""
            return value + 1

        self.assertEqual(_test_params_decorator.__name__, '_test_params_decorator')
        self.assertEqual(_test_params_decorator.__doc__, 'Test params decorator')
        self.assertEqual(_test_params_decorator(2), 7)

As you can see I've opted for a design with hooks for modifying the arguments and responses in methods. Hopefully, most of the time this would prevent needing to touch __call__ or __new__.

I couldn't think of a way to attach params to ParamsDecorator after returning the partial, so I had to opt for putting it into the SimpleDecorator but not using it.

I think that this does a good job of keeping the content flat instead of nested. I also like that this can take care of functools.wraps for you so you shouldn't need to worry about including that on these decorators. The downside to writing a decorator this way is you're now introducing a new module that you would need to install or maintain and then import every time you write a decorator.

Anthony · Accepted Answer · 2020-09-10 13:24:31Z

Finally got it!

As I wrote, the problem that a method can't have two self, so we need to capture both values in some way. Descriptors and closures to the rescue!

Here is complete example:

class decorator_with_args(object):
    def __init__(self, arg):
        self._arg = arg

    def __call__(self, fn):
        self._fn = fn

        return self

    def __get__(self, instance, owner):
        if instance is None:
            return self

        def _decorator(self_, *args, **kwargs):
            print(f'decorated! arg: {self._arg}')

            return self._fn(self_, *args, **kwargs)

        return _decorator.__get__(instance, owner)

Let's break it down to pieces!

It starts exactly as my previous attempt. In __init__ we just capture decorator arguments to it's private attribute(s).

Things get more interested in next part: a __call__ method.

def __call__(self, fn):
    self._fn = fn

    return self

As before, we capture decorated method to decorator's private attribute. But then, instead of returning actual decorator method (def _decorator in previous example), we return self. So decorated method becomes instance of decorator. This is required to allow it to act as descriptor. According to docs:

a descriptor is an object attribute with "binding behavior"

Confusing, uh? Actually, it is easier than it looks. Descriptor is just an object with "magic" (dunder) methods which is assigned to another's object attribute. When you try to access this attribute, those dunder methods will be invoked with some calling convention. And we'll return to "binding behavior" a little bit later.

Let's look at the details.

def __get__(self, instance, owner):

Descriptor must implement at least __get__ dunder (and __set__ & __delete__ optionally). This is called "descriptor protocol" (similar to "context manager protocol", "collection protocol" an so on).

    if instance is None:
        return self

This is by convention. When descriptor accessed on class rather than instance, it should return itself.

Next part is most interesting.

        def _decorator(self_, *args, **kwargs):
            print(f'decorated! arg: {self._arg}')

            return self._fn(self_, *args, **kwargs)

        return _decorator.__get__(instance, owner)

We need to capture decorator's self as well as decorated instance's self in some way. As we can't define function with two self (even if we can, Python couldn't understand us), so we enclose decorator's self with closure - an inner function. In this closure, we actually do alter behavior of decorated method (print('decorated! arg: {self._arg}')) and then call original one. Again, as there is already argument named self, we need to choose another name for instance's self - in this example I named it self_, but actually it is self' - "self prime" (kinda math humor).

        return _decorator.__get__(instance, owner)

And finally, usually, when we define closures, we just return it: def inner(): pass; return inner. But here we can't do that. Because of "binding behavior". What we need is returned closure to be bound to decorated instance in order it to work properly. Let me explain with an example.

class Foo(object):
    def foo(self):
        print(self)

Foo.foo
# <function Foo.foo at 0x7f5b1f56dcb0>
Foo().foo
# <bound method Foo.foo of <__main__.Foo object at 0x7f5b1f586910>>

When you access method on class, it is just a plain Python function. What makes it a method instead is binding. Binding is an act of linking object's methods with instance which is passed implicitly as firs argument. By convention, it is called self, but roughly speaking this is not required. You can even store method in other variable and call it, and will still have reference to instance:

f = Foo()
f.foo()
# <__main__.Foo object at 0x7f5b1f5868d0>
other_foo = f.foo
other_foo()
# <__main__.Foo object at 0x7f5b1f5868d0>

So, we need to bind our returned closure to decorated instance. How to do that? Remember when we was looking at method? It could be something like that:

# <bound method Foo.foo of <__main__.Foo object at 0x7f5b1f586910>>

Let's look at it's type:

type(f.foo)
# <class 'method'>

Wow! It actually even a class! Let's create it!

method()
# Traceback (most recent call last):
#  File "<stdin>", line 1, in <module>
# NameError: name 'method' is not defined

Unfortunately, we can't do that directly. But, there is types.MethodType:

types.MethodType
# <class 'method'>

Seems like we finally found what we wanted! But, actually, we don't need to create method manually!. All we need to do is to delegate to standard mechanics of method creation. And this is how actually methods work in Python - they are just descriptors which bind themselves to an instances when accessed as instance's attribute!

To support method calls, functions include the __get__() method for binding methods during attribute access.

So, we need to just delegate binding machinery to function itself:

_decorator.__get__(instance, owner)

and get method with right binding!

Collectives™ on Stack Overflow

How to implement Python decorator with arguments as a class?

Note

UPDATE

2 Answers 2

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

Note

UPDATE

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related