From 2699f883f92108f64516a55bc9b6e8de5621b74d Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Mon, 6 Apr 2026 18:57:37 -0700 Subject: [PATCH] Type lp.memoize_on_disk. --- loopy/tools.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/loopy/tools.py b/loopy/tools.py index d645abc98..6c6963042 100644 --- a/loopy/tools.py +++ b/loopy/tools.py @@ -30,7 +30,7 @@ import numpy as np from constantdict import constantdict -from typing_extensions import override +from typing_extensions import ParamSpec, override import islpy as isl from pytools import Hash, ProcessLogger, memoize_method @@ -936,7 +936,14 @@ def clear_in_mem_caches() -> None: # {{{ memoize_on_disk -def memoize_on_disk(func, key_builder_t=LoopyKeyBuilder): +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def memoize_on_disk( + func: abc.Callable[_P, _R], + key_builder_t: type[KeyBuilderBase] = LoopyKeyBuilder, +) -> abc.Callable[_P, _R]: from functools import wraps from pytools.persistent_dict import WriteOncePersistentDict @@ -956,7 +963,7 @@ def memoize_on_disk(func, key_builder_t=LoopyKeyBuilder): caches.append(transform_cache) @wraps(func) - def wrapper(*args, **kwargs): + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: from loopy import CACHING_ENABLED if (not CACHING_ENABLED