Reputation: 23
I'm trying to annotate a decorator implemented as a class, but mypy seems to either lose the annotation or lose the type and think it's an Any. What I am trying to annotate:
class my_decorator:
def __init__(self, func):
self.func = func
def __call__(self, *args, **kwargs):
return self.func(*args, **kwargs)
@my_decorator
def func():
return 2
How do I annotate this so func is detected as returning an int after decorating? I realize the above looks simple and I could convert my_decorator to a function, but in reality it is subclassed to have more specialized options.
Upvotes: 1
Views: 1133
Reputation: 64348
You will need to make my_decorator
a generic class and do something like the following:
from typing import Any, Callable, Generic, TypeVar
T = TypeVar('T')
class my_decorator(Generic[T]):
def __init__(self, func: Callable[..., T]) -> None:
self.func = func
def __call__(self, *args: Any, **kwargs: Any) -> T:
return self.func(*args, **kwargs)
@my_decorator
def func() -> int:
return 2
That is, capture the return type of your function using a TypeVar, which is scoped to your my_decorator
class. This ensures the value bound to TypeVar is "available" when we try making use of it when analyzing any invocations of __call__
.
Unfortunately, it is not possible to ensure that the parameters of __call__
will match the parameters of func()
. So, mypy won't report an error if you try doing something like func(1, 2, 3)
.
This may become possible once mypy adds support for PEP 612, which adds better support for typing decorator-related code in general.
Upvotes: 2