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

108 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-16 10:44 +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, offset: int | None = 0) -> DataCoordinateQueryResults: 

147 # Docstring inherited. 

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

149 

150 

151class DirectDatasetQueryResults(DatasetQueryResults): 

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

153 obtained from registry. 

154 

155 Parameters 

156 ---------- 

157 registry_query_result : \ 

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

159 Query result from Registry. 

160 """ 

161 

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

163 self._registry_query_result = registry_query_result 

164 

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

166 return iter(self._registry_query_result) 

167 

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

169 # Docstring inherited. 

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

171 yield DirectSingleTypeDatasetQueryResults(by_parent) 

172 

173 @contextlib.contextmanager 

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

175 # Docstring inherited. 

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

177 yield DirectDatasetQueryResults(result) 

178 

179 def expanded(self) -> DatasetQueryResults: 

180 # Docstring inherited. 

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

182 

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

184 # Docstring inherited. 

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

186 

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

188 # Docstring inherited. 

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

190 

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

192 # Docstring inherited. 

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

194 

195 

196class DirectSingleTypeDatasetQueryResults(SingleTypeDatasetQueryResults): 

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

198 obtained from registry. 

199 

200 Parameters 

201 ---------- 

202 registry_query_result : \ 

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

204 Query result from Registry. 

205 """ 

206 

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

208 self._registry_query_result = registry_query_result 

209 

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

211 return iter(self._registry_query_result) 

212 

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

214 # Docstring inherited. 

215 yield self 

216 

217 @contextlib.contextmanager 

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

219 # Docstring inherited. 

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

221 yield DirectSingleTypeDatasetQueryResults(result) 

222 

223 @property 

224 def dataset_type(self) -> DatasetType: 

225 # Docstring inherited. 

226 return self._registry_query_result.parentDatasetType 

227 

228 @property 

229 def data_ids(self) -> DataCoordinateQueryResults: 

230 # Docstring inherited. 

231 return DirectDataCoordinateQueryResults(self._registry_query_result.dataIds) 

232 

233 def expanded(self) -> SingleTypeDatasetQueryResults: 

234 # Docstring inherited. 

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

236 

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

238 # Docstring inherited. 

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

240 

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

242 # Docstring inherited. 

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

244 

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

246 # Docstring inherited. 

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

248 

249 

250class DirectDimensionRecordQueryResults(DimensionRecordQueryResults): 

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

252 obtained from registry. 

253 

254 Parameters 

255 ---------- 

256 registry_query_result : \ 

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

258 Query result from Registry. 

259 """ 

260 

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

262 self._registry_query_result = registry_query_result 

263 

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

265 return iter(self._registry_query_result) 

266 

267 @property 

268 def element(self) -> DimensionElement: 

269 # Docstring inherited. 

270 return self._registry_query_result.element 

271 

272 def run(self) -> DimensionRecordQueryResults: 

273 # Docstring inherited. 

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

275 

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

277 # Docstring inherited. 

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

279 

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

281 # Docstring inherited. 

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

283 

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

285 # Docstring inherited. 

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

287 

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

289 # Docstring inherited. 

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

291 

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

293 # Docstring inherited. 

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