How can the return annotation of decorated method be reflected when decoration (which changes the return type) is done dynamically

Viewed 68

Using ParamSpec and generic types it is possible to create a decorator that preserves the annotations of the decorated method and optionally also change them. The works as expected when the method is directly decorated inline using the @decorator approach however, when the decorator is applied dynamically on a class method the change in return type is not reflected.

Using the following contrived example:

import functools
import inspect
from typing import TYPE_CHECKING, Callable, List, ParamSpec, TypeVar

R = TypeVar("R")
P = ParamSpec("P")


def listifyresponse(func: Callable[P, R]) -> Callable[P, List[R]]:
    @functools.wraps(func)
    def wrapped(*args: P.args, **kwargs: P.kwargs) -> List[R]:
        return [func(*args, **kwargs)]

    return wrapped


def listifyclassresponses(cls):
    for name, method in inspect.getmembers(cls):
        if inspect.isfunction(method):
            setattr(cls, name, listifyresponse(method))

    return cls


class Base:
    def check(self, v: str, w: str) -> bool:
        return v == w


class SubA(Base):
    pass


@listifyclassresponses
class SubB(Base):
    pass


class SubC(Base):
    @listifyresponse
    def check(self, v: str, w: str) -> bool:
        return super().check(v, w)


if TYPE_CHECKING:
    reveal_type(SubA().check("a", "b")) # False
    reveal_type(SubB().check("a", "b")) # [False]
    reveal_type(SubC().check("a", "b")) # [False]

print(SubA().check("a", "b")) # Revealed type is "builtins.bool"
print(SubB().check("a", "b")) # Revealed type is "builtins.bool"
print(SubC().check("a", "b")) # Revealed type is "builtins.list[builtins.bool*]"

I observed the same behavior when using a metaclass to update the methods instead of a class decorator.

0 Answers
Related