Search code examples
pythonpython-3.xfunctionscopedecorator

Create decorator with arguments that doesn't modify the function


Edit: The decorator should do something with the function at runtime, not modify it.

I want to be able to add a decorator to the top of functions to execute with arguments at the script start, and print the result (as an example).

Something similar to the code below, except the code below doesn't work because I don't know of a concise way to pass arguments to it. I don't want to modify the original function, just have it like a "test" section when the function is interpreted.

def test_decorator(func, *args, **kwargs):  # I'm aware this doesn't work. I'm modeling it after class methods for simplicity sake.
    print("Passed:", args, kwargs,
          "\nResult:", func(*args, **kwargs))

    return func


@test_decorator(70, 20, 10, "last_arg_example:")
def test_example(arg1, arg2, arg3=None, arg4="arg4_example:"):
    print(arg4, arg1 + arg2 + (arg3 if arg3 else 0))

I know that the first line of the decorator definition is messed up, but I'm also not clear on how to fix it. All of the things I could find use a "wrapper" of sorts, but I don't want to modify the original, like I said.

If the function requires no arguments, I should also be able to decorate it with just @test_decorator above it. Or even if it does require arguments, do nothing with the actual function and print something else entirely.

How would I do this?

Edit: To be more concise, I want the decorator to execute the function that follows it with the argument passed to the decorator and print the output.

>>> @decorator(4, 5)
>>> def test1(arg1, arg2):
...    return arg1 + arg2
...
9
>>>

Further functionality would be to handle no arguments by just passing no arguments to the function. Here are two cases:

>>> @decorator
>>> def test2(arg1=4, arg2=5):
...    return arg1 + arg2
...
9
>>>

Here is if the arguments have no defaults:

>>> @decorator
>>> def test1(arg1, arg2):
...    return arg1 + arg2
...
Traceback (most recent call last):
 File "<stdin>", line 1 in <module>
TypeError: test1() takes exactly 2 arguments (0 given)
>>>

Keep in mind that I'm aware the interactive interpreter doesn't work like this. I didn't actually try these examples; I typed then out. I wouldn't have been able to test the examples anyways because I don't have a working decorator for them.


Solution

  • @martineau's answer covers how to make a decorator factory that accepts arguments and returns a proper decorator. If you use this pattern, your actual decorator can return the original function unmodified:

    def test_decorator(*args, **kwargs):
        def decorator(func):
            return func
        return decorator
    

    The only difference from your original spec is that you will have to invoke the factory explicitly in the no-arg case: @test_decorator() instead of just @test_decorator.

    Now you can do whatever you want in the factory or the decorator, including registering the function call somewhere. For example:

    call_list = []
    
    def test_decorator(*args, **kwargs):
        def decorator(func):
            def make_call():
                return func(*args, **kwargs)
            call_list.append(make_call)
            # Decorated function is completely unchanged
            return func
        return decorator
    
    @test_decorator(70, 20, 10, "last_arg_example:")
    @test_decorator(60, 10, 20, "middle_arg_example:")
    @test_decorator(50, 0, 30, "first_arg_example:")
    def test_example(arg1, arg2, arg3=None, arg4="arg4_example:"):
        print(arg4, arg1 + arg2 + (arg3 if arg3 else 0))
    
    for item in call_list:
        item()
    

    This registers every decorated function call in call_list when the module is first imported, then runs the decorated calls one after the other. It does not modify the actual decorated function(s) in any way. Another nice side effect of that is that you can nest the factory an arbitrary number of times, as shown. The return value of each inner decoration is passed to the outer one, which means that the order in which the function calls are registered is "backwards". The result of the above is:

    first_arg_example: 80
    middle_arg_example: 90
    last_arg_example: 100
    

    More Info on Class Decorators

    You can reformat the decorator factory into a class. A decorator (for your purpose) is any callable that accepts a function and returns a function. This means that an instance of a class with a __call__ method can be a decorator too. The example above can be rewritten as

    call_list = []
    
    class TestDecorator:
        def __init__(self, *args, **kwargs):
            self.args = args
            self.kwargs = kwargs
    
        def __call__(self, func):
            def make_call():
                return func(*self.args, **self.kwargs)
            call_list.append(make_call)
            # Decorated function is completely unchanged
            return func
    
    @TestDecorator(70, 20, 10, "last_arg_example:")
    @TestDecorator(60, 10, 20, "middle_arg_example:")
    @TestDecorator(50, 0, 30, "first_arg_example:")
    def test_example(arg1, arg2, arg3=None, arg4="arg4_example:"):
        print(arg4, arg1 + arg2 + (arg3 if arg3 else 0))
    
    for item in call_list:
        item()