pythonpython-decoratorstype-hintingmypypython-lru-cache

TypeVar inference broken by lru_cache decorator


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)

Solution

  • 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.