1
1
import asyncio
2
+ import inspect
2
3
from collections import OrderedDict
3
4
import functools
4
5
import logging
5
6
import os
6
7
import pickle
7
8
import sqlite3
8
9
from pathlib import Path
9
- from typing import Callable , Any , Awaitable , Hashable
10
-
10
+ from typing import Callable , Any , Awaitable , Hashable , Optional
11
11
12
12
USE_CACHE = True if os .getenv ("NO_CACHE" ) != "1" else False
13
13
CACHE_LOCATION = (
@@ -208,22 +208,33 @@ def __init__(
208
208
self ,
209
209
max_size : int ,
210
210
method : Callable [..., Awaitable [Any ]],
211
- cache_key_index : int = 0 ,
211
+ cache_key_index : Optional [ int ] = 0 ,
212
212
):
213
213
"""
214
214
Args:
215
215
max_size: max size of the cache (in items)
216
216
method: the function to cache
217
- cache_key_index: if the method takes multiple args, only one will be used as the cache key. This is the
218
- index of that cache key in the args list (default is the first arg)
217
+ cache_key_index: if the method takes multiple args, this is the index of that cache key in the args list
218
+ (default is the first arg). By setting this to `None`, it will use all args as the cache key.
219
219
"""
220
220
self ._inflight : dict [Hashable , asyncio .Future ] = {}
221
221
self ._method = method
222
222
self ._cache = LRUCache (max_size = max_size )
223
223
self ._cache_key_index = cache_key_index
224
224
225
- async def __call__ (self , * args : Any ) -> Any :
226
- key = args [self ._cache_key_index ]
225
+ def make_cache_key (self , args : tuple , kwargs : dict ) -> Hashable :
226
+ bound = inspect .signature (self ._method ).bind (* args , ** kwargs )
227
+ bound .apply_defaults ()
228
+
229
+ if self ._cache_key_index is not None :
230
+ key_name = list (bound .arguments )[self ._cache_key_index ]
231
+ return bound .arguments [key_name ]
232
+
233
+ return (tuple (bound .arguments .items ()),)
234
+
235
+ async def __call__ (self , * args : Any , ** kwargs : Any ) -> Any :
236
+ key = self .make_cache_key (args , kwargs )
237
+
227
238
if item := self ._cache .get (key ):
228
239
return item
229
240
@@ -235,7 +246,7 @@ async def __call__(self, *args: Any) -> Any:
235
246
self ._inflight [key ] = future
236
247
237
248
try :
238
- result = await self ._method (* args )
249
+ result = await self ._method (* args , ** kwargs )
239
250
self ._cache .set (key , result )
240
251
future .set_result (result )
241
252
return result
0 commit comments