Coverage for src/dataknobs_data/query_logic.py: 26%

189 statements  

« prev     ^ index     » next       coverage.py v7.10.3, created at 2025-08-17 18:54 -0500

1"""Boolean logic support for complex queries.""" 

2 

3from abc import ABC, abstractmethod 

4from dataclasses import dataclass, field 

5from enum import Enum 

6from typing import Any, Dict, List, Union 

7 

8from .query import Filter, Operator 

9 

10 

11class LogicOperator(Enum): 

12 """Logical operators for combining conditions.""" 

13 AND = "and" 

14 OR = "or" 

15 NOT = "not" 

16 

17 

18class Condition(ABC): 

19 """Abstract base class for query conditions.""" 

20 

21 @abstractmethod 

22 def matches(self, record: Any) -> bool: 

23 """Check if a record matches this condition.""" 

24 pass 

25 

26 @abstractmethod 

27 def to_dict(self) -> Dict[str, Any]: 

28 """Convert condition to dictionary representation.""" 

29 pass 

30 

31 @classmethod 

32 @abstractmethod 

33 def from_dict(cls, data: Dict[str, Any]) -> "Condition": 

34 """Create condition from dictionary representation.""" 

35 pass 

36 

37 

38@dataclass 

39class FilterCondition(Condition): 

40 """A single filter condition.""" 

41 filter: Filter 

42 

43 def matches(self, record: Any) -> bool: 

44 """Check if a record matches this filter.""" 

45 from .records import Record 

46 

47 if isinstance(record, Record): 

48 value = record.get_value(self.filter.field) 

49 elif isinstance(record, dict): 

50 # Support nested field access for dicts 

51 value = record 

52 for part in self.filter.field.split('.'): 

53 if isinstance(value, dict): 

54 value = value.get(part) 

55 else: 

56 value = None 

57 break 

58 else: 

59 value = getattr(record, self.filter.field, None) 

60 

61 return self.filter.matches(value) 

62 

63 def to_dict(self) -> Dict[str, Any]: 

64 """Convert to dictionary representation.""" 

65 return { 

66 "type": "filter", 

67 "filter": self.filter.to_dict() 

68 } 

69 

70 @classmethod 

71 def from_dict(cls, data: Dict[str, Any]) -> "FilterCondition": 

72 """Create from dictionary representation.""" 

73 return cls(filter=Filter.from_dict(data["filter"])) 

74 

75 

76@dataclass 

77class LogicCondition(Condition): 

78 """A logical combination of conditions.""" 

79 operator: LogicOperator 

80 conditions: List[Condition] = field(default_factory=list) 

81 

82 def matches(self, record: Any) -> bool: 

83 """Check if a record matches this logical condition.""" 

84 if self.operator == LogicOperator.AND: 

85 # All conditions must match 

86 return all(cond.matches(record) for cond in self.conditions) 

87 elif self.operator == LogicOperator.OR: 

88 # At least one condition must match 

89 return any(cond.matches(record) for cond in self.conditions) 

90 elif self.operator == LogicOperator.NOT: 

91 # No conditions should match (or negate single condition) 

92 if len(self.conditions) == 1: 

93 return not self.conditions[0].matches(record) 

94 else: 

95 # NOT with multiple conditions = none should match 

96 return not any(cond.matches(record) for cond in self.conditions) 

97 return False 

98 

99 def to_dict(self) -> Dict[str, Any]: 

100 """Convert to dictionary representation.""" 

101 return { 

102 "type": "logic", 

103 "operator": self.operator.value, 

104 "conditions": [cond.to_dict() for cond in self.conditions] 

105 } 

106 

107 @classmethod 

108 def from_dict(cls, data: Dict[str, Any]) -> "LogicCondition": 

109 """Create from dictionary representation.""" 

110 conditions = [] 

111 for cond_data in data.get("conditions", []): 

112 if cond_data["type"] == "filter": 

113 conditions.append(FilterCondition.from_dict(cond_data)) 

114 elif cond_data["type"] == "logic": 

115 conditions.append(LogicCondition.from_dict(cond_data)) 

116 

117 return cls( 

118 operator=LogicOperator(data["operator"]), 

119 conditions=conditions 

120 ) 

121 

122 

123def condition_from_dict(data: Dict[str, Any]) -> Condition: 

124 """Factory function to create condition from dictionary.""" 

125 if data["type"] == "filter": 

126 return FilterCondition.from_dict(data) 

127 elif data["type"] == "logic": 

128 return LogicCondition.from_dict(data) 

129 else: 

130 raise ValueError(f"Unknown condition type: {data['type']}") 

131 

132 

133class QueryBuilder: 

134 """Builder for complex queries with boolean logic.""" 

135 

136 def __init__(self): 

137 """Initialize empty query builder.""" 

138 self.root_condition = None 

139 self.sort_specs = [] 

140 self.limit_value = None 

141 self.offset_value = None 

142 self.fields = None 

143 

144 def where(self, field: str, operator: Union[str, Operator], value: Any = None) -> "QueryBuilder": 

145 """Add a filter condition (defaults to AND with existing conditions).""" 

146 filter_cond = FilterCondition(Filter(field, operator, value)) 

147 

148 if self.root_condition is None: 

149 self.root_condition = filter_cond 

150 elif isinstance(self.root_condition, LogicCondition) and self.root_condition.operator == LogicOperator.AND: 

151 self.root_condition.conditions.append(filter_cond) 

152 else: 

153 # Wrap existing condition in AND 

154 self.root_condition = LogicCondition( 

155 operator=LogicOperator.AND, 

156 conditions=[self.root_condition, filter_cond] 

157 ) 

158 

159 return self 

160 

161 def and_(self, *conditions: Union["QueryBuilder", Filter, Condition]) -> "QueryBuilder": 

162 """Add AND conditions.""" 

163 logic_cond = LogicCondition(operator=LogicOperator.AND) 

164 

165 for cond in conditions: 

166 if isinstance(cond, QueryBuilder): 

167 if cond.root_condition: 

168 logic_cond.conditions.append(cond.root_condition) 

169 elif isinstance(cond, Filter): 

170 logic_cond.conditions.append(FilterCondition(cond)) 

171 elif isinstance(cond, Condition): 

172 logic_cond.conditions.append(cond) 

173 

174 if self.root_condition is None: 

175 self.root_condition = logic_cond 

176 elif isinstance(self.root_condition, LogicCondition) and self.root_condition.operator == LogicOperator.AND: 

177 self.root_condition.conditions.extend(logic_cond.conditions) 

178 else: 

179 self.root_condition = LogicCondition( 

180 operator=LogicOperator.AND, 

181 conditions=[self.root_condition, logic_cond] 

182 ) 

183 

184 return self 

185 

186 def or_(self, *conditions: Union["QueryBuilder", Filter, Condition]) -> "QueryBuilder": 

187 """Add OR conditions.""" 

188 logic_cond = LogicCondition(operator=LogicOperator.OR) 

189 

190 for cond in conditions: 

191 if isinstance(cond, QueryBuilder): 

192 if cond.root_condition: 

193 logic_cond.conditions.append(cond.root_condition) 

194 elif isinstance(cond, Filter): 

195 logic_cond.conditions.append(FilterCondition(cond)) 

196 elif isinstance(cond, Condition): 

197 logic_cond.conditions.append(cond) 

198 

199 if self.root_condition is None: 

200 self.root_condition = logic_cond 

201 else: 

202 # Always wrap in OR at top level 

203 if isinstance(self.root_condition, LogicCondition) and self.root_condition.operator == LogicOperator.OR: 

204 self.root_condition.conditions.extend(logic_cond.conditions) 

205 else: 

206 self.root_condition = LogicCondition( 

207 operator=LogicOperator.OR, 

208 conditions=[self.root_condition] + logic_cond.conditions 

209 ) 

210 

211 return self 

212 

213 def not_(self, condition: Union["QueryBuilder", Filter, Condition]) -> "QueryBuilder": 

214 """Add NOT condition.""" 

215 if isinstance(condition, QueryBuilder): 

216 not_cond = LogicCondition( 

217 operator=LogicOperator.NOT, 

218 conditions=[condition.root_condition] if condition.root_condition else [] 

219 ) 

220 elif isinstance(condition, Filter): 

221 not_cond = LogicCondition( 

222 operator=LogicOperator.NOT, 

223 conditions=[FilterCondition(condition)] 

224 ) 

225 else: 

226 not_cond = LogicCondition( 

227 operator=LogicOperator.NOT, 

228 conditions=[condition] 

229 ) 

230 

231 if self.root_condition is None: 

232 self.root_condition = not_cond 

233 elif isinstance(self.root_condition, LogicCondition) and self.root_condition.operator == LogicOperator.AND: 

234 self.root_condition.conditions.append(not_cond) 

235 else: 

236 self.root_condition = LogicCondition( 

237 operator=LogicOperator.AND, 

238 conditions=[self.root_condition, not_cond] 

239 ) 

240 

241 return self 

242 

243 def sort_by(self, field: str, order: str = "asc") -> "QueryBuilder": 

244 """Add sort specification.""" 

245 from .query import SortSpec, SortOrder 

246 

247 sort_order = SortOrder.ASC if order.lower() == "asc" else SortOrder.DESC 

248 self.sort_specs.append(SortSpec(field=field, order=sort_order)) 

249 return self 

250 

251 def limit(self, value: int) -> "QueryBuilder": 

252 """Set result limit.""" 

253 self.limit_value = value 

254 return self 

255 

256 def offset(self, value: int) -> "QueryBuilder": 

257 """Set result offset.""" 

258 self.offset_value = value 

259 return self 

260 

261 def select(self, *fields: str) -> "QueryBuilder": 

262 """Set field projection.""" 

263 self.fields = list(fields) if fields else None 

264 return self 

265 

266 def build(self) -> "ComplexQuery": 

267 """Build the final query.""" 

268 return ComplexQuery( 

269 condition=self.root_condition, 

270 sort_specs=self.sort_specs, 

271 limit_value=self.limit_value, 

272 offset_value=self.offset_value, 

273 fields=self.fields 

274 ) 

275 

276 

277@dataclass 

278class ComplexQuery: 

279 """A query with complex boolean logic support.""" 

280 

281 condition: Condition | None = None 

282 sort_specs: List = field(default_factory=list) 

283 limit_value: int | None = None 

284 offset_value: int | None = None 

285 fields: List[str] | None = None 

286 

287 def matches(self, record: Any) -> bool: 

288 """Check if a record matches this query.""" 

289 if self.condition is None: 

290 return True 

291 return self.condition.matches(record) 

292 

293 def to_simple_query(self) -> "Query": 

294 """Convert to simple Query if possible (AND filters only).""" 

295 from .query import Query 

296 

297 filters = [] 

298 

299 # Try to extract simple filters if all are AND conditions 

300 if self.condition is None: 

301 pass 

302 elif isinstance(self.condition, FilterCondition): 

303 filters.append(self.condition.filter) 

304 elif isinstance(self.condition, LogicCondition) and self.condition.operator == LogicOperator.AND: 

305 # Check if all sub-conditions are simple filters 

306 all_filters = True 

307 for cond in self.condition.conditions: 

308 if isinstance(cond, FilterCondition): 

309 filters.append(cond.filter) 

310 else: 

311 all_filters = False 

312 break 

313 

314 if not all_filters: 

315 # Can't convert complex logic to simple query 

316 raise ValueError("Cannot convert complex boolean logic to simple Query") 

317 else: 

318 raise ValueError("Cannot convert complex boolean logic to simple Query") 

319 

320 return Query( 

321 filters=filters, 

322 sort_specs=self.sort_specs, 

323 limit_value=self.limit_value, 

324 offset_value=self.offset_value, 

325 fields=self.fields 

326 ) 

327 

328 def to_dict(self) -> Dict[str, Any]: 

329 """Convert to dictionary representation.""" 

330 result = {} 

331 

332 if self.condition: 

333 result["condition"] = self.condition.to_dict() 

334 

335 if self.sort_specs: 

336 result["sort"] = [s.to_dict() for s in self.sort_specs] 

337 

338 if self.limit_value is not None: 

339 result["limit"] = self.limit_value 

340 

341 if self.offset_value is not None: 

342 result["offset"] = self.offset_value 

343 

344 if self.fields is not None: 

345 result["fields"] = self.fields 

346 

347 return result 

348 

349 @classmethod 

350 def from_dict(cls, data: Dict[str, Any]) -> "ComplexQuery": 

351 """Create from dictionary representation.""" 

352 from .query import SortSpec 

353 

354 condition = None 

355 if "condition" in data: 

356 condition = condition_from_dict(data["condition"]) 

357 

358 sort_specs = [] 

359 for sort_data in data.get("sort", []): 

360 sort_specs.append(SortSpec.from_dict(sort_data)) 

361 

362 return cls( 

363 condition=condition, 

364 sort_specs=sort_specs, 

365 limit_value=data.get("limit"), 

366 offset_value=data.get("offset"), 

367 fields=data.get("fields") 

368 )