Coverage for slidge/db/track.py: 100%

1from __future__ import annotations 

2 

3import functools 

4from collections.abc import Callable, Sequence 

5from typing import Any, cast 

6 

7from ..util.types import ExtraAttrMapping, _AttrT 

8 

9 

10def _notify_after[F: Callable[..., Any]](method: F) -> F: 

11 @functools.wraps(method) 

12 def wrapper( 

13 self: TrackedDict | TrackedList, *args: object, **kwargs: object 

14 ) -> object: 

15 result = method(self, *args, **kwargs) 

16 self._notify() 

17 return result 

18 

19 return cast(F, wrapper) 

20 

21 

22def _notify_on[C: type](*names: str) -> Callable[[C], C]: 

23 def decorator(cls: C) -> C: 

24 for name in names: 

25 setattr(cls, name, _notify_after(getattr(cls, name))) 

26 return cls 

27 

28 return decorator 

29 

30 

31def _snapshot(value: object) -> object: 

32 """Recursively copy tracked containers into plain dicts and lists.""" 

33 if isinstance(value, dict): 

34 return {k: _snapshot(v) for k, v in value.items()} 

35 if isinstance(value, list): 

36 return list(value) # TrackedList only holds scalars 

37 return value 

38 

39 

40@_notify_on("__delitem__", "popitem", "clear", "pop") 

41class TrackedDict[T: ExtraAttrMapping = Any](dict[str, _AttrT]): 

42 def __init__(self, data: T, on_change: Callable[[T], None]) -> None: 

43 self._on_change = on_change 

44 super().__init__(self._convert(data)) 

45 

46 def _wrap(self, value: _AttrT) -> _AttrT: 

47 if isinstance(value, dict): 

48 # Nested dicts report changes to this dict, which then notifies its own parent. 

49 return cast(_AttrT, TrackedDict(value, lambda _: self._notify())) 

50 if isinstance(value, list): 

51 return cast(_AttrT, TrackedList(value, self._notify)) 

52 return value 

53 

54 def _convert(self, data: ExtraAttrMapping) -> dict[str, _AttrT]: 

55 return {k: self._wrap(v) for k, v in data.items()} 

56 

57 def _notify(self) -> None: 

58 self._on_change(cast(T, _snapshot(self))) 

59 

60 def __setitem__(self, key: str, value: _AttrT) -> None: 

61 super().__setitem__(key, self._wrap(value)) 

62 self._notify() 

63 

64 def setdefault(self, *args: object) -> None: 

65 raise NotImplementedError 

66 

67 def update(self, data: ExtraAttrMapping) -> None: # type:ignore[override] # ty:ignore 

68 super().update(self._convert(data)) 

69 self._notify() 

70 

71 def get(self, key: str, default: _AttrT = None) -> _AttrT | None: # type:ignore[override] # ty:ignore 

72 value = super().get(key) 

73 if value is None: 

74 self[key] = default # __setitem__ wraps and notifies 

75 return self[key] 

76 

77 

78@_notify_on( 

79 "__setitem__", 

80 "__delitem__", 

81 "__iadd__", 

82 "__imul__", 

83 "append", 

84 "extend", 

85 "insert", 

86 "remove", 

87 "pop", 

88 "clear", 

89 "sort", 

90 "reverse", 

91) 

92class TrackedList[T: str | float = Any](list[T]): 

93 def __init__(self, data: Sequence[T], on_change: Callable[[], None]) -> None: 

94 super().__init__(data) 

95 self._notify = on_change