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

43 statements  

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

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, xxh3_128_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 # Serialise the data. 

56 data = serialize(data) 

57 

58 # Hash the data and affix its length, preceded by a hyphen (to reduce the likelihood of collisions). 

59 return f'{xxh3_128_hexdigest(data)}{len(data)}' 

60 

61def shorthash(data: Any) -> str: 

62 """Hash the given data.""" 

63 

64 # Serialise the data. 

65 data = serialize(data) 

66 

67 # Hash the data and affix its length, preceded by a hyphen (to reduce the likelihood of collisions). 

68 return f'{xxh3_64_hexdigest(data)}{len(data)}' 

69 

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

71 """Delete the provided cache.""" 

72 

73 # Remove the cache directory and all its contents. 

74 shutil.rmtree(dir, ignore_errors=True) 

75 

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

77 """Clear the provided cache.""" 

78 

79 # Delete the cache. 

80 delete(dir) 

81 

82 # Recreate the cache directory. 

83 os.makedirs(dir, exist_ok=True) 

84 

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

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

87 

88 # Iterate over keys in the cache. 

89 for file in os.listdir(dir): 

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

91 

92 # Lock the entry before reading it. 

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

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

95 timestamp = os.path.getmtime(path) 

96 

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

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

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

100 os.remove(path)