Coverage for python/lsst/daf/relation/_columns/_predicate.py: 51%

1# This file is part of daf_relation. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ( 

25 "flatten_logical_and", 

26 "LogicalNot", 

27 "LogicalAnd", 

28 "LogicalOr", 

29 "Predicate", 

30 "PredicateLiteral", 

31 "PredicateReference", 

32) 

33 

34import dataclasses 

35from abc import ABC, abstractmethod 

36from collections.abc import Set 

37from typing import TYPE_CHECKING, ClassVar, Literal 

38 

39from lsst.utils.classes import cached_getter 

40 

41from ._tag import ColumnTag 

42 

43if TYPE_CHECKING: 

44 from .._engine import Engine 

45 

46 

47class Predicate(ABC): 

48 """An abstract base class and factory for boolean column expressions. 

49 

50 `Predicate` inheritance is closed to the types already provided by this 

51 package, but considerable custom behavior can still be provided via the 

52 `PredicateFunction` class and an `Engine` that knows how to interpret its 

53 `~PredicateFunction.name` value. These concrete types can all be 

54 constructed via factory methods on `Predicate` itself, `ColumnExpression`, 

55 or `ColumnContainer`, so the derived types themselves only need to be 

56 referenced when writing `match` expressions that process an expression 

57 tree. See :ref:`lsst.daf.relation-overview-extensibility` for rationale 

58 and details. 

59 """ 

60 

61 def __init_subclass__(cls) -> None: 

62 assert cls.__name__ in { 

63 "PredicateLiteral", 

64 "PredicateReference", 

65 "PredicateFunction", 

66 "ColumnInContainer", 

67 "LogicalNot", 

68 "LogicalAnd", 

69 "LogicalOr", 

70 }, "Predicate inheritance is closed to predefined types in daf_relation." 

71 

72 dtype: ClassVar[type[bool]] = bool 

73 """The Python type this expression evaluates to (`type` [ `bool` ]). 

74 """ 

75 

76 @property 

77 @abstractmethod 

78 def columns_required(self) -> Set[ColumnTag]: 

79 """Columns required by this expression 

80 (`~collections.abc.Set` [ `ColumnTag` ]). 

81 

82 This includes columns required by expressions nested within this one. 

83 """ 

84 raise NotImplementedError() 

85 

86 @abstractmethod 

87 def is_supported_by(self, engine: Engine) -> bool: 

88 """Test whether the given engine is capable of evaluating this 

89 expression. 

90 

91 Parameters 

92 ---------- 

93 engine : `Engine` 

94 Engine to test. 

95 

96 Returns 

97 ------- 

98 supported : `bool` 

99 Whether the engine supports this expression and all expressions 

100 nested within it. 

101 """ 

102 raise NotImplementedError() 

103 

104 @classmethod 

105 def literal(cls, value: bool) -> PredicateLiteral: 

106 """Construct a boolean expression that is a constant `True` or `False`. 

107 

108 Parameters 

109 ---------- 

110 value : `bool` 

111 Value for the expression. 

112 

113 Returns 

114 ------- 

115 literal : `PredicateLiteral` 

116 A boolean column expression set to the given value. 

117 """ 

118 return PredicateLiteral(value) 

119 

120 @classmethod 

121 def reference(cls, tag: ColumnTag) -> PredicateReference: 

122 """Construct an expression that refers to a boolean column in a 

123 relation. 

124 

125 Parameters 

126 ---------- 

127 tag : `ColumnTag` 

128 Identifier for the column to reference. 

129 

130 Returns 

131 ------- 

132 reference : `PredicateReference` 

133 A column expression that refers the given relation column. 

134 """ 

135 return PredicateReference(tag) 

136 

137 def logical_not(self) -> LogicalNot: 

138 """Return a boolean expression that is the logical NOT of this one. 

139 

140 Returns 

141 ------- 

142 logical_not : `Predicate` 

143 Logical NOT expression. 

144 """ 

145 return LogicalNot(self) 

146 

147 def logical_and(*operands: Predicate) -> Predicate: 

148 """Return a boolean expression that is the logical AND of the given 

149 ones. 

150 

151 Parameters 

152 ---------- 

153 *operands : `Predicate` 

154 Existing boolean expressions to AND together. 

155 

156 Returns 

157 ------- 

158 logical_and : `Predicate` 

159 Logical AND expression. If no operands are provided, a 

160 `PredicateLiteral` for `True` is returned. If one operand is 

161 provided, it is returned directly. 

162 """ 

163 if not operands: 

164 return Predicate.literal(True) 

165 elif len(operands) == 1: 

166 return operands[0] 

167 return LogicalAnd(operands) 

168 

169 def logical_or(*operands: Predicate) -> Predicate: 

170 """Return a boolean expression that is the logical OR of the given 

171 ones. 

172 

173 Parameters 

174 ---------- 

175 *operands : `Predicate` 

176 Existing boolean expressions to OR together. 

177 

178 Returns 

179 ------- 

180 logical_and : `Predicate` 

181 Logical OR expression. If no operands are provided, a 

182 `PredicateLiteral` for `False` is returned. If one operand is 

183 provided, it is returned directly. 

184 """ 

185 if not operands: 

186 return Predicate.literal(False) 

187 elif len(operands) == 1: 

188 return operands[0] 

189 return LogicalOr(operands) 

190 

191 @abstractmethod 

192 def as_trivial(self) -> bool | None: 

193 """Attempt to simplify this expression into a constant boolean. 

194 

195 Returns 

196 ------- 

197 trivial : `bool` or `None` 

198 If `True` or `False`, the expression always evaluates to exactly 

199 that constant value. If `None`, the expression is nontrivial (or 

200 at least could not easily be simplified into a trivial expression). 

201 """ 

202 return None 

203 

204 

205@dataclasses.dataclass(frozen=True) 

206class PredicateLiteral(Predicate): 

207 """A concrete boolean column expression that is a constant `True` or 

208 `False`. 

209 """ 

210 

211 value: bool 

212 """Constant value for the expression (`bool`).""" 

213 

214 @property 

215 def columns_required(self) -> Set[ColumnTag]: 

216 # Docstring inherited. 

217 return frozenset() 

218 

219 def __str__(self) -> str: 

220 return str(self.value) 

221 

222 def is_supported_by(self, engine: Engine) -> bool: 

223 # Docstring inherited. 

224 return True 

225 

226 def as_trivial(self) -> bool: 

227 # Docstring inherited. 

228 return self.value 

229 

230 

231@dataclasses.dataclass(frozen=True) 

232class PredicateReference(Predicate): 

233 """A concrete boolean column expression that refers to a boolean relation 

234 column. 

235 """ 

236 

237 tag: ColumnTag 

238 """Identifier for the column this expression refers to (`ColumnTag`).""" 

239 

240 @property 

241 def columns_required(self) -> Set[ColumnTag]: 

242 # Docstring inherited. 

243 return {self.tag} 

244 

245 def __str__(self) -> str: 

246 return str(self.tag) 

247 

248 def is_supported_by(self, engine: Engine) -> bool: 

249 # Docstring inherited. 

250 return True 

251 

252 def as_trivial(self) -> None: 

253 # Docstring inherited. 

254 return None 

255 

256 

257@dataclasses.dataclass(frozen=True) 

258class LogicalNot(Predicate): 

259 """A concrete boolean column expression that inverts its operand.""" 

260 

261 operand: Predicate 

262 """Boolean expression to invert (`Predicate`). 

263 """ 

264 

265 @property 

266 def columns_required(self) -> Set[ColumnTag]: 

267 # Docstring inherited. 

268 return self.operand.columns_required 

269 

270 def __str__(self) -> str: 

271 return f"not ({self.operand!s})" 

272 

273 def is_supported_by(self, engine: Engine) -> bool: 

274 # Docstring inherited. 

275 return self.operand.is_supported_by(engine) 

276 

277 def as_trivial(self) -> bool | None: 

278 # Docstring inherited. 

279 if (operand_as_trivial := self.operand.as_trivial()) is None: 

280 return None 

281 return not operand_as_trivial 

282 

283 

284@dataclasses.dataclass(frozen=True) 

285class LogicalAnd(Predicate): 

286 """A concrete boolean column expression that ANDs its operands.""" 

287 

288 operands: tuple[Predicate, ...] 

289 """Boolean expressions to combine (`tuple` [ `Predicate`, ... ]).""" 

290 

291 @property 

292 @cached_getter 

293 def columns_required(self) -> Set[ColumnTag]: 

294 # Docstring inherited. 

295 result: set[ColumnTag] = set() 

296 for operand in self.operands: 

297 result.update(operand.columns_required) 

298 return result 

299 

300 def __str__(self) -> str: 

301 return " and ".join( 

302 str(operand) if not isinstance(operand, LogicalOr) else f"({operand!s})" 

303 for operand in self.operands 

304 ) 

305 

306 def is_supported_by(self, engine: Engine) -> bool: 

307 # Docstring inherited. 

308 return all(operand.is_supported_by(engine) for operand in self.operands) 

309 

310 def as_trivial(self) -> bool | None: 

311 # Docstring inherited. 

312 result: bool | None = True 

313 for operand in self.operands: 

314 if (operand_as_trivial := operand.as_trivial()) is False: 

315 return False 

316 elif operand_as_trivial is None: 

317 result = None 

318 return result 

319 

320 

321@dataclasses.dataclass(frozen=True) 

322class LogicalOr(Predicate): 

323 """A concrete boolean column expression that ORs its operands.""" 

324 

325 operands: tuple[Predicate, ...] 

326 """Boolean expressions to combine (`tuple` [ `Predicate`, ... ]).""" 

327 

328 def __str__(self) -> str: 

329 return " or ".join(str(operand) for operand in self.operands) 

330 

331 @property 

332 @cached_getter 

333 def columns_required(self) -> Set[ColumnTag]: 

334 # Docstring inherited. 

335 result: set[ColumnTag] = set() 

336 for operand in self.operands: 

337 result.update(operand.columns_required) 

338 return result 

339 

340 def is_supported_by(self, engine: Engine) -> bool: 

341 # Docstring inherited. 

342 return all(operand.is_supported_by(engine) for operand in self.operands) 

343 

344 def as_trivial(self) -> bool | None: 

345 # Docstring inherited. 

346 result: bool | None = False 

347 for operand in self.operands: 

348 if (operand_as_trivial := operand.as_trivial()) is True: 

349 return True 

350 elif operand_as_trivial is None: 

351 result = None 

352 return result 

353 

354 

355def flatten_logical_and(predicate: Predicate) -> list[Predicate] | Literal[False]: 

356 """Flatten all logical AND operations in predicate into a `list`. 

357 

358 Parameters 

359 ---------- 

360 predicate : `Predicate` 

361 Original expression to flatten. 

362 

363 Returns 

364 ------- 

365 flat : `list` [ `Predicate` ] or `False` 

366 A list of predicates that could be combined with AND to reproduce the 

367 original expression, or `False` if the predicate is 

368 `trivially false <is_trivial>`. 

369 

370 Notes 

371 ----- 

372 This algorithm is not guaranteed to descend into nested OR or NOT 

373 operations, but it does descend into nested AND operations. 

374 """ 

375 match predicate: 

376 case LogicalAnd(operands=operands): 

377 result: list[Predicate] = [] 

378 for operand in operands: 

379 if (nested_result := flatten_logical_and(operand)) is False: 

380 return False 

381 result.extend(nested_result) 

382 return result 

383 case PredicateLiteral(value=value): 

384 if value: 

385 return [] 

386 else: 

387 return False 

388 return [predicate]