Add cache transform

Author: mikeshardmind

src/async_utils/__init__.py

@@ -9,7 +9,7 @@
 __author__ = "Michael Hall"
 __license__ = "Apache-2.0"
 __copyright__ = "Copyright 2020-Present Michael Hall"
-__version__ = "2025.01.06"
+__version__ = "2025.01.12"
 
 import os
 import sys

src/async_utils/corofunc_cache.py

@@ -28,32 +28,48 @@
 type CoroFunc[**P, R] = Callable[P, Coroutine[t.Any, t.Any, R]]
 type CoroLike[**P, R] = Callable[P, Awaitable[R]]
 
+type _CT_RET = tuple[tuple[t.Any, ...], dict[t.str, t.Any]]
+type CacheTransformer = Callable[[tuple[t.Any, ...], dict[t.str, t.Any]], _CT_RET]
+
 type Deco[**P, R] = Callable[[CoroLike[P, R]], CoroFunc[P, R]]
 
 
-def corocache[**P, R](ttl: float | None = None) -> Deco[P, R]:
+def corocache[**P, R](
+    ttl: float | None = None,
+    *,
+    cache_transform: CacheTransformer | None = None,
+) -> Deco[P, R]:
     """Cache the results of the decorated coroutine.
 
     This is less powerful than the version in task_cache.py but may work better
     for some cases where typing of libraries this interacts with is too
     restrictive.
 
-    Note: This uses the args and kwargs of the original coroutine function as a
-    cache key. This includes instances (self) when wrapping methods.
+    Note: This by default uses the args and kwargs of the original coroutine
+    function as a cache key. This includes instances (self) when wrapping methods.
     Consider not wrapping instance methods, but what those methods call when
-    feasible in cases where this may matter.
+    feasible in cases where this may matter, or using a cache transform.
 
     The ordering of args and kwargs matters.
 
     Parameters
     ----------
     ttl: float | None
         The time to live in seconds for cached results. Defaults to None (forever)
+    cache_transform: CacheTransformer | None
+        An optional callable that transforms args and kwargs used
+        as a cache key.
 
     Returns
     -------
     A decorator which wraps coroutine-like functions with preemptive caching.
     """
+    if cache_transform is None:
+        key_func = make_key
+    else:
+
+        def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any]) -> Hashable:
+            return make_key(*cache_transform(args, kwds))
 
     def wrapper(coro: CoroLike[P, R]) -> CoroFunc[P, R]:
         internal_cache: dict[Hashable, asyncio.Future[R]] = {}
@@ -64,7 +80,7 @@ def _internal_cache_evict(key: Hashable, _ignored_task: object) -> None:
                 loop.call_later(ttl, internal_cache.pop, key)
 
         async def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
-            key = make_key(args, kwargs)
+            key = key_func(args, kwargs)
             if (cached := internal_cache.get(key)) is not None:
                 return await cached
 
@@ -79,17 +95,22 @@ async def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
     return wrapper
 
 
-def lrucorocache[**P, R](ttl: float | None = None, maxsize: int = 1024) -> Deco[P, R]:
+def lrucorocache[**P, R](
+    ttl: float | None = None,
+    maxsize: int = 1024,
+    *,
+    cache_transform: CacheTransformer | None = None,
+) -> Deco[P, R]:
     """Cache the results of the decorated coroutine.
 
     This is less powerful than the version in task_cache.py but may work better
     for some cases where typing of libraries this interacts with is too
     restrictive.
 
-    Note: This uses the args and kwargs of the original coroutine function as a
-    cache key. This includes instances (self) when wrapping methods.
+    Note: This by default uses the args and kwargs of the original coroutine
+    function as a cache key. This includes instances (self) when wrapping methods.
     Consider not wrapping instance methods, but what those methods call when
-    feasible in cases where this may matter.
+    feasible in cases where this may matter, or using a cache transform.
 
     The ordering of args and kwargs matters.
 
@@ -103,11 +124,20 @@ def lrucorocache[**P, R](ttl: float | None = None, maxsize: int = 1024) -> Deco[
         The maximum number of items to retain no matter if they have reached
         expiration by ttl or not.
         Items evicted by this policy are evicted by least recent use.
+    cache_transform: CacheTransformer | None
+        An optional callable that transforms args and kwargs used
+        as a cache key.
 
     Returns
     -------
     A decorator which wraps coroutine-like functions with preemptive caching.
     """
+    if cache_transform is None:
+        key_func = make_key
+    else:
+
+        def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any]) -> Hashable:
+            return make_key(*cache_transform(args, kwds))
 
     def wrapper(coro: CoroLike[P, R]) -> CoroFunc[P, R]:
         internal_cache: LRU[Hashable, asyncio.Future[R]] = LRU(maxsize)
@@ -119,7 +149,7 @@ def _internal_cache_evict(key: Hashable, _ignored_task: object) -> None:
 
         @wraps(coro)
         async def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
-            key = make_key(args, kwargs)
+            key = key_func(args, kwargs)
             if (cached := internal_cache.get(key, None)) is not None:
                 return await cached
 

src/async_utils/task_cache.py

@@ -31,13 +31,20 @@
 type TaskFunc[**P, R] = CoroFunc[P, R] | Callable[P, asyncio.Task[R]]
 type TaskCoroFunc[**P, R] = CoroFunc[P, R] | TaskFunc[P, R]
 
+type _CT_RET = tuple[tuple[t.Any, ...], dict[t.str, t.Any]]
+type CacheTransformer = Callable[[tuple[t.Any, ...], dict[t.str, t.Any]], _CT_RET]
+
 type Deco[**P, R] = Callable[[TaskCoroFunc[P, R]], TaskFunc[P, R]]
 
 # Non-annotation assignments for transformed functions
 _WRAP_ASSIGN = ("__module__", "__name__", "__qualname__", "__doc__")
 
 
-def taskcache[**P, R](ttl: float | None = None) -> Deco[P, R]:
+def taskcache[**P, R](
+    ttl: float | None = None,
+    *,
+    cache_transform: CacheTransformer | None = None,
+) -> Deco[P, R]:
     """Cache the results of the decorated coroutine.
 
     Decorator to modify coroutine functions to instead act as functions
@@ -46,23 +53,32 @@ def taskcache[**P, R](ttl: float | None = None) -> Deco[P, R]:
     For general use, this leaves the end user API largely the same,
     while leveraging tasks to allow preemptive caching.
 
-    Note: This uses the args and kwargs of the original coroutine function as a
-    cache key. This includes instances (self) when wrapping methods.
+    Note: This by default uses the args and kwargs of the original coroutine
+    function as a cache key. This includes instances (self) when wrapping methods.
     Consider not wrapping instance methods, but what those methods call when
-    feasible in cases where this may matter.
+    feasible in cases where this may matter, or using a cache transform.
 
     The ordering of args and kwargs matters.
 
     Parameters
     ----------
     ttl: float | None
         The time to live in seconds for cached results. Defaults to None (forever)
+    cache_transform: CacheTransformer | None
+        An optional callable that transforms args and kwargs used
+        as a cache key.
 
     Returns
     -------
     A decorator which wraps coroutine-like objects in functions that return
     preemptively cached tasks.
     """
+    if cache_transform is None:
+        key_func = make_key
+    else:
+
+        def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any]) -> Hashable:
+            return make_key(*cache_transform(args, kwds))
 
     def wrapper(coro: TaskCoroFunc[P, R]) -> TaskFunc[P, R]:
         internal_cache: dict[Hashable, asyncio.Task[R]] = {}
@@ -74,7 +90,7 @@ def _internal_cache_evict(key: Hashable, _ignored_task: object) -> None:
 
         @wraps(coro, assigned=_WRAP_ASSIGN)
         def wrapped(*args: P.args, **kwargs: P.kwargs) -> asyncio.Task[R]:
-            key = make_key(args, kwargs)
+            key = key_func(args, kwargs)
             if (cached := internal_cache.get(key)) is not None:
                 return cached
 
@@ -100,7 +116,12 @@ def wrapped(*args: P.args, **kwargs: P.kwargs) -> asyncio.Task[R]:
     return wrapper
 
 
-def lrutaskcache[**P, R](ttl: float | None = None, maxsize: int = 1024) -> Deco[P, R]:
+def lrutaskcache[**P, R](
+    ttl: float | None = None,
+    maxsize: int = 1024,
+    *,
+    cache_transform: CacheTransformer | None = None,
+) -> Deco[P, R]:
     """Cache the results of the decorated coroutine.
 
     Decorator to modify coroutine functions to instead act as functions
@@ -109,10 +130,10 @@ def lrutaskcache[**P, R](ttl: float | None = None, maxsize: int = 1024) -> Deco[
     For general use, this leaves the end user API largely the same,
     while leveraging tasks to allow preemptive caching.
 
-    Note: This uses the args and kwargs of the original coroutine function as a
-    cache key. This includes instances (self) when wrapping methods.
+    Note: This by default uses the args and kwargs of the original coroutine
+    function as a cache key. This includes instances (self) when wrapping methods.
     Consider not wrapping instance methods, but what those methods call when
-    feasible in cases where this may matter.
+    feasible in cases where this may matter, or using a cache transform.
 
     The ordering of args and kwargs matters.
 
@@ -127,12 +148,21 @@ def lrutaskcache[**P, R](ttl: float | None = None, maxsize: int = 1024) -> Deco[
         The maximum number of items to retain no matter if they have reached
         expiration by ttl or not.
         Items evicted by this policy are evicted by least recent use.
+    cache_transform: CacheTransformer | None
+        An optional callable that transforms args and kwargs used
+        as a cache key.
 
     Returns
     -------
     A decorator which wraps coroutine-like objects in functions that return
     preemptively cached tasks.
     """
+    if cache_transform is None:
+        key_func = make_key
+    else:
+
+        def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any]) -> Hashable:
+            return make_key(*cache_transform(args, kwds))
 
     def wrapper(coro: TaskCoroFunc[P, R]) -> TaskFunc[P, R]:
         internal_cache: LRU[Hashable, asyncio.Task[R]] = LRU(maxsize)
@@ -144,7 +174,7 @@ def _internal_cache_evict(key: Hashable, _ignored_task: object) -> None:
 
         @wraps(coro, assigned=_WRAP_ASSIGN)
         def wrapped(*args: P.args, **kwargs: P.kwargs) -> asyncio.Task[R]:
-            key = make_key(args, kwargs)
+            key = key_func(args, kwargs)
             if (cached := internal_cache.get(key, None)) is not None:
                 return cached