Skip to content

Commit 901cf24

Browse files
committed
fix(typing): preserve decorated function type signatures with ParamSpec
1 parent 4602401 commit 901cf24

4 files changed

Lines changed: 590 additions & 26 deletions

File tree

src/cachier/core.py

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from concurrent.futures import ThreadPoolExecutor
1717
from datetime import datetime, timedelta
1818
from functools import wraps
19-
from typing import Any, Callable, Optional, Union
19+
from typing import Any, Callable, Optional, ParamSpec, TypeVar, Union
2020
from warnings import warn
2121

2222
from ._types import RedisClient, S3Client
@@ -31,6 +31,9 @@
3131
from .metrics import CacheMetrics, MetricsContext
3232
from .util import parse_bytes
3333

34+
_P = ParamSpec("_P")
35+
_R = TypeVar("_R")
36+
3437
MAX_WORKERS_ENVAR_NAME = "CACHIER_MAX_WORKERS"
3538
DEFAULT_MAX_WORKERS = 8
3639
ZERO_TIMEDELTA = timedelta(seconds=0)
@@ -221,7 +224,7 @@ def cachier(
221224
allow_non_static_methods: Optional[bool] = None,
222225
enable_metrics: bool = False,
223226
metrics_sampling_rate: float = 1.0,
224-
):
227+
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
225228
"""Wrap as a persistent, stale-free memoization decorator.
226229
227230
The positional and keyword arguments to the wrapped function must be
@@ -400,7 +403,7 @@ def cachier(
400403
else:
401404
raise ValueError("specified an invalid core: %s" % backend)
402405

403-
def _cachier_decorator(func):
406+
def _cachier_decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
404407
core.set_func(func)
405408

406409
# Guard: raise TypeError when decorating an instance method unless
@@ -513,7 +516,7 @@ def _call(*args, max_age: Optional[timedelta] = None, **kwds):
513516
from .config import _global_params
514517

515518
if ignore_cache or not _global_params.caching_enabled:
516-
return func(args[0], **kwargs) if core.func_is_method else func(**kwargs)
519+
return func(args[0], **kwargs) if core.func_is_method else func(**kwargs) # type: ignore[call-arg]
517520

518521
with MetricsContext(cache_metrics) as _mctx:
519522
key, entry = core.get_entry((), kwargs)
@@ -629,7 +632,7 @@ async def _call_async(*args, max_age: Optional[timedelta] = None, **kwds):
629632
from .config import _global_params
630633

631634
if ignore_cache or not _global_params.caching_enabled:
632-
return await func(args[0], **kwargs) if core.func_is_method else await func(**kwargs)
635+
return await func(args[0], **kwargs) if core.func_is_method else await func(**kwargs) # type: ignore[call-arg,misc]
633636

634637
with MetricsContext(cache_metrics) as _mctx:
635638
key, entry = await core.aget_entry((), kwargs)
@@ -699,14 +702,14 @@ async def _call_async(*args, max_age: Optional[timedelta] = None, **kwds):
699702
if is_coroutine:
700703

701704
@wraps(func)
702-
async def func_wrapper(*args, **kwargs):
703-
return await _call_async(*args, **kwargs)
705+
async def func_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
706+
return await _call_async(*args, **kwargs) # type: ignore[arg-type]
704707

705708
else:
706709

707710
@wraps(func)
708-
def func_wrapper(*args, **kwargs):
709-
return _call(*args, **kwargs)
711+
def func_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
712+
return _call(*args, **kwargs) # type: ignore[arg-type]
710713

711714
def _clear_cache():
712715
"""Clear the cache."""
@@ -751,13 +754,13 @@ def _precache_value(*args, value_to_cache, **kwds):
751754
kwargs = _convert_args_kwargs(func, _is_method=core.func_is_method, args=args, kwds=kwds)
752755
return core.precache_value((), kwargs, value_to_cache)
753756

754-
func_wrapper.clear_cache = _clear_cache
755-
func_wrapper.clear_being_calculated = _clear_being_calculated
756-
func_wrapper.aclear_cache = _aclear_cache
757-
func_wrapper.aclear_being_calculated = _aclear_being_calculated
758-
func_wrapper.cache_dpath = _cache_dpath
759-
func_wrapper.precache_value = _precache_value
760-
func_wrapper.metrics = cache_metrics # Expose metrics object
761-
return func_wrapper
757+
func_wrapper.clear_cache = _clear_cache # type: ignore[attr-defined]
758+
func_wrapper.clear_being_calculated = _clear_being_calculated # type: ignore[attr-defined]
759+
func_wrapper.aclear_cache = _aclear_cache # type: ignore[attr-defined]
760+
func_wrapper.aclear_being_calculated = _aclear_being_calculated # type: ignore[attr-defined]
761+
func_wrapper.cache_dpath = _cache_dpath # type: ignore[attr-defined]
762+
func_wrapper.precache_value = _precache_value # type: ignore[attr-defined]
763+
func_wrapper.metrics = cache_metrics # type: ignore[attr-defined]
764+
return func_wrapper # type: ignore[return-value]
762765

763766
return _cachier_decorator

tests/requirements.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ pytest-rerunfailures # for retrying flaky tests
77
coverage
88
pytest-cov
99
birch
10+
# type checking
11+
mypy
1012
# to be able to run `python setup.py checkdocs`
1113
collective.checkdocs
1214
pygments

tests/test_typing.py

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
"""Tests that the @cachier decorator preserves function type signatures.
2+
3+
These tests invoke mypy programmatically and assert that decorated functions retain their parameter types and return
4+
types as seen by static analysis.
5+
6+
"""
7+
8+
import textwrap
9+
10+
import pytest
11+
12+
mypy_api = pytest.importorskip("mypy.api", reason="mypy is required for typing tests")
13+
14+
15+
def _run_mypy(code: str) -> tuple[list[str], list[str]]:
16+
"""Run mypy on a code snippet and return (notes, errors).
17+
18+
Parameters
19+
----------
20+
code : str
21+
Python source code to type-check.
22+
23+
Returns
24+
-------
25+
tuple[list[str], list[str]]
26+
A tuple of (note lines, error lines) from mypy output.
27+
28+
"""
29+
result = mypy_api.run(
30+
[
31+
"-c",
32+
textwrap.dedent(code),
33+
"--no-error-summary",
34+
"--hide-error-context",
35+
]
36+
)
37+
stdout = result[0]
38+
notes = []
39+
errors = []
40+
for line in stdout.splitlines():
41+
if ": note:" in line:
42+
notes.append(line)
43+
elif ": error:" in line:
44+
errors.append(line)
45+
return notes, errors
46+
47+
48+
class TestSyncTyping:
49+
"""Verify that synchronous decorated functions preserve types."""
50+
51+
def test_return_type_preserved(self) -> None:
52+
"""Mypy should infer the original return type through @cachier."""
53+
notes, errors = _run_mypy("""
54+
from cachier import cachier
55+
56+
@cachier()
57+
def my_func(x: int) -> str:
58+
return str(x)
59+
60+
reveal_type(my_func(5))
61+
""")
62+
assert not errors
63+
assert any('"str"' in n for n in notes)
64+
65+
def test_param_types_preserved(self) -> None:
66+
"""Mypy should see the original parameter types through @cachier."""
67+
notes, errors = _run_mypy("""
68+
from cachier import cachier
69+
70+
@cachier()
71+
def my_func(x: int, y: str) -> list[str]:
72+
return [y] * x
73+
74+
reveal_type(my_func)
75+
""")
76+
assert not errors
77+
assert any("int" in n and "str" in n for n in notes)
78+
79+
def test_wrong_arg_type_is_error(self) -> None:
80+
"""Mypy should reject calls with wrong argument types."""
81+
_notes, errors = _run_mypy("""
82+
from cachier import cachier
83+
84+
@cachier()
85+
def add(a: int, b: int) -> int:
86+
return a + b
87+
88+
add("not", "ints")
89+
""")
90+
assert errors
91+
92+
def test_return_type_mismatch_is_error(self) -> None:
93+
"""Mypy should catch assigning the result to an incompatible type."""
94+
_notes, errors = _run_mypy("""
95+
from cachier import cachier
96+
97+
@cachier()
98+
def get_name() -> str:
99+
return "hello"
100+
101+
x: int = get_name()
102+
""")
103+
assert errors
104+
105+
106+
class TestAsyncTyping:
107+
"""Verify that async decorated functions preserve types."""
108+
109+
def test_async_return_type_preserved(self) -> None:
110+
"""Mypy should infer the awaited return type for async functions."""
111+
notes, errors = _run_mypy("""
112+
import asyncio
113+
from cachier import cachier
114+
115+
@cachier()
116+
async def fetch(url: str) -> bytes:
117+
return b"data"
118+
119+
async def main() -> None:
120+
result = await fetch("http://example.com")
121+
reveal_type(result)
122+
123+
asyncio.run(main())
124+
""")
125+
assert not errors
126+
assert any('"bytes"' in n for n in notes)
127+
128+
def test_async_signature_preserved(self) -> None:
129+
"""Mypy should see the async function as a coroutine."""
130+
notes, errors = _run_mypy("""
131+
from cachier import cachier
132+
133+
@cachier()
134+
async def fetch(url: str) -> bytes:
135+
return b"data"
136+
137+
reveal_type(fetch)
138+
""")
139+
assert not errors
140+
assert any("Coroutine" in n for n in notes)
141+
142+
def test_async_wrong_arg_type_is_error(self) -> None:
143+
"""Mypy should reject calls with wrong argument types for async."""
144+
_notes, errors = _run_mypy("""
145+
from cachier import cachier
146+
147+
@cachier()
148+
async def fetch(url: str) -> bytes:
149+
return b"data"
150+
151+
async def main() -> None:
152+
await fetch(123)
153+
""")
154+
assert errors
155+
156+
157+
class TestComplexSignatures:
158+
"""Verify preservation of more complex type signatures."""
159+
160+
def test_optional_params(self) -> None:
161+
"""Mypy should preserve Optional parameter types."""
162+
notes, errors = _run_mypy("""
163+
from typing import Optional
164+
from cachier import cachier
165+
166+
@cachier()
167+
def greet(name: str, greeting: Optional[str] = None) -> str:
168+
return f"{greeting or 'Hello'}, {name}"
169+
170+
reveal_type(greet)
171+
""")
172+
assert not errors
173+
assert any("str" in n for n in notes)
174+
175+
def test_generic_return_type(self) -> None:
176+
"""Mypy should preserve generic return types like dict."""
177+
notes, errors = _run_mypy("""
178+
from cachier import cachier
179+
180+
@cachier()
181+
def make_mapping(keys: list[str], value: int) -> dict[str, int]:
182+
return {k: value for k in keys}
183+
184+
reveal_type(make_mapping(["a"], 1))
185+
""")
186+
assert not errors
187+
assert any("dict[str, int]" in n for n in notes)
188+
189+
def test_none_return_type(self) -> None:
190+
"""Mypy should preserve None return type."""
191+
notes, errors = _run_mypy("""
192+
from cachier import cachier
193+
194+
@cachier()
195+
def side_effect(x: int) -> None:
196+
pass
197+
198+
reveal_type(side_effect(1))
199+
""")
200+
assert not errors
201+
assert any('"None"' in n for n in notes)
202+
203+
204+
class TestDecoratorWithArgs:
205+
"""Verify typing works with various decorator arguments."""
206+
207+
def test_with_backend_arg(self) -> None:
208+
"""Type preservation should work with explicit backend selection."""
209+
notes, errors = _run_mypy("""
210+
from cachier import cachier
211+
212+
@cachier(backend="memory")
213+
def compute(x: float) -> float:
214+
return x * 2.0
215+
216+
reveal_type(compute(1.0))
217+
""")
218+
assert not errors
219+
assert any('"float"' in n for n in notes)
220+
221+
def test_with_stale_after_arg(self) -> None:
222+
"""Type preservation should work with stale_after parameter."""
223+
notes, errors = _run_mypy("""
224+
from datetime import timedelta
225+
from cachier import cachier
226+
227+
@cachier(stale_after=timedelta(hours=1))
228+
def lookup(key: str) -> list[int]:
229+
return [1, 2, 3]
230+
231+
reveal_type(lookup("x"))
232+
""")
233+
assert not errors
234+
assert any("list[int]" in n for n in notes)

0 commit comments

Comments
 (0)