Search code examples
pythonpython-decorators

Python: How to decorate a special (dunder) method


Wrapping a special method works, but doesn't have the desired effect on the behavior of the instance.

For example, decorating the a.__call__ method (of an instance a) will indeed take effect if I call a.__call__(x), but not if I call a(x).

Consider the following function that makes a decorator that preprocesses the input:

def input_wrap_decorator(preprocess):
    def decorator(func):
        def func_wrapper(*args, **kwargs):
            return func(preprocess(*args, **kwargs))
        return func_wrapper
    return decorator

Consider this simple class:

class A:
    def __call__(self, k):
        return "{}({})".format(self.__class__.__name__, k)

Demo of its amazing functionality:

>>> a = A()
>>> a(7)
'A(7)'

Now say I want to do something critical: Multiply all inputs to __call__ by 10, using input_wrap_decorator. Here's what happens:

>>> a = A()
>>> a.__call__ = input_wrap_decorator(preprocess=lambda x: x * 10)(a.__call__)
>>> a.__call__(7)  # __call__ works as expected
'A(70)'
>>> a(7)  # but a(.) does not!
'A(7)'

Something obscure is happening that only a python grown-up would know...


Solution

  • As stated in Special method lookup,

    For custom classes, implicit invocations of special methods are only guaranteed to work correctly if defined on an object’s type, not in the object’s instance dictionary

    So, you could do it like this:

    def input_wrap_decorator(preprocess):
        def decorator(func):
            def func_wrapper(self, *args, **kwargs):
                return func(self, preprocess(*args, **kwargs))
            return func_wrapper
        return decorator
    
    class A:
        def __call__(self, k):
            return "{}({})".format(self.__class__.__name__, k)
    
    a = A()
    
    # A.__call__ will be used by a(7), not a.__call__
    A.__call__ = input_wrap_decorator(preprocess=lambda x: x * 10)(A.__call__)
    
    print(a.__call__(7))
    # A(70)
    print(a(7))
    # A(70)
    

    Note that I isolated self in func_wrapper, so that it doesn't get passed to preprocess with the other args.

    And of course, you can use the syntactic sugar for decorators:

    class A:
        @input_wrap_decorator(preprocess=lambda x: x * 10)
        def __call__(self, k):
            return "{}({})".format(self.__class__.__name__, k)
    
    a = A()    
    print(a.__call__(7))
    # A(70)
    print(a(7))
    # A(70)