Coverage for r11k/cache.py: 54%

1"""Simple Key-Value cache backed by disk.""" 

2 

3import os 

4 

5import json 

6import atexit 

7import hashlib 

8from r11k import util 

9 

10from typing import Any, Optional 

11 

12 

13def _encode(string: str) -> str: 

14 """Calculate cheksum string for string.""" 

15 return hashlib.sha256(string.encode('UTF-8')).hexdigest() 

16 

17 

18class KeyValueStore: 

19 """ 

20 A simple key-value store. 

21 

22 Serializes its data to disk as one json file for each key. 

23 Serialization to disk happens upon program end, through an 

24 `atexit` handler. 

25 

26 ### Example 

27 

28 If the data is 

29 >>> { 'a': 1, 'b': { 'c': 3 } } 

30 

31 then the data is stored as: 

32 

33 - `<path>` 

34 - `a.json`, *content*: `1` 

35 - `b.json`, *content*: `{"c":3}` 

36 """ 

37 

38 def __init__(self, path: str = '/tmp'): 

39 self._path: str = path 

40 self.data: dict[str, Any] = {} 

41 self._index: dict[str, str] = {} 

42 

43 util.ensure_directory(path) 

44 

45 atexit.register(self.save) 

46 

47 def save(self) -> None: 

48 """Serialize ourselves to disk.""" 

49 for key, value in self.data.items(): 

50 filename = os.path.join(self._path, key) 

51 with open(filename, 'w') as f: 

52 json.dump(value, f) 

53 

54 index_file = os.path.join(self._path, 'index.json') 

55 # Load existing index 

56 try: 

57 with open(index_file) as f: 

58 index = json.load(f) 

59 except FileNotFoundError: 

60 index = {} 

61 # Merge with current index 

62 index |= self._index 

63 

64 # Write new index 

65 with open(index_file, 'w') as f: 

66 json.dump(index, f) 

67 

68 def path(self, key: str) -> str: 

69 """Return filesystem path where this key would be stored.""" 

70 key = _encode(key) 

71 return os.path.join(self._path, key) 

72 

73 def get(self, key: str) -> Optional[Any]: 

74 """ 

75 Get the value for key, possibly loading it from disk. 

76 

77 Lazily loads it from disk. Return None if not found. 

78 """ 

79 key = _encode(key) 

80 if it := self.data.get(key): 

81 return it 

82 else: 

83 filename = os.path.join(self._path, key) 

84 if os.path.exists(filename): 84 ↛ 85line 84 didn't jump to line 85, because the condition on line 84 was never true

85 with open(filename) as f: 

86 data = json.load(f) 

87 self.data[key] = data 

88 return data 

89 else: 

90 return None 

91 

92 def put(self, key: str, value: Any) -> None: 

93 """Store the given value under key.""" 

94 _key = _encode(key) 

95 self.data[_key] = value 

96 self._index[key] = _key