Coverage for C:\Python311\Lib\site-packages\persist_cache\persist_cache.py: 57%

94 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-05-06 20:49 +1000

1import inspect 

2import os 

3from datetime import timedelta 

4from functools import wraps 

5from typing import Any, Callable, Union 

6 

7from . import caching 

8from .caching import NOT_IN_CACHE 

9from .helpers import inflate_arguments, is_async, signaturize 

10 

11 

12def cache( 

13 name: Union[str, Callable, None] = None, 

14 dir: Union[str, None] = None, 

15 expiry: Union[int, float, timedelta, None] = None, 

16 ) -> Callable: 

17 """Persistently and locally cache the returns of a function. 

18  

19 The function to be cached must accept and return dillable objects only (with the exception of methods' `self` argument, which is always ignored). Additionally, for consistent caching across subsequent sessions, arguments and returns should also be hashable. 

20  

21 Arguments: 

22 name (`str | Callable`, optional): The name of the cache (or, if `cache()` is being called as an argument-less decorator (ie, as `@cache` instead of `@cache(...)`), the function to be cached). Defaults to the qualified name of the function. If `dir` is set, this argument will be ignored. 

23  

24 dir (`str`, optional): The directory in which the cache should be stored. Defaults to a subdirectory named after the hash of the cache's name in a parent folder named '.persist_cache' in the current working directory. 

25  

26 expiry (`int | float | timedelta`, optional): How long, in seconds or as a `timedelta`, function calls should persist in the cache. Defaults to `None`. 

27  

28 Returns: 

29 `Callable`: If `cache()` is called with arguments, a decorator that wraps the function to be cached, otherwise, the wrapped function itself. Once wrapped, the function will have the following methods attached to it: 

30 - `set_expiry(value: int | float | timedelta) -> None`: Set the expiry of the cache. 

31 - `flush_cache() -> None`: Flush out any expired cached returns. 

32 - `clear_cache() -> None`: Clear out all cached returns. 

33 - `delete_cache() -> None`: Delete the cache.""" 

34 

35 def decorator(func: Callable) -> Callable: 

36 nonlocal name, dir, expiry 

37 

38 # If the cache directory has not been set, and the name of the cache has, set it to a subdirectory by the name of the hash of that name in a directory named '.persist_cache' in the current working directory, or, if the name of the cache has not been set, set the name of that subdirectory to the hash of the qualified name of the function. 

39 if dir is None: 

40 name = name if name is not None else func.__qualname__ 

41 

42 dir = f'.persist_cache/{caching.shorthash(name)}' 

43 

44 # Create the cache directory and any other necessary directories if it does not exist. 

45 if not os.path.exists(dir): 

46 os.makedirs(dir, exist_ok=True) 

47 

48 # If an expiry has been set, flush out any expired cached returns. 

49 if expiry is not None: 

50 caching.flush(dir, expiry) 

51 

52 # Flag whether the function is a method to enable the exclusion of the first argument (which will be the instance of the function's class) from being hashed to produce the cache key. 

53 is_method = inspect.ismethod(func) 

54 

55 # Preserve a map of the function's arguments to their default values and the name and index of the args parameter if such a parameter exists to enable the remapping of positional arguments to their keywords, which thereby allows for the consistent caching of function calls where positional arguments are used on some occasions and keyword arguments are used on others. 

56 signature, args_parameter, args_i = signaturize(func) 

57 

58 # Initialise a wrapper for synchronous functions. 

59 def sync_wrapper(*args: tuple[Any], **kwargs: dict[str, Any]) -> Any: 

60 nonlocal dir, expiry, is_method 

61 

62 # Map arguments to their keywords or the keyword of the args parameter where necessary, filtering out the first argument if the function is a method, to enable the consistent caching of function calls where positional arguments are used on some occasions and keyword arguments are used on others. 

63 arguments = inflate_arguments(signature, args_parameter, args_i, args[is_method:], kwargs) 

64 

65 # Hash the arguments to produce the cache key. 

66 key = caching.hash(arguments) 

67 

68 # Get the value of the key from the cache if it is not expired, otherwise, call the function and set the value of the key in the cache to the result of that call. 

69 if (value := caching.get(key, dir, expiry)) is NOT_IN_CACHE: 

70 value = func(*args, **kwargs) 

71 caching.set(key, value, dir) 

72 

73 return value 

74 

75 # Initialise a wrapper for asynchronous functions. 

76 async def async_wrapper(*args: tuple[Any], **kwargs: dict[str, Any]) -> Any: 

77 nonlocal dir, expiry, is_method 

78 

79 # Map arguments to their keywords or the keyword of the args parameter where necessary, filtering out the first argument if the function is a method, to enable the consistent caching of function calls where positional arguments are used on some occasions and keyword arguments are used on others. 

80 arguments = inflate_arguments(signature, args_parameter, args_i, args[is_method:], kwargs) 

81 

82 # Hash the arguments to produce the cache key. 

83 key = caching.hash(arguments) 

84 

85 # Get the value of the key from the cache if it is not expired, otherwise, call the function and set the value of the key in the cache to the result of that call. 

86 if (value := caching.get(key, dir, expiry)) is NOT_IN_CACHE: 

87 value = await func(*args, **kwargs) 

88 caching.set(key, value, dir) 

89 

90 return value 

91 

92 # Initialise a wrapper for generator functions. 

93 def generator_wrapper(*args: tuple[Any], **kwargs: dict[str, Any]) -> Any: 

94 nonlocal dir, expiry, is_method 

95 

96 # Map arguments to their keywords or the keyword of the args parameter where necessary, filtering out the first argument if the function is a method, to enable the consistent caching of function calls where positional arguments are used on some occasions and keyword arguments are used on others. 

97 arguments = inflate_arguments(signature, args_parameter, args_i, args[is_method:], kwargs) 

98 

99 # Hash the arguments to produce the cache key. 

100 key = caching.hash(arguments) 

101 

102 # Get the value of the key from the cache if it is not expired, otherwise, call the function and set the value of the key in the cache to the result of that call. 

103 if (value := caching.get(key, dir, expiry)) is NOT_IN_CACHE: 

104 value = [] 

105 

106 for item in func(*args, **kwargs): 

107 value.append(item) 

108 

109 yield item 

110 

111 caching.set(key, value, dir) 

112 

113 return 

114 

115 for item in value: 

116 yield item 

117 

118 # Initialise a wrapper for asynchronous generator functions. 

119 async def async_generator_wrapper(*args: tuple[Any], **kwargs: dict[str, Any]) -> Any: 

120 nonlocal dir, expiry, is_method 

121 

122 # Map arguments to their keywords or the keyword of the args parameter where necessary, filtering out the first argument if the function is a method, to enable the consistent caching of function calls where positional arguments are used on some occasions and keyword arguments are used on others. 

123 arguments = inflate_arguments(signature, args_parameter, args_i, args[is_method:], kwargs) 

124 

125 # Hash the arguments to produce the cache key. 

126 key = caching.hash(arguments) 

127 

128 # Get the value of the key from the cache if it is not expired, otherwise, call the function and set the value of the key in the cache to the result of that call. 

129 if (value := caching.get(key, dir, expiry)) is NOT_IN_CACHE: 

130 value = [] 

131 

132 async for item in func(*args, **kwargs): 

133 value.append(item) 

134 

135 yield item 

136 

137 caching.set(key, value, dir) 

138 

139 return 

140 

141 for item in value: 

142 yield item 

143 

144 # Identify the appropriate wrapper for the function. 

145 if is_async(func): 

146 if inspect.isasyncgenfunction(func): 

147 wrapper = async_generator_wrapper 

148 

149 else: 

150 wrapper = async_wrapper 

151 

152 elif inspect.isgeneratorfunction(func): 

153 wrapper = generator_wrapper 

154 

155 else: 

156 wrapper = sync_wrapper 

157 

158 # Attach convenience functions to the wrapper for modifying the cache. 

159 def delete_cache() -> None: 

160 """Delete the cache.""" 

161 nonlocal dir 

162 

163 caching.delete(dir) 

164 

165 def clear_cache() -> None: 

166 """Clear the cache.""" 

167 nonlocal dir 

168 

169 caching.clear(dir) 

170 

171 def flush_cache() -> None: 

172 """Flush expired keys from the cache.""" 

173 nonlocal dir, expiry, is_method 

174 

175 caching.flush(dir, expiry) 

176 

177 def set_expiry(value: Union[int, float, timedelta, None]) -> None: 

178 """Set the expiry of the cache. 

179  

180 Arguments: 

181 expiry (`int | float | timedelta`): How long, in seconds or as a `timedelta`, function calls should persist in the cache.""" 

182 

183 nonlocal expiry 

184 

185 expiry = value 

186 

187 wrapper.delete_cache = delete_cache 

188 wrapper.clear_cache = clear_cache 

189 wrapper.cache_clear = wrapper.clear_cache # Add an alias for cache_clear which is used by lru_cache. 

190 wrapper.flush_cache = flush_cache 

191 wrapper.set_expiry = set_expiry 

192 

193 # Preserve the original function. 

194 wrapper.__wrapped__ = func 

195 

196 # Preserve the function's original signature. 

197 wrapper = wraps(func)(wrapper) 

198 

199 return wrapper 

200 

201 # If the first argument is a function and all of the other arguments are `None`, indicating that this decorator factory was invoked without passing any arguments, return the result of passing that function to the decorator while also emptying the first argument to avoid it being used by the decorator. 

202 if callable(name) and dir is expiry is None: 

203 func = name 

204 name = None 

205 

206 return decorator(func) 

207 

208 return decorator 

209 

210def delete(function_or_name: Union[str, Callable]) -> None: 

211 """Delete the cache of the given function or name. 

212  

213 Arguments: 

214 function_or_name (`str | Callable`): The function or name of the cache to be deleted.""" 

215 

216 name = function_or_name if isinstance(function_or_name, str) else function_or_name.__qualname__ 

217 caching.delete(f'.persist_cache/{caching.shorthash(name)}') 

218 

219def clear(function_or_name: Union[str, Callable]) -> None: 

220 """Clear the cache of the given function or name. 

221  

222 Arguments: 

223 function_or_name (`str | Callable`): The function or name of the cache to be cleared.""" 

224 

225 name = function_or_name if isinstance(function_or_name, str) else function_or_name.__qualname__ 

226 caching.clear(f'.persist_cache/{caching.shorthash(name)}') 

227 

228def flush(function_or_name: Union[str, Callable], expiry: Union[int, float, timedelta]) -> None: 

229 """Flush expired keys from the cache of the given function or name. 

230  

231 Arguments: 

232 function_or_name (`str | Callable`): The function or name of the cache to be flushed. 

233 expiry (`int | float | timedelta`): How long, in seconds or as a `timedelta`, function calls should persist in the cache.""" 

234 

235 name = function_or_name if isinstance(function_or_name, str) else function_or_name.__qualname__ 

236 caching.flush(f'.persist_cache/{caching.shorthash(name)}', expiry)