python - How to "teach" typecheckers that the decorator extends function's argument types? - Stack Ove

admin管理员组

文章数量:1122832

from functools import wraps
from typing import Any, Callable, Generic, ParamSpec, TypeVar

T = TypeVar("T")
P = ParamSpec("P")


class LazyOp(Generic[T]):
    def __init__(self, func: Callable[..., T], args: tuple[Any, ...], kwargs: dict[str, Any]):
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __call__(self) -> T:
        args = [a() if isinstance(a, LazyOp) else a for a in self.args]
        kwargs = {k: v() if isinstance(v, LazyOp) else v for k, v in self.kwargs.items()}
        return self.func(*args, **kwargs)


def lazy_decorator(func: Callable[P, T]) -> Callable[P, LazyOp[T]]:
    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> LazyOp[T]:
        return LazyOp(func, args, kwargs)

    return wrapper


@lazy_decorator
def lazy_sum(a: int, b: int) -> int:
    return a + b


lazy_result = lazy_sum(1, 2)

lazy_result = lazy_sum(
    lazy_result, 4
)  # Argument of type "LazyOp[int]" cannot be assigned to parameter "a" of type "int". "LazyOp[int]" is not assignable to "int"
result = lazy_result()
print(result)

How to enable automatic support for lazy arguments in @lazy_decorator-decorated functions with mypy?

I need a way to enable @lazy_decorator functions to:

1. Preserve original type annotations: Functions should still appear to mypy as if they accept and return their original types (e.g., int).
2. Allow LazyOp arguments transparently: Decorated functions should accept both the original type (e.g., int) and its lazy variant (LazyOp[int]) without triggering mypy errors.
3. Avoid manual adjustments: There should be no need to rewrite type annotations for each decorated function or introduce additional wrapping types in the function signature.
4. Integrate seamlessly: The solution should not require external plugins, major restructuring, or sacrificing runtime efficiency.

Essentially, I want a way to "teach" mypy that the decorator extends argument types to include their lazy equivalents while maintaining the function's original type annotation for clarity and developer experience.

本文标签：