CODEWITHSUNDEEP
pythondecoratormypy
sixtyfootersdude
sixtyfootersdude

Reputation: 27221

Mypy type annotations for a decorator

I have a decorator class but am having trouble adding type annotations to it.

import functools

class LogInfo:

    def __init__(self, environment: str):
        self.environment = environment

    def __call__(self, func):
        @functools.wraps(func)
        def decorated(*args, **kwargs):
            # My Stuff goes here...
            return func(*args, **kwargs)
        return decorated

Closest I can get is this:

import functools
from collections import Callable
from typing import TypeVar, Any

GenericReturn = TypeVar("GenericReturn")
GenericCallable = TypeVar("GenericCallable", bound=Callable[..., GenericReturn])


class LogInfo:
    def __init__(self, environment: str) -> None:
        self.environment = environment

    def __call__(self, func: GenericCallable) -> GenericCallable:
        @functools.wraps(func)
        def decorated(*args: Any, **kwargs: Any) -> GenericReturn:
            # My Stuff goes here...
            return func(*args, **kwargs)
        return decorated  # LINE 29

But I still get this error:

29: error: Incompatible return value type (got "Callable[..., Any]", expected "GenericCallable")

Removing the @functools.wraps(func) changes the error to:

29: error: Incompatible return value type (got "Callable[[VarArg(Any), KwArg(Any)], GenericReturn]", expected "GenericCallable")

Upvotes: 4

Views: 3264

Answers (1)

sixtyfootersdude
sixtyfootersdude

Reputation: 27221

This is a decent solution:

import functools
from collections import Callable
from typing import TypeVar, cast, Any

T = TypeVar("T", bound=Callable[..., Any])


class LogInfo:
    def __init__(self, environment: str):
        self.environment = environment

    def __call__(self, func: T) -> T:
        @functools.wraps(func)
        def decorated(*args, **kwargs):  # type: ignore
            # My Stuff goes here...
            return func(*args, **kwargs)

        return cast(T, decorated)

We can test this with the follow code:

@LogInfo(environment="HI")
def foo(input: str) -> str:
    return f"{input}{input}"

# NOTE: Intentional error here to trigger mypy
def bar() -> int:
    return foo("jake")

As expected, we get this mypy error:

error: Incompatible return value type (got "str", expected "int")

Remaining things that could be improved:

  • The return cast(T, decorated) is ugly.
  • Typing args & kwargs would be good.

Upvotes: 3

Related Questions