Coverage for databooks/data_models/base.py: 91%

1"""Data models - Base Pydantic model with custom methods.""" 

2 

3from __future__ import annotations 

4 

5from abc import abstractmethod 

6from collections import UserList 

7from typing import Any, Dict, Generic, Iterable, List, TypeVar, cast, overload 

8 

9from pydantic import BaseModel, ConfigDict, create_model 

10from typing_extensions import Protocol, runtime_checkable 

11 

12T = TypeVar("T") 

13 

14 

15@runtime_checkable 

16class DiffModel(Iterable, Protocol): 

17 """Protocol for static type checking.""" 

18 

19 is_diff: bool 

20 

21 def resolve(self, *args: Any, **kwargs: Any) -> DatabooksBase: 

22 """Protocol method that returns a valid base object.""" 

23 

24 

25class BaseCells(UserList, Generic[T]): 

26 """Base abstract class for notebook cells.""" 

27 

28 @abstractmethod 

29 def resolve(self, **kwargs: Any) -> list: 

30 """Return valid notebook cells from differences.""" 

31 raise NotImplementedError 

32 

33 

34@overload 

35def resolve( 

36 model: DiffModel, 

37 **kwargs: Any, 

38) -> DatabooksBase: ... 

39 

40 

41@overload 

42def resolve( 

43 model: BaseCells, 

44 **kwargs: Any, 

45) -> List[T]: ... 

46 

47 

48def resolve( 

49 model: DiffModel | BaseCells, 

50 *, 

51 keep_first: bool = True, 

52 ignore_none: bool = True, 

53 **kwargs: Any, 

54) -> DatabooksBase | List[T]: 

55 """ 

56 Resolve differences for 'diff models'. 

57 

58 Return instance alike the parent class `databooks.data_models.base.DatabooksBase`. 

59 :param model: DiffModel that is to be resolved (self when added as a method to a 

60 class 

61 :param keep_first: Whether to keep the information from the prior in the 

62 'diff model' or the latter 

63 :param ignore_none: Whether to ignore `None` values if encountered, and use the 

64 other field value 

65 :return: Model with selected fields from the differences 

66 """ 

67 field_d = dict(model) 

68 is_diff = field_d.pop("is_diff") 

69 if not is_diff: 

70 raise TypeError("Can only resolve dynamic 'diff models' (when `is_diff=True`).") 

71 

72 res_vals: Dict[str, Any] = {} 

73 for name, value in field_d.items(): 

74 if isinstance(value, (DiffModel, BaseCells)): 

75 res_vals[name] = value.resolve( 

76 keep_first=keep_first, ignore_none=ignore_none, **kwargs 

77 ) 

78 else: 

79 res_vals[name] = ( 

80 value[keep_first] 

81 if value[not keep_first] is None and ignore_none 

82 else value[not keep_first] 

83 ) 

84 

85 return type(model).mro()[1](**res_vals) 

86 

87 

88class DatabooksBase(BaseModel): 

89 """Base Pydantic class with extras on managing fields.""" 

90 

91 model_config = ConfigDict(extra="allow") 

92 

93 def remove_fields( 

94 self, 

95 fields: Iterable[str], 

96 *, 

97 recursive: bool = False, 

98 missing_ok: bool = False, 

99 ) -> None: 

100 """ 

101 Remove selected fields. 

102 

103 :param fields: Fields to remove 

104 :param recursive: Whether to remove the fields recursively in case of nested 

105 models 

106 :param missing_ok: Whether to raise errors in case field is missing 

107 :return: 

108 """ 

109 d_model = dict(self) 

110 for field in fields: 

111 field_val = d_model.get(field) if missing_ok else d_model[field] 

112 if recursive and isinstance(field_val, DatabooksBase): 

113 field_val.remove_fields(fields) 

114 elif field in d_model: 

115 delattr(self, field) 

116 

117 def __str__(self) -> str: 

118 """Return outputs of __repr__.""" 

119 return repr(self) 

120 

121 def __sub__(self, other: DatabooksBase) -> DiffModel: 

122 """ 

123 Subtraction between `databooks.data_models.base.DatabooksBase` objects. 

124 

125 The difference basically return models that replace each fields by a tuple, 

126 where for each field we have `field = (self_value, other_value)` 

127 """ 

128 if type(self) is not type(other): 

129 raise TypeError( 

130 f"Unsupported operand types for `-`: `{type(self).__name__}` and" 

131 f" `{type(other).__name__}`" 

132 ) 

133 

134 # Get field and values for each instance 

135 self_d = dict(self) 

136 other_d = dict(other) 

137 

138 # Build dict with {field: (type, value)} for each field 

139 fields_d: Dict[str, Any] = {} 

140 for name in self_d.keys() | other_d.keys(): 

141 self_val = self_d.get(name) 

142 other_val = other_d.get(name) 

143 if type(self_val) is type(other_val) and all( 

144 isinstance(val, (DatabooksBase, BaseCells)) 

145 for val in (self_val, other_val) 

146 ): 

147 # Recursively get the diffs for nested models; the `all(...)` check 

148 # above does not narrow the values for the type checker 

149 diff = self_val - other_val # ty: ignore[unsupported-operator] 

150 fields_d[name] = (Any, diff) 

151 else: 

152 fields_d[name] = (tuple, (self_val, other_val)) 

153 

154 # Build Pydantic models dynamically 

155 DiffInstance = create_model( 

156 "Diff" + type(self).__name__, 

157 __base__=type(self), 

158 __namespace__={"resolve": resolve}, 

159 is_diff=(bool, True), 

160 **fields_d, 

161 ) 

162 

163 return cast(DiffModel, DiffInstance()) # it'll be filled in with the defaults