python decorating instance method with parameter changes

Viewed 78

I'm trying to make a decorator which can be used for any methods or functions, and I want to modify the arguments passed to the func at the same time. Here's an example:

import math

def chunk(l, n):
    for i in range(0, len(l), n):
        yield l[i:i + n]

def divide_and_run(func):
    def wrapper(self, entries, *args, **kwargs):
        output = []
        for c in chunk(entries, 3):
            output.append(func(self, c, *args, **kwargs))
        return output
    return wrapper

class Dummy:
    @divide_and_run
    def multiply_all(self, mylist):
        return math.prod(mylist)

a = [1,2,3,4,5,6,7,8,9,10]
d = Dummy()

print(d.multiply_all(a))

The decorator divide_and_run divides mylist, runs func on each chunk, and returns the list of outputs. But this one cannot be used for a free function (no self argument)

I want two things:

  1. Modify the argument passed to func (mylist -> chunk)
  2. Make this decorator usable regardless of the method or function type (e.g., instance method or free function)

What's the best practice? Is decorator a good use when I modify the arguments?

2 Answers

The reason that it's hard to write a decorator that works on both top-level functions and on methods is that you can't tell which argument you want to be modifying will be. In your example, you always modify the entries argument to wrapper, which is the second argument. But if you were running on a top-level function, there would be no self argument, so the argument you'd want to change would be the first one.

One way to solve this is to write a callable descriptor type, rather than making your decorator return a function (which is a descriptor with method-binding behavior). This way you can react differently if you're called normally than if you're called as a method.

Here's a simple implementation that handles methods by creating a new instance on a bound method if it gets looked up as a method.

class divide_and_run:
    def __init__(self, func):
        self.func = func

    def __call__(self, entries, *args, **kwargs):  # non-method calls
        return [self.func(c, *args, **kwargs) for c in chunk(entries, 3)]

    def __get__(self, obj, obj_type):    # handles method lookups on an instance
        return divide_and_run(self.func.__get__(obj, obj_type))

This implementation has one flaw, which is that it doesn't support unbound method lookups on a class (e.g. Dummy.multiply_all(d, a)). That is not impossible to solve, you just need to check that obj is not None in __get__ and do something appropriate if it is. But the "something appropriate" is a bit fiddly, and I decided it would distract from the current code which is very concise.

The types module provides a way to distinguish between method and function.

import types

# ...

def divide_and_run(func):
    if isinstance(func, types.MethodType):
        def wrapper(self, entries, *args, **kwargs):
            output = []
            for c in chunk(entries, 3):
                output.append(func(self, c, *args, **kwargs))
            return output
    elif isinstance(func, types.FunctionType):
        def wrapper(entries, *args, **kwargs):
            output = []
            for c in chunk(entries, 3):
                output.append(func(c, *args, **kwargs))
            return output
    return wrapper

A hacky and less robust way could be

if func.__code__.co_varnames[0] == 'self': # <- hardcoded!
   ...
else:
   ...

The disadvantage is that when declaring the instance method the reference to the object must be done with self identifier


EDIT

There are many many possible ways to implement the "same" decorator. Each is tight together with a specific signature hence on the way you want to call it. I just tried to be more faithful to the requirements and I discuss only two cases.

Summary: divide_and_run_DIY always works; check remark 3. for the reason that divide_and_run_SS is not universal

Remarks:

  1. notation: SS is used for Syntactic Sugar and DIY for Do It Yourself (without the @)

  2. classification of callables. It is quite an hard task and I haven't found a solution yet. It will play an important role for the @-decoration. The following pairs of methods are equivalents types.MethodType and inspect.ismethod, types.FunctionType and inspect.isfunction. Only useful for the DIY solution.

  3. syntactic sugar decoration. I realize that @decoration change the type of the callable, kind of automatic casting, which then make its classification harder for the purpose of the question (see previous remark). Here an example:

2A: The @-decoration cast the decorated function from an instance method to a function:

class Dummy:
    @divide_and_run_SS
    def multiply_all(self, mylist): return math.prod(mylist)

ss = Dummy().multiply_all
print(ss.__closure__[0].cell_contents)
#<function Dummy.multiply_all at 0x7f3ed8be25e0>

2B: The DIY-decoration of the decorated function doesn't change its type:

class Dummy:
    def multiply_all(self, mylist): return math.prod(mylist)

diy = divide_and_run_SS(Dummy().multiply_all)
print(print(diy.__closure__[0].cell_contents))
#<bound method Dummy.multiply_all of <__main__.Dummy object at 0x7f3ed8c4ce80>>

Notice that printing func.__closure__[0].cell_contents has the same effect as adding to the source code a print(func) statement.

  1. two possible answers are presented. divide_and_run_DIY is stable for DIY-decoration (+SS-decoration of function only!). Instead, the other, divide_and_run_SS, is not universal but needs to ad-hoc implementation each time. The causes are mentioned in the previous remarks.

About the issues with divide_and_run_SS:

def divide_and_run_SS(func):
    def wrapper(entries, *args, **kwargs):
        # do smt with func(..., *args, **kwargs)
    return wrapper

class Dummy:
    @divide_and_run_SS
    def multiply_all(self, mylist):
        return math.prod(mylist)

From 2.A: @-decoration change type

multiply_all: bound method --divide_and_run_SS--> multiply_all: function

... and the problems arises:

  • such change of type pass the reference of the instance as a parameter (the usual self) to the wrapper, in this case args[0] will be self
  • the decorator is expected to work both with instance method and function but the @-decoration cast to function so how could the decorator distinguish them? No idea... here I presented an hardcoded solution. A speculative approach is already mentioned in the original answer (the hacky way and it is probably the best tmp-fix).
  1. the post seems long but it's just few lines... I used verbose doc-strings in order to be checked with doctest. I found it very useful to deal with such problem

  2. extension of the problem to Part 2 has exactly the same issues.


Header

import math
import types
#import inspect #   in alternative to types
import doctest  #   check correctness of the code in the comments


def chunk(l, n):
    for i in range(0, len(l), n):
        yield l[i:i + n]

Part 2 Make this decorator usable regardless of the method or function type (e.g., instance method or free function)

DIY-Syntax: divide_and_run_SS(func)(args_func)

def divide_and_run_DIY(func):
    """DIY (Manual) decoration. It works for SS of functions BUT it fails with methods
    >>> a = [1,2,3,4,5,6,7,8,9,10]
    >>> print('DIY (Manual) decoration')
    DIY (Manual) decoration
    >>> def f(mylist): return math.prod(mylist)
    >>> divide_and_run_DIY(f)(a)
    [6, 120, 504, 10]
    >>> # <function f at 0x7f7e47977ee0>
    >>> class Dummy:
    ...     def multiply_all(self, mylist): return math.prod(mylist)
    >>> divide_and_run_DIY(Dummy().multiply_all)(a)
    [6, 120, 504, 10]
    >>> # <bound method Dummy.multiply_all of <__main__.Dummy object at 0x7f7e47939820>>
    >>> print('='*20)
    ====================
    >>> print('SS decoration')
    SS decoration
    >>> print('Test with a function')
    Test with a function
    >>> @divide_and_run_DIY
    ... def f(mylist): return math.prod(mylist)
    >>> f(a)
    [6, 120, 504, 10]
    >>> # <function f at 0x7f8877ab5670>
    >>> print('Test with a method')
    Test with a method
    >>> class Dummy:
    ...     @divide_and_run_DIY
    ...     def multiply_all(self, mylist): return math.prod(mylist)
    >>> Dummy().multiply_all(a)
    Traceback (most recent call last):
    TypeError: object of type 'Dummy' has no len()
    >>> # <function Dummy.multiply_all at 0x7f8877ab5700>
    """
    if isinstance(func, types.MethodType):
        def wrapper(entries, *args, **kwargs):
            output = []
            for c in chunk(entries, 3):
                output.append(func(c, *args, **kwargs))
            return output
    elif isinstance(func, types.FunctionType):
        def wrapper(entries, *args, **kwargs):
            output = []
            for c in chunk(entries, 3):
                output.append(func(c, *args, **kwargs))
            return output
    return wrapper


def divide_and_run_SS(func):
    """syntactic sugar decoration. It works both DIY and SS decoration BUT it is not universal! Check the DANGER comments and remark 3.
    >>> print('Syntactic sugar decoration')
    Syntactic sugar decoration
    >>> a = [1,2,3,4,5,6,7,8,9,10]
    >>> print('instance method')
    instance method
    >>> class Dummy:
    ...     @divide_and_run_SS
    ...     def multiply_all(self, mylist): return math.prod(mylist)
    >>> Dummy().multiply_all(a)
    [6, 120, 504, 10]
    >>> # <function Dummy.multiply_all at 0x7f4c8b5a70d0> # DANGER: it is a function!!
    >>> print('function')
    function
    >>> @divide_and_run_SS
    ... def f(mylist): return math.prod(mylist)
    >>> f(a)
    [6, 120, 504, 10]
    >>> # <function f at 0x7f4c8b58dee0>
    >>> print('='*20)
    ====================
    >>> print('SS (Manual) decoration')
    SS (Manual) decoration
    >>> print('function')
    function
    >>> def f(mylist): return math.prod(mylist)
    >>> divide_and_run_SS(f)(a)
    [6, 120, 504, 10]
    >>> # <function f at 0x7f0725fd5670>
    >>> print('instance method')
    instance method
    >>> class Dummy:
    ...     def multiply_all(self, mylist): return math.prod(mylist)
    >>> divide_and_run_SS(Dummy().multiply_all)(a)
    [6, 120, 504, 10]
    >>> # <bound method Dummy.multiply_all of <__main__.Dummy object at 0x7f0726056f40>>
    """
    def wrapper(entries, *args, **kwargs):
        output = []
        # unnatural fix
        if len(args) > 0: # func <- "method" is a function due to @SS
            # reparametrization
            __self = entries     # <- DANGER
            __entries = args[0]  # <- DANGER
            for c in chunk(__entries, 3):
                output.append(func(__self, c, *args[1:], **kwargs))
        else: # func <- function
            for c in chunk(entries, 3):
                output.append(func(c, *args, **kwargs))
        return output
    return wrapper

Check output with

if __name__ == "__main__":
    # run doctest
    doctest.testmod()

    # print __doc__
    print(divide_and_run_SS.__name__)
    print('#' * len(divide_and_run_SS.__name__))
    print(divide_and_run_SS.__doc__.replace('\t', ''))

    print(divide_and_run_DIY.__name__)
    print('#' * len(divide_and_run_DIY.__name__))
    print(divide_and_run_DIY.__doc__.replace('\t', ''))

Part 2 + Part 1 Modify the argument passed to func (mylist -> chunk)

DIY-Syntax: Dec(f1)(args_f1).divide_and_run_DIY(f2)(args_f2)

class Dec:
    def __init__(self, func):
        self.func = func

    def __call__(self, *func_args, **func_kwargs):
        self.func_args, self.func_kwargs = func_args, func_kwargs
        return self

    def divide_and_run_DIY(self, f_target):
        """DIY (Manual) decoration. It works also for SS of functions BUT it fails with methods
        >>> a = [1,2,3,4,5,6,7,8,9,10]
        >>> print('DIY (Manual) decoration')
        DIY (Manual) decoration
        >>> print('Test with a function')
        Test with a function
        >>> def f(mylist): return math.prod(mylist)
        >>> Dec(chunk)(3).divide_and_run_DIY(f)(a)
        [6, 120, 504, 10]
        >>> # <function f at 0x7fa1fe6785e0>
        >>> print('Test with a method')
        Test with a method
        >>> class Dummy:
        ...     def multiply_all(self, mylist): return math.prod(mylist)
        >>> Dec(chunk)(3).divide_and_run_DIY(Dummy().multiply_all)(a)
        [6, 120, 504, 10]
        >>> # <bound method Dummy.multiply_all of <__main__.Dummy object at 0x7fa1fe6f7df0>>
        >>> print('='*20)
        ====================
        >>> print('SS decoration')
        SS decoration
        >>> print('Test with a function')
        Test with a function
        >>> @Dec(chunk)(3).divide_and_run_DIY
        ... def f(mylist): return math.prod(mylist)
        >>> f(a)
        [6, 120, 504, 10]
        >>> # <function f at 0x7fa1fe678940>
        >>> print('Test with a method')
        Test with a method
        >>> class Dummy:
        ...     @Dec(chunk)(3).divide_and_run_DIY
        ...     def multiply_all(self, mylist): return math.prod(mylist)
        >>> Dummy().multiply_all(a)
        Traceback (most recent call last):
        TypeError: object of type 'Dummy' has no len()
        >>> # <function Dummy.multiply_all at 0x7fa1fe678af0> # <-- it become a function!
        """
        if isinstance(f_target, types.MethodType):
            def wrapper(entries, *args, **kwargs):
                output = []
                for c in self.func(entries, *self.func_args, **self.func_kwargs):
                    output.append(f_target(c, *args, **kwargs))
                return output
        return wrapper

    def divide_and_run_SS(self, f_target):
        """Not universal solution, check the DANGER comments and remark 3.
        >>> print('Syntactic sugar decoration')
        Syntactic sugar decoration
        >>> print('function')
        function
        >>> @Dec(chunk)(3).divide_and_run_SS
        ... def f(mylist): return math.prod(mylist)
        >>> f(a)
        [6, 120, 504, 10]
        >>> # <function f at 0x7fd4de72d550>
        >>> print('instance method')
        instance method
        >>> class Dummy:
        ...     @Dec(chunk)(3).divide_and_run_SS
        ...     def multiply_all(self, mylist): return math.prod(mylist)
        >>> Dummy().multiply_all(a)
        [6, 120, 504, 10]
        >>> # <function Dummy.multiply_all at 0x7fd4de72da60>
        >>> print('='*20)
        ====================
        >>> print('SS (Manual) decoration')
        SS (Manual) decoration
        >>> print('function')
        function
        >>> def f(mylist): return math.prod(mylist)
        >>> Dec(chunk)(3).divide_and_run_SS(f)(a)
        [6, 120, 504, 10]
        >>> # <function f at 0x7f1221206550>
        >>> print('instance method')
        instance method
        >>> class Dummy:
        ...     def multiply_all(self, mylist): return math.prod(mylist)
        >>> Dec(chunk)(3).divide_and_run_SS(Dummy().multiply_all)(a)
        [6, 120, 504, 10]
        >>> # <bound method Dummy.multiply_all of <__main__.Dummy object at 0x7f1221284100>>
        """
        def wrapper(entries, *args, **kwargs):
            output = []
            # unnatural fix
            if len(args) > 0: # func <- "method" is a function due to @SS
                # reparametrization
                __self = entries    # <- DANGER
                __entries = args[0] # <- DANGER
                for c in self.func(__entries, *self.func_args, **self.func_kwargs):
                    output.append(f_target(__self, c, *args[1:], **kwargs))
            else: # func <- function
                for c in self.func(entries, *self.func_args, **self.func_kwargs):
                    output.append(f_target(c, *args, **kwargs))
            return output
        return wrapper

Check output with

if __name__ == "__main__":
    # run doctest
    doctest.testmod()

    # print __doc__
    print(Dec.divide_and_run_SS.__name__)
    print('#' * len(Dec.divide_and_run_SS.__name__))
    print(Dec.divide_and_run_SS.__doc__.replace('\t', ''))

    print(Dec.divide_and_run_DIY.__name__)
    print('#' * len(Dec.divide_and_run_DIY.__name__))
    print(Dec.divide_and_run_DIY.__doc__.replace('\t', ''))

EDIT: divide_and_run_SS_hack is the re-implementation of divide_and_run_SS. It is much more robust, still lake of universality: it depends only on the identifier chosen to refer to the instance of a decorated method. Default is 'self' string. It could (not tested yet + see answer of Blckknght + remark 1.) be made also compatible with class method, default 'cls', by adding a/some conditional branch/es and extending vs[0] in {'self', 'cls'}.

Here an example of a failure

class Dummy:
    @Dec(chunk)(3).divide_and_run_SS_hack
    def multiply_all(SELF, mylist): return math.prod(mylist) # capital case

Part 1 + Part2: divide_and_run_SS_hack

class Dec:
    # ...

    def divide_and_run_SS_hack(self, f_target):
        """SS decoration hacky way. For method: assumed that in the implementation the "self" string is used to refer to  the object.
        >>> a = [1,2,3,4,5,6,7,8,9,10]
        >>> print('Syntactic sugar decoration - hacky way')
        Syntactic sugar decoration - hacky way
        >>> print('function')
        function
        >>> @Dec(chunk)(3).divide_and_run_SS_hack
        ... def f(mylist): return math.prod(mylist)
        >>> f(a)
        [6, 120, 504, 10]
        >>> print('instance method')
        instance method
        >>> class Dummy:
        ...     @Dec(chunk)(3).divide_and_run_SS_hack
        ...     def multiply_all(self, mylist): return math.prod(mylist)
        >>> Dummy().multiply_all(a)
        [6, 120, 504, 10]
        >>> print('='*20)
        ====================
        >>> print('SS (Manual) decoration')
        SS (Manual) decoration
        >>> print('function')
        function
        >>> def f(mylist): return math.prod(mylist)
        >>> Dec(chunk)(3).divide_and_run_SS_hack(f)(a)
        [6, 120, 504, 10]
        >>> print('instance method')
        instance method
        >>> class Dummy:
        ...     def multiply_all(self, mylist): return math.prod(mylist)
        >>> Dec(chunk)(3).divide_and_run_SS_hack(Dummy().multiply_all)(a)
        [6, 120, 504, 10]
        """
        vs = f_target.__code__.co_varnames
        if len(vs) > 0:
            # is method
            if vs[0] in {'self'}: # <- DANGER hardcoded
                # for @-decoration, method --> function
                if isinstance(f_target, types.FunctionType):
                    def wrapper(target_self, entries, *args, **kwargs): # <- "extra" parameter
                        output = []
                        for c in self.func(entries, *self.func_args, **self.func_kwargs):
                            output.append(f_target(target_self, c, *args, **kwargs)) # <- "extra" parameter
                        return output
                # for DIY-decoration
                elif isinstance(f_target, types.MethodType):
                    def wrapper(entries, *args, **kwargs):
                        output = []
                        for c in self.func(entries, *self.func_args, **self.func_kwargs):
                            output.append(f_target(c, *args, **kwargs))
                        return output
                return wrapper
        # is function
        def wrapper(entries, *args, **kwargs):
            output = []
            for c in self.func(entries, *self.func_args, **self.func_kwargs):
                output.append(f_target(c, *args, **kwargs))
            return output
        return wrapper
Related