Coverage for src / taipanstack / utils / cache.py: 100%

1""" 

2Intelligent Cache decorator. 

3 

4Provides in-memory caching that respects the Result monad and TTL, 

5ignoring caching for Err() results. 

6""" 

7 

8import functools 

9import inspect 

10import time 

11from collections.abc import Callable, Coroutine 

12from typing import Any, ParamSpec, Protocol, TypeVar, cast, overload 

13 

14from taipanstack.core.result import Err, Ok, Result 

15 

16P = ParamSpec("P") 

17T = TypeVar("T") 

18E = TypeVar("E", bound=Exception) 

19 

20 

21class CacheDecorator(Protocol): 

22 """Protocol for the cache decorator.""" 

23 

24 @overload 

25 def __call__( 

26 self, func: Callable[P, Result[T, E]] 

27 ) -> Callable[P, Result[T, E]]: ... # pragma: no cover 

28 

29 @overload 

30 def __call__( 

31 self, func: Callable[P, Coroutine[Any, Any, Result[T, E]]] 

32 ) -> Callable[P, Coroutine[Any, Any, Result[T, E]]]: ... # pragma: no cover 

33 

34 

35def cached(ttl: float) -> CacheDecorator: 

36 """Cache the Ok() results of a function for a given TTL. 

37 

38 Err() results are not cached. Supports both async and sync functions. 

39 

40 Args: 

41 ttl: Time to live in seconds. 

42 

43 Returns: 

44 Decorator function. 

45 

46 """ 

47 _cache: dict[tuple[Any, ...], tuple[float, Any]] = {} 

48 

49 def get_cache_key( 

50 func_name: str, args: tuple[Any, ...], kwargs: dict[str, Any] 

51 ) -> tuple[Any, ...]: 

52 kwargs_tuple = tuple(sorted(kwargs.items())) 

53 return (func_name, args, kwargs_tuple) 

54 

55 def decorator( 

56 func: ( 

57 Callable[P, Result[T, E]] | Callable[P, Coroutine[Any, Any, Result[T, E]]] 

58 ), 

59 ) -> Callable[P, Result[T, E]] | Callable[P, Coroutine[Any, Any, Result[T, E]]]: 

60 if inspect.iscoroutinefunction(func): 

61 

62 @functools.wraps(func) 

63 async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, E]: 

64 cache_key = get_cache_key( 

65 func.__name__, args, cast(dict[str, Any], kwargs) 

66 ) 

67 now = time.monotonic() 

68 

69 if cache_key in _cache: 

70 expiry, value = _cache[cache_key] 

71 if now < expiry: 

72 return Ok(value) 

73 del _cache[cache_key] 

74 

75 func_coro = cast(Callable[P, Coroutine[Any, Any, Result[T, E]]], func) 

76 result = await func_coro(*args, **kwargs) 

77 

78 match result: 

79 case Ok(value): 

80 _cache[cache_key] = (now + ttl, value) 

81 case Err(_): 

82 pass 

83 

84 return result 

85 

86 return async_wrapper 

87 

88 @functools.wraps(func) 

89 def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, E]: 

90 cache_key = get_cache_key(func.__name__, args, cast(dict[str, Any], kwargs)) 

91 now = time.monotonic() 

92 

93 if cache_key in _cache: 

94 expiry, value = _cache[cache_key] 

95 if now < expiry: 

96 return Ok(value) 

97 del _cache[cache_key] 

98 

99 func_sync = cast(Callable[P, Result[T, E]], func) 

100 result = func_sync(*args, **kwargs) 

101 

102 match result: 

103 case Ok(value): 

104 _cache[cache_key] = (now + ttl, value) 

105 case Err(_): 

106 pass 

107 

108 return result 

109 

110 return sync_wrapper 

111 

112 return decorator # type: ignore[return-value]