fix: cachetools.__wrapped__ attribute typing for cached functions

stubs/cachetools/cachetools/__init__.pyi

@@ -1,7 +1,7 @@
-from _typeshed import IdentityFunction, Unused
+from _typeshed import Unused
 from collections.abc import Callable, Iterator, MutableMapping, Sequence
 from contextlib import AbstractContextManager
-from typing import Any, TypeVar, overload
+from typing import Any, Protocol, TypeVar, overload
 from typing_extensions import deprecated
 
 __all__ = ("Cache", "FIFOCache", "LFUCache", "LRUCache", "MRUCache", "RRCache", "TLRUCache", "TTLCache", "cached", "cachedmethod")
@@ -10,6 +10,7 @@ __version__: str
 _KT = TypeVar("_KT")
 _VT = TypeVar("_VT")
 _T = TypeVar("_T")
+_R = TypeVar("_R")
 
 class Cache(MutableMapping[_KT, _VT]):
     @overload
@@ -99,14 +100,18 @@ class TLRUCache(_TimedCache[_KT, _VT]):
     def ttu(self) -> Callable[[_KT, _VT, float], float]: ...
     def expire(self, time: float | None = None) -> list[tuple[_KT, _VT]]: ...
 
+class _Cached(Protocol[_R]):
+    __wrapped__: Callable[..., _R]
+    def __call__(self, *args: Any, **kwargs: Any) -> _R: ...
+
 def cached(
     cache: MutableMapping[_KT, Any] | None,
     key: Callable[..., _KT] = ...,
     lock: AbstractContextManager[Any] | None = None,
     info: bool = False,
-) -> IdentityFunction: ...
+) -> Callable[[Callable[..., _R]], _Cached[_R]]: ...
 def cachedmethod(
     cache: Callable[[Any], MutableMapping[_KT, Any] | None],
     key: Callable[..., _KT] = ...,
     lock: Callable[[Any], AbstractContextManager[Any]] | None = None,
-) -> IdentityFunction: ...
+) -> Callable[[Callable[..., _R]], _Cached[_R]]: ...

stubs/cachetools/cachetools/func.pyi

@@ -1,19 +1,24 @@
-from _typeshed import IdentityFunction
 from collections.abc import Callable, Sequence
-from typing import TypeVar
+from typing import TypeVar, overload
 from typing_extensions import deprecated
 
+from . import _Cached
+
 __all__ = ("fifo_cache", "lfu_cache", "lru_cache", "mru_cache", "rr_cache", "ttl_cache")
 _T = TypeVar("_T")
+_S = TypeVar("_S")
 
-def fifo_cache(maxsize: float | None = 128, typed: bool = False) -> IdentityFunction: ...
-def lfu_cache(maxsize: float | None = 128, typed: bool = False) -> IdentityFunction: ...
-def lru_cache(maxsize: float | None = 128, typed: bool = False) -> IdentityFunction: ...
+def fifo_cache(maxsize: float | None = 128, typed: bool = False) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...
+def lfu_cache(maxsize: float | None = 128, typed: bool = False) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...
+def lru_cache(maxsize: float | None = 128, typed: bool = False) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...
 @deprecated("@mru_cache is deprecated")
-def mru_cache(maxsize: float | None = 128, typed: bool = False) -> IdentityFunction: ...
+def mru_cache(maxsize: float | None = 128, typed: bool = False) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...
+@overload
+def rr_cache(maxsize: float | None = 128, *, typed: bool = False) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...
+@overload
 def rr_cache(
-    maxsize: float | None = 128, choice: Callable[[Sequence[_T]], _T] | None = ..., typed: bool = False
-) -> IdentityFunction: ...
+    maxsize: float | None, choice: Callable[[Sequence[_S]], _S], typed: bool = False
+) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...
 def ttl_cache(
     maxsize: float | None = 128, ttl: float = 600, timer: Callable[[], float] = ..., typed: bool = False
-) -> IdentityFunction: ...
+) -> Callable[[Callable[..., _T]], _Cached[_T]]: ...

tests/cachetools_wrapped_test.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+from cachetools import LRUCache, cached
+
+cache: LRUCache[str, Any] = LRUCache(maxsize=128)
+
+
+@cached(cache)
+def example_function(x: int) -> str:
+    return str(x)
+
+
+original_function = example_function.__wrapped__
+result = original_function(42)