Coverage for python/lsst/daf/butler/direct_query_results.py: 55%

110 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-13 10:57 +0000

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

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

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

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

18# (at your option) any later version. 

19# 

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

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

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

23# GNU General Public License for more details. 

24# 

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

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

27 

28from __future__ import annotations 

29 

30__all__ = [ 

31 "DirectDataCoordinateQueryResults", 

32 "DirectDatasetQueryResults", 

33 "DirectDimensionRecordQueryResults", 

34 "DirectSingleTypeDatasetQueryResults", 

35] 

36 

37import contextlib 

38from collections.abc import Iterable, Iterator 

39from typing import TYPE_CHECKING, Any 

40 

41from ._query_results import ( 

42 DataCoordinateQueryResults, 

43 DatasetQueryResults, 

44 DimensionRecordQueryResults, 

45 SingleTypeDatasetQueryResults, 

46) 

47from .registry import queries as registry_queries 

48 

49if TYPE_CHECKING: 

50 from ._dataset_ref import DatasetRef 

51 from ._dataset_type import DatasetType 

52 from .dimensions import DataCoordinate, DimensionElement, DimensionGroup, DimensionRecord 

53 

54 

55class DirectDataCoordinateQueryResults(DataCoordinateQueryResults): 

56 """Implementation of `DataCoordinateQueryResults` using query result 

57 obtained from registry. 

58 

59 Parameters 

60 ---------- 

61 registry_query_result : \ 

62 `~lsst.daf.butler.registry.queries.DataCoordinateQueryResults` 

63 Query result from Registry. 

64 """ 

65 

66 def __init__(self, registry_query_result: registry_queries.DataCoordinateQueryResults): 

67 self._registry_query_result = registry_query_result 

68 

69 def __iter__(self) -> Iterator[DataCoordinate]: 

70 return iter(self._registry_query_result) 

71 

72 @property 

73 def dimensions(self) -> DimensionGroup: 

74 # Docstring inherited. 

75 return self._registry_query_result.dimensions 

76 

77 def has_full(self) -> bool: 

78 # Docstring inherited. 

79 return self._registry_query_result.hasFull() 

80 

81 def has_records(self) -> bool: 

82 # Docstring inherited. 

83 return self._registry_query_result.hasRecords() 

84 

85 @contextlib.contextmanager 

86 def materialize(self) -> Iterator[DataCoordinateQueryResults]: 

87 with self._registry_query_result.materialize() as result: 

88 yield DirectDataCoordinateQueryResults(result) 

89 

90 def expanded(self) -> DataCoordinateQueryResults: 

91 # Docstring inherited. 

92 if self.has_records(): 

93 return self 

94 return DirectDataCoordinateQueryResults(self._registry_query_result.expanded()) 

95 

96 def subset( 

97 self, 

98 dimensions: DimensionGroup | Iterable[str] | None = None, 

99 *, 

100 unique: bool = False, 

101 ) -> DataCoordinateQueryResults: 

102 # Docstring inherited. 

103 return DirectDataCoordinateQueryResults(self._registry_query_result.subset(dimensions, unique=unique)) 

104 

105 def find_datasets( 

106 self, 

107 dataset_type: DatasetType | str, 

108 collections: Any, 

109 *, 

110 find_first: bool = True, 

111 ) -> DatasetQueryResults: 

112 # Docstring inherited. 

113 return DirectDatasetQueryResults( 

114 self._registry_query_result.findDatasets(dataset_type, collections, findFirst=find_first) 

115 ) 

116 

117 def find_related_datasets( 

118 self, 

119 dataset_type: DatasetType | str, 

120 collections: Any, 

121 *, 

122 find_first: bool = True, 

123 dimensions: DimensionGroup | Iterable[str] | None = None, 

124 ) -> Iterable[tuple[DataCoordinate, DatasetRef]]: 

125 # Docstring inherited. 

126 return self._registry_query_result.findRelatedDatasets( 

127 dataset_type, collections, findFirst=find_first, dimensions=dimensions 

128 ) 

129 

130 def count(self, *, exact: bool = True, discard: bool = False) -> int: 

131 # Docstring inherited. 

132 return self._registry_query_result.count(exact=exact, discard=discard) 

133 

134 def any(self, *, execute: bool = True, exact: bool = True) -> bool: 

135 # Docstring inherited. 

136 return self._registry_query_result.any(execute=execute, exact=exact) 

137 

138 def explain_no_results(self, execute: bool = True) -> Iterable[str]: 

139 # Docstring inherited. 

140 return self._registry_query_result.explain_no_results(execute=execute) 

141 

142 def order_by(self, *args: str) -> DataCoordinateQueryResults: 

143 # Docstring inherited. 

144 return DirectDataCoordinateQueryResults(self._registry_query_result.order_by(*args)) 

145 

146 def limit(self, limit: int | None = None, offset: int = 0) -> DataCoordinateQueryResults: 

147 # Docstring inherited. 

148 if limit is None: 

149 raise NotImplementedError("Offset without limit is temporarily unsupported.") 

150 return DirectDataCoordinateQueryResults(self._registry_query_result.limit(limit, offset)) 

151 

152 

153class DirectDatasetQueryResults(DatasetQueryResults): 

154 """Implementation of `DatasetQueryResults` using query result 

155 obtained from registry. 

156 

157 Parameters 

158 ---------- 

159 registry_query_result : \ 

160 `~lsst.daf.butler.registry.queries.DatasetQueryResults` 

161 Query result from Registry. 

162 """ 

163 

164 def __init__(self, registry_query_result: registry_queries.DatasetQueryResults): 

165 self._registry_query_result = registry_query_result 

166 

167 def __iter__(self) -> Iterator[DatasetRef]: 

168 return iter(self._registry_query_result) 

169 

170 def by_dataset_type(self) -> Iterator[SingleTypeDatasetQueryResults]: 

171 # Docstring inherited. 

172 for by_parent in self._registry_query_result.byParentDatasetType(): 

173 yield DirectSingleTypeDatasetQueryResults(by_parent) 

174 

175 @contextlib.contextmanager 

176 def materialize(self) -> Iterator[DatasetQueryResults]: 

177 # Docstring inherited. 

178 with self._registry_query_result.materialize() as result: 

179 yield DirectDatasetQueryResults(result) 

180 

181 def expanded(self) -> DatasetQueryResults: 

182 # Docstring inherited. 

183 return DirectDatasetQueryResults(self._registry_query_result.expanded()) 

184 

185 def count(self, *, exact: bool = True, discard: bool = False) -> int: 

186 # Docstring inherited. 

187 return self._registry_query_result.count(exact=exact, discard=discard) 

188 

189 def any(self, *, execute: bool = True, exact: bool = True) -> bool: 

190 # Docstring inherited. 

191 return self._registry_query_result.any(execute=execute, exact=exact) 

192 

193 def explain_no_results(self, execute: bool = True) -> Iterable[str]: 

194 # Docstring inherited. 

195 return self._registry_query_result.explain_no_results(execute=execute) 

196 

197 

198class DirectSingleTypeDatasetQueryResults(SingleTypeDatasetQueryResults): 

199 """Implementation of `SingleTypeDatasetQueryResults` using query result 

200 obtained from registry. 

201 

202 Parameters 

203 ---------- 

204 registry_query_result : \ 

205 `~lsst.daf.butler.registry.queries.ParentDatasetQueryResults` 

206 Query result from Registry. 

207 """ 

208 

209 def __init__(self, registry_query_result: registry_queries.ParentDatasetQueryResults): 

210 self._registry_query_result = registry_query_result 

211 

212 def __iter__(self) -> Iterator[DatasetRef]: 

213 return iter(self._registry_query_result) 

214 

215 def by_dataset_type(self) -> Iterator[SingleTypeDatasetQueryResults]: 

216 # Docstring inherited. 

217 yield self 

218 

219 @contextlib.contextmanager 

220 def materialize(self) -> Iterator[SingleTypeDatasetQueryResults]: 

221 # Docstring inherited. 

222 with self._registry_query_result.materialize() as result: 

223 yield DirectSingleTypeDatasetQueryResults(result) 

224 

225 @property 

226 def dataset_type(self) -> DatasetType: 

227 # Docstring inherited. 

228 return self._registry_query_result.parentDatasetType 

229 

230 @property 

231 def data_ids(self) -> DataCoordinateQueryResults: 

232 # Docstring inherited. 

233 return DirectDataCoordinateQueryResults(self._registry_query_result.dataIds) 

234 

235 def expanded(self) -> SingleTypeDatasetQueryResults: 

236 # Docstring inherited. 

237 return DirectSingleTypeDatasetQueryResults(self._registry_query_result.expanded()) 

238 

239 def count(self, *, exact: bool = True, discard: bool = False) -> int: 

240 # Docstring inherited. 

241 return self._registry_query_result.count(exact=exact, discard=discard) 

242 

243 def any(self, *, execute: bool = True, exact: bool = True) -> bool: 

244 # Docstring inherited. 

245 return self._registry_query_result.any(execute=execute, exact=exact) 

246 

247 def explain_no_results(self, execute: bool = True) -> Iterable[str]: 

248 # Docstring inherited. 

249 return self._registry_query_result.explain_no_results(execute=execute) 

250 

251 

252class DirectDimensionRecordQueryResults(DimensionRecordQueryResults): 

253 """Implementation of `DimensionRecordQueryResults` using query result 

254 obtained from registry. 

255 

256 Parameters 

257 ---------- 

258 registry_query_result : \ 

259 `~lsst.daf.butler.registry.queries.DimensionRecordQueryResults` 

260 Query result from Registry. 

261 """ 

262 

263 def __init__(self, registry_query_result: registry_queries.DimensionRecordQueryResults): 

264 self._registry_query_result = registry_query_result 

265 

266 def __iter__(self) -> Iterator[DimensionRecord]: 

267 return iter(self._registry_query_result) 

268 

269 @property 

270 def element(self) -> DimensionElement: 

271 # Docstring inherited. 

272 return self._registry_query_result.element 

273 

274 def run(self) -> DimensionRecordQueryResults: 

275 # Docstring inherited. 

276 return DirectDimensionRecordQueryResults(self._registry_query_result.run()) 

277 

278 def count(self, *, exact: bool = True, discard: bool = False) -> int: 

279 # Docstring inherited. 

280 return self._registry_query_result.count(exact=exact, discard=discard) 

281 

282 def any(self, *, execute: bool = True, exact: bool = True) -> bool: 

283 # Docstring inherited. 

284 return self._registry_query_result.any(execute=execute, exact=exact) 

285 

286 def order_by(self, *args: str) -> DimensionRecordQueryResults: 

287 # Docstring inherited. 

288 return DirectDimensionRecordQueryResults(self._registry_query_result.order_by(*args)) 

289 

290 def limit(self, limit: int | None = None, offset: int = 0) -> DimensionRecordQueryResults: 

291 # Docstring inherited. 

292 if limit is None: 

293 raise NotImplementedError("Offset without limit is temporarily unsupported.") 

294 return DirectDimensionRecordQueryResults(self._registry_query_result.limit(limit, offset)) 

295 

296 def explain_no_results(self, execute: bool = True) -> Iterable[str]: 

297 # Docstring inherited. 

298 return self._registry_query_result.explain_no_results(execute=execute)