Allow TypedDict for cache implementation (#105)

Author: Dreamsorcerer

.coveragerc

@@ -13,7 +13,10 @@ _site-packages-to-src-mapping =
 
 [report]
 exclude_also =
-  ^\s*@pytest\.mark\.xfail
+    if TYPE_CHECKING
+    assert False
+    : \.\.\.(\s*#.*)?$
+    ^ +\.\.\.$
 # small library, don't fail when running without C-extension
 fail_under = 50.00
 skip_covered = true

src/propcache/_helpers_py.py

@@ -1,6 +1,7 @@
 """Various helper functions."""
 
 import sys
+from collections.abc import Mapping
 from functools import cached_property
 from typing import Any, Callable, Generic, Optional, Protocol, TypeVar, Union, overload
 
@@ -13,10 +14,13 @@
     Self = Any
 
 _T = TypeVar("_T")
+# We use Mapping to make it possible to use TypedDict, but this isn't
+# technically type safe as we need to assign into the dict.
+_Cache = TypeVar("_Cache", bound=Mapping[str, Any])
 
 
-class _CacheImpl(Protocol):
-    _cache: dict[str, Any]
+class _CacheImpl(Protocol[_Cache]):
+    _cache: _Cache
 
 
 class under_cached_property(Generic[_T]):
@@ -29,7 +33,7 @@ class under_cached_property(Generic[_T]):
     variable.  It is, in Python parlance, a data descriptor.
     """
 
-    def __init__(self, wrapped: Callable[..., _T]) -> None:
+    def __init__(self, wrapped: Callable[[Any], _T]) -> None:
         self.wrapped = wrapped
         self.__doc__ = wrapped.__doc__
         self.name = wrapped.__name__
@@ -38,10 +42,10 @@ def __init__(self, wrapped: Callable[..., _T]) -> None:
     def __get__(self, inst: None, owner: Optional[type[object]] = None) -> Self: ...
 
     @overload
-    def __get__(self, inst: _CacheImpl, owner: Optional[type[object]] = None) -> _T: ...
+    def __get__(self, inst: _CacheImpl[Any], owner: Optional[type[object]] = None) -> _T: ...  # type: ignore[misc]
 
     def __get__(
-        self, inst: Optional[_CacheImpl], owner: Optional[type[object]] = None
+        self, inst: Optional[_CacheImpl[Any]], owner: Optional[type[object]] = None
     ) -> Union[_T, Self]:
         if inst is None:
             return self
@@ -52,5 +56,5 @@ def __get__(
             inst._cache[self.name] = val
             return val
 
-    def __set__(self, inst: _CacheImpl, value: _T) -> None:
+    def __set__(self, inst: _CacheImpl[Any], value: _T) -> None:
         raise AttributeError("cached property is read-only")

tests/test_under_cached_property.py

@@ -1,6 +1,6 @@
 import sys
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, Protocol, TypeVar
+from typing import TYPE_CHECKING, Any, Protocol, TypedDict, TypeVar
 
 import pytest
 
@@ -40,6 +40,34 @@ def prop2(self) -> str:
     assert a.prop2 == "foo"
 
 
+def test_under_cached_property_typeddict(propcache_module: APIProtocol) -> None:
+    """Test static typing passes with TypedDict."""
+
+    class _Cache(TypedDict, total=False):
+        prop: int
+        prop2: str
+
+    class A:
+        def __init__(self) -> None:
+            self._cache: _Cache = {}
+
+        @propcache_module.under_cached_property
+        def prop(self) -> int:
+            return 1
+
+        @propcache_module.under_cached_property
+        def prop2(self) -> str:
+            return "foo"
+
+    a = A()
+    if sys.version_info >= (3, 11):
+        assert_type(a.prop, int)
+    assert a.prop == 1
+    if sys.version_info >= (3, 11):
+        assert_type(a.prop2, str)
+    assert a.prop2 == "foo"
+
+
 def test_under_cached_property_assignment(propcache_module: APIProtocol) -> None:
     class A:
         def __init__(self) -> None: