TypeVar inference broken by lru_cache decorator

Question:

python’s TypeVar inference broken when using lru_cache decorator. For example, after applying mypy the following example, only function with lru_cache causes error like:

main.py:14: error: Incompatible types in assignment (expression has type "T", variable has type "int")
Found 1 error in 1 file (checked 1 source file)

and pyright’s editor support also warn the same thing. Is this lru_cache‘s own limitation or is there some good workaround?

from functools import lru_cache
from typing import TypeVar

T = TypeVar("T")

def working(foo: T) -> T:
    return foo

@lru_cache(maxsize=None)
def not_working(foo: T) -> T:
    return foo

a: int = working(1)
b: int = not_working(1)
Asked By: orematasaburo

||

Answers:

Here’s the relevant parts of the lru_cache type hints

_T = TypeVar("_T")

class _lru_cache_wrapper(Generic[_T]):
    __wrapped__: Callable[..., _T]
    def __call__(self, *args: Hashable, **kwargs: Hashable) -> _T: ...

def lru_cache(
    maxsize: int | None = ..., typed: bool = ...
) -> Callable[[Callable[..., _T]], _lru_cache_wrapper[_T]]: ...

so it appears that, in its attempts to allow for any argument set, it loses any connection between the input and output types and so is unable to refine T to int. You may have to wrap lru_cache locally to fix this. You may be able to use ParamSpec, but you might find difficulties with that, see the note below. If you only need it for a small set of function types (unary, binary, ternary), you could wrap it for those.

Apparently they did actually fix this with ParamSpec but from a cursory reading it appears to break other things so they reverted it. This issue also discusses it.

Answered By: joel