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

39 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-03-14 21:04 +1100

1import os 

2import shutil 

3from datetime import datetime, timedelta 

4from typing import Any, Union 

5 

6from filelock import FileLock 

7from xxhash import xxh3_64_hexdigest 

8 

9from .serialization import deserialize, serialize 

10 

11NOT_IN_CACHE = object() 

12"""A sentinel object that flags that a key is not in the cache.""" 

13 

14def set(key: str, value: Any, dir: str) -> None: 

15 """Set the given key of the provided cache to the specified value.""" 

16 

17 path = f'{dir}/{key}.msgpack' 

18 

19 # Lock the entry before writing to it. 

20 with FileLock(f'{path}.lock'), \ 

21 open(path, 'wb') as file: 

22 file.write(serialize(value)) 

23 

24def get(key: str, dir: str, expiry: Union[int, float, timedelta, None] = None) -> Any: 

25 """Get the value of the given key from the provided cache if it is not expired.""" 

26 

27 path = f'{dir}/{key}.msgpack' 

28 

29 # If the key does not exist in the cache, return `NOT_IN_CACHE`. 

30 if not os.path.exists(path): 

31 return NOT_IN_CACHE 

32 

33 # Lock the entry. 

34 with FileLock(f'{path}.lock'): 

35 # Handle expiry if necessary. 

36 if expiry is not None: 

37 # Get the time at which the key was last set. 

38 timestamp = os.path.getmtime(path) 

39 

40 # If the entry is expired, remove it from the cache and return `NOT_IN_CACHE`. 

41 if isinstance(expiry, timedelta) and datetime.fromtimestamp(timestamp) + expiry < datetime.now() \ 

42 or timestamp + expiry < datetime.now().timestamp(): 

43 # Remove the entry. 

44 os.remove(path) 

45 

46 return NOT_IN_CACHE 

47 

48 # Read, deserialize and return the value. 

49 with open(path, 'rb') as file: 

50 return deserialize(file.read()) 

51 

52def hash(data: Any) -> str: 

53 """Hash the given data.""" 

54 

55 return xxh3_64_hexdigest(serialize(data)) 

56 

57def delete(dir: str) -> None: 

58 """Delete the provided cache.""" 

59 

60 # Remove the cache directory and all its contents. 

61 shutil.rmtree(dir, ignore_errors=True) 

62 

63def clear(dir: str) -> None: 

64 """Clear the provided cache.""" 

65 

66 # Delete the cache. 

67 delete(dir) 

68 

69 # Recreate the cache directory. 

70 os.makedirs(dir, exist_ok=True) 

71 

72def flush(dir: str, expiry: Union[int, float, timedelta, None]) -> None: 

73 """Flush expired keys from the provided cache.""" 

74 

75 # Iterate over keys in the cache. 

76 for file in os.listdir(dir): 

77 path = f'{dir}/{file}' 

78 

79 # Lock the entry before reading it. 

80 with FileLock(f'{path}.lock'): 

81 # Get the time at which the key was last set. 

82 timestamp = os.path.getmtime(path) 

83 

84 # If the entry is expired, remove it from the cache. 

85 if (isinstance(expiry, timedelta) and datetime.fromtimestamp(timestamp) + expiry < datetime.now()) \ 

86 or (timestamp + expiry < datetime.now().timestamp()): 

87 os.remove(path)