Search code examples
pythongenericsclosuresdecoratorpython-decorators

How to apply decorators to an arbitrary function and call it inside the closure?


I was learning about how to use decorators in Python and as a practice, I made a decorator function that can print the execution time of a function it decorates. To do this, I had to call the decorated function in the closure function, which is where I had a problem. My code looks like this:

import time
import typing

T = typing.TypeVar("T")

def print_elapsed_time(func: typing.Callable[..., T]) -> typing.Callable[..., typing.Any]:
    def closure() -> T:
        start: float = time.time()
        returnValue: T = func()
        end: float = time.time()
        print(end - start)
        return returnValue
    return closure

@print_elapsed_time
def foo() -> int:
    print("foo")
    return 1

i: int = foo()
print(i)

This outputs:

foo
0.0
1

As you can see, I can only call func inside closure only if foo has 0 arguments. So my question is: is it possible to modify print_elapsed_time so that it will work correctly no matter what function it decorates and call the decorated function with the originally supplied arguments? What I mean is that if foo accepts 1 int argument, how do I modify print_elapsed_time so that when I write for example foo(10), the call to func inside closure can "adapt" to that and call func(10)?

Another strange problem is that at line 6, Pylance gives me this error:

TypeVar "T" appears only once in generic function signature

But I need to use T inside print_elapsed_time to know what to return. Is this an error caused by Plyance because it doesn't look into function bodies or is there some way to fix this error?


Solution

  • You can make the closure function take variable arguments and variable keyword arguments, and pass them to the call to func:

    def closure(*args, **kwargs) -> T:
        start: float = time.time()
        returnValue: T = func(*args, **kwargs)
        end: float = time.time()
        print(end - start)
        return returnValue
    

    To answer the second question of yours, since the decorator should return a function that returns the same type as that of the returning value of the function that's passed to the decorator, you should specify the returning type of print_elapsed_time to be a callable that returns the type T instead of Any:

    def print_elapsed_time(func: typing.Callable[..., T]) -> typing.Callable[..., T]: