user2084911
user2084911

Reputation: 23

Python mypy Annotation Decorators __call__

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

Answers (1)

Michael0x2a
Michael0x2a

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

Related Questions