Search code examples
pythoninstancedecoratorpython-decorators

How to implement Python decorator with arguments as a class?


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.


Solution

  • 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!