Coverage for python/lsst/daf/relation/_columns/_container.py: 71%

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 "ColumnContainer", 

26 "ColumnRangeLiteral", 

27 "ColumnExpressionSequence", 

28 "ColumnInContainer", 

29) 

30 

31import dataclasses 

32from abc import ABC, abstractmethod 

33from collections.abc import Sequence, Set 

34from typing import TYPE_CHECKING 

35 

36from lsst.utils.classes import cached_getter 

37 

38from ._predicate import Predicate 

39from ._tag import ColumnTag 

40 

41if TYPE_CHECKING: 41 ↛ 42line 41 didn't jump to line 42, because the condition on line 41 was never true

42 from .._engine import Engine 

43 from ._expression import ColumnExpression 

44 

45 

46class ColumnContainer(ABC): 

47 """A abstract base class and factory for expressions that represent 

48 containers of multiple column values. 

49 

50 `ColumnContainer` inheritance is closed to the types already provided by 

51 this package. These concrete types can all be constructed via factory 

52 methods on `ColumnContainer` itself, so the derived types themselves only 

53 need to be referenced when writing `match` expressions that process an 

54 expression tree. See :ref:`lsst.daf.relation-overview-extensibility` for 

55 rationale and details. 

56 """ 

57 

58 def __init_subclass__(cls) -> None: 

59 assert cls.__name__ in { 

60 "ColumnRangeLiteral", 

61 "ColumnExpressionSequence", 

62 }, "ColumnContainer inheritance is closed to predefined types in daf_relation." 

63 

64 dtype: type | None 

65 """The Python type of the elements in the container (`type` or `None`). 

66 

67 Interpretation of this attribute is up to the `Engine` or other algorithms 

68 that operate on the expression tree; it is ignored by all code in the 

69 `lsst.daf.relation` package. 

70 """ 

71 

72 @property 

73 @abstractmethod 

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

75 """Columns required by this expression 

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

77 

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

79 """ 

80 raise NotImplementedError() 

81 

82 @abstractmethod 

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

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

85 expression. 

86 

87 Parameters 

88 ---------- 

89 engine : `Engine` 

90 Engine to test. 

91 

92 Returns 

93 ------- 

94 supported : `bool` 

95 Whether the engine supports this expression and all expressions 

96 nested within it. 

97 """ 

98 raise NotImplementedError() 

99 

100 def contains(self, item: ColumnExpression) -> ColumnInContainer: 

101 """Construct a boolean column expression that tests whether a scalar 

102 expression is present in this container expression. 

103 

104 Parameters 

105 ---------- 

106 item : `ColumnExpression` 

107 Item expression to test. 

108 

109 Returns 

110 ------- 

111 contains : `ColumnInContainer` 

112 Boolean column expression that tests for membership in the 

113 container. 

114 """ 

115 return ColumnInContainer(item, self) 

116 

117 @classmethod 

118 def range_literal(cls, r: range) -> ColumnRangeLiteral: 

119 """Construct a container expression from a range of indices. 

120 

121 Parameters 

122 ---------- 

123 r : `range` 

124 Range object. 

125 

126 Returns 

127 ------- 

128 container : `ColumnRangeLiteral` 

129 Container expression object representing the range. 

130 """ 

131 return ColumnRangeLiteral(r) 

132 

133 @classmethod 

134 def sequence( 

135 cls, items: Sequence[ColumnExpression], dtype: type | None = None 

136 ) -> ColumnExpressionSequence: 

137 """Construct a container expression from a sequence of item 

138 expressions. 

139 

140 Parameters 

141 ---------- 

142 items : `~collections.abc.Sequence` [ `ColumnExpression` ] 

143 Sequence of item expressions. 

144 dtype : `type`, optional 

145 The Python type of the elements in the container. 

146 

147 Returns 

148 ------- 

149 container : `ColumnExpressionSequence` 

150 Container expression object backed by the given items. 

151 """ 

152 return ColumnExpressionSequence(items, dtype) 

153 

154 

155@dataclasses.dataclass(frozen=True) 

156class ColumnRangeLiteral(ColumnContainer): 

157 """A container expression backed by a range of integer indices.""" 

158 

159 value: range 

160 """Range value that backs the expression (`range`). 

161 """ 

162 

163 dtype: type[int] = int 

164 """The Python type of the elements in the container (`type` [ `int` ] ).""" 

165 

166 @property 

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

168 # Docstring inherited. 

169 return frozenset() 

170 

171 def __str__(self) -> str: 

172 return f"[{self.value.start}:{self.value.stop}:{self.value.step}]" 

173 

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

175 # Docstring inherited. 

176 return True 

177 

178 

179@dataclasses.dataclass(frozen=True) 

180class ColumnExpressionSequence(ColumnContainer): 

181 """A container expression backed by a sequence of scalar column 

182 expressions. 

183 """ 

184 

185 items: Sequence[ColumnExpression] 

186 """Sequence of item expressions 

187 (`~collections.abc.Sequence` [ `ColumnExpression` ]). 

188 """ 

189 

190 dtype: type | None 

191 """The Python type of the elements in the container (`type` or `None`).""" 

192 

193 @property 

194 @cached_getter 

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

196 # Docstring inherited. 

197 result: set[ColumnTag] = set() 

198 for item in self.items: 

199 result.update(item.columns_required) 

200 return result 

201 

202 def __str__(self) -> str: 

203 return f"[{', '.join(str(i) for i in self.items)}]" 

204 

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

206 # Docstring inherited. 

207 return all(item.is_supported_by(engine) for item in self.items) 

208 

209 

210@dataclasses.dataclass(frozen=True) 

211class ColumnInContainer(Predicate): 

212 """A boolean column expression that tests whether a scalar column 

213 expression is present in a container expression. 

214 """ 

215 

216 item: ColumnExpression 

217 """Item to be tested (`ColumnExpression`).""" 

218 

219 container: ColumnContainer 

220 """Container to be tested (`ColumnContainer`).""" 

221 

222 @property 

223 @cached_getter 

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

225 # Docstring inherited. 

226 return self.item.columns_required | self.container.columns_required 

227 

228 def __str__(self) -> str: 

229 return f"{self.item}∈{self.container}" 

230 

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

232 # Docstring inherited. 

233 return self.item.is_supported_by(engine) and self.container.is_supported_by(engine) 

234 

235 def as_trivial(self) -> None: 

236 # Docstring inherited. 

237 return None