CODEWITHSUNDEEP
python
Nijan
Nijan

Reputation: 638

How to pass an argument to a method decorator

I have a method decorator like this.

class MyClass:
    def __init__(self):
        self.start = 0

    class Decorator:
        def __init__(self, f):
            self.f = f
            self.msg = msg

        def __get__(self, instance, _):
            def wrapper(test):
                print(self.msg)
                print(instance.start)    
                self.f(instance, test)
                return self.f
            return wrapper

    @Decorator
    def p1(self, sent):
        print(sent)

c = MyClass()
c.p1('test')

This works fine. However, If I want to pass an argument to the decorator, the method is no longer passed as an argument, and I get this error:

TypeError: init() missing 1 required positional argument: 'f'

class MyClass:
    def __init__(self):
        self.start = 0

    class Decorator:
        def __init__(self, f, msg):
            self.f = f
            self.msg = msg

        def __get__(self, instance, _):
            def wrapper(test):
                print(self.msg)
                print(instance.start)    
                self.f(instance, test)
                return self.f
            return wrapper

    @Decorator(msg='p1')
    def p1(self, sent):
        print(sent)

    @Decorator(msg='p2')
    def p2(self, sent):
        print(sent)

How do I pass an argument to the decorator class, and why is it overriding the method?

Upvotes: 3

Views: 526

Answers (3)

mementum
mementum

Reputation: 3203

A decorator will be called.

In your case you receive the function as a parameter in the __call__ method

class MyClass:
    def __init__(self):
        self.start = 0

    class Decorator:
        def __init__(self, msg):
            self.msg = msg

        def __call__(self, f):
            self.f = f
            return self

        def __get__(self, instance, _):
            def wrapper(test):
                print(self.msg)
                self.f(instance, test)
                return self.f
            return wrapper

    @Decorator(msg='p1')
    def p1(self, sent):
        print(sent)

    @Decorator(msg='p2')
    def p2(self, sent):
        print(sent)

Your first example works because calling the Class creates an instance and the function is the parameter.

But in your second example you call the Class manually to set the msg parameter, so you the decoration process calls what's left, i.e.: the instance and that goes to the __call__ method.

Upvotes: 2

Blckknght
Blckknght

Reputation: 104752

When you call a decorator with arguments, the function you call isn't actually working as a decorator itself. Rather, it's a decorator factory (a function or other callable that returns something that will act as the decorator). Usually you solve this by adding an extra level of nested functions. Since you're defining your decorator with a class, that's a bit awkward to do directly (though you probably could make it work). But there doesn't really seem to be any need for your decorator to be a class, as long as you handle self in the wrapper function (it will be the instance of MyClass now, rather than an instance of a Decorator class):

class MyClass:
    def __init__(self):
        self.start = 0

    def decorator_factory(msg):
        def decorator(f):
            def wrapper(self, test): # you might want to replace test with *args and **kwargs
                print(msg)
                print(self.start)
                return f(self, test)
            return wrapper
        return decorator

    @decorator_factory(msg='p1')
    def p1(self, sent):
        print(sent)

    @decorator_factory(msg='p2')
    def p2(self, sent):
        print(sent)

I named the decorator factory the way I did to be explicit about the different levels of nested functions, but you should of course use something that's actually meaningful for your use case as the top level name. You might also want to move it out of the class namespace, since it will be available to call on all instances of MyClass (with possibly silly results, since it's not intended to be a method).

Upvotes: 1

user2390182
user2390182

Reputation: 73470

The descriptor protocol doesn't serve much of a purpose here. You can simply pass the function itself to __call__ and return the wrapper function without losing access to the instance:

class MyClass:
    def __init__(self):
        self.start = 0

    class Decorator:
        def __init__(self, msg):
            self.msg = msg

        def __call__(self, f):
            def wrapper(instance, *args, **kwargs):
                print(self.msg)
                # access any other instance attributes
                return f(instance, *args, **kwargs)
            return wrapper

    @Decorator(msg='p1')
    def p1(self, sent):
        print(sent)

>>> c = MyClass()
>>> c.p1('test')
p1
test

Upvotes: 2

Related Questions