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?

CodePudding user response：

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.

CodePudding user response：

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