Coverage for python/lsst/daf/butler/_registry_shim.py: 41%

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__ = ("RegistryShim",) 

31 

32import contextlib 

33from collections.abc import Iterable, Iterator, Mapping, Sequence 

34from typing import TYPE_CHECKING, Any 

35 

36from ._collection_type import CollectionType 

37from ._dataset_ref import DatasetId, DatasetIdGenEnum, DatasetRef 

38from ._dataset_type import DatasetType 

39from ._exceptions import CalibrationLookupError 

40from ._storage_class import StorageClassFactory 

41from ._timespan import Timespan 

42from .dimensions import ( 

43 DataCoordinate, 

44 DataId, 

45 DimensionElement, 

46 DimensionGroup, 

47 DimensionRecord, 

48 DimensionUniverse, 

49) 

50from .registry._collection_summary import CollectionSummary 

51from .registry._defaults import RegistryDefaults 

52from .registry._exceptions import NoDefaultCollectionError 

53from .registry._registry_base import RegistryBase 

54from .registry.queries._query_common import resolve_collections 

55 

56if TYPE_CHECKING: 

57 from .direct_butler import DirectButler 

58 from .registry._registry import CollectionArgType 

59 from .registry.interfaces import ObsCoreTableManager 

60 

61 

62class RegistryShim(RegistryBase): 

63 """Implementation of `Registry` interface exposed to clients by `Butler`. 

64 

65 Parameters 

66 ---------- 

67 butler : `DirectButler` 

68 Data butler instance. 

69 

70 Notes 

71 ----- 

72 This shim implementation of `Registry` forwards all methods to an actual 

73 Registry instance which is internal to Butler or to Butler methods. Its 

74 purpose is to provide a stable interface to many client-visible operations 

75 while we perform re-structuring of Registry and Butler implementations. 

76 """ 

77 

78 def __init__(self, butler: DirectButler): 

79 super().__init__(butler) 

80 self._registry = butler._registry 

81 

82 def isWriteable(self) -> bool: 

83 # Docstring inherited from a base class. 

84 return self._registry.isWriteable() 

85 

86 @property 

87 def dimensions(self) -> DimensionUniverse: 

88 # Docstring inherited from a base class. 

89 return self._registry.dimensions 

90 

91 @property 

92 def defaults(self) -> RegistryDefaults: 

93 # Docstring inherited from a base class. 

94 return self._registry.defaults 

95 

96 @defaults.setter 

97 def defaults(self, value: RegistryDefaults) -> None: 

98 # Docstring inherited from a base class. 

99 self._registry.defaults = value 

100 

101 def refresh(self) -> None: 

102 # Docstring inherited from a base class. 

103 self._registry.refresh() 

104 

105 def refresh_collection_summaries(self) -> None: 

106 # Docstring inherited from a base class. 

107 self._registry.refresh_collection_summaries() 

108 

109 def caching_context(self) -> contextlib.AbstractContextManager[None]: 

110 # Docstring inherited from a base class. 

111 return self._butler._caching_context() 

112 

113 @contextlib.contextmanager 

114 def transaction(self, *, savepoint: bool = False) -> Iterator[None]: 

115 # Docstring inherited from a base class. 

116 with self._registry.transaction(savepoint=savepoint): 

117 yield 

118 

119 def resetConnectionPool(self) -> None: 

120 # Docstring inherited from a base class. 

121 self._registry.resetConnectionPool() 

122 

123 def registerCollection( 

124 self, name: str, type: CollectionType = CollectionType.TAGGED, doc: str | None = None 

125 ) -> bool: 

126 # Docstring inherited from a base class. 

127 return self._registry.registerCollection(name, type, doc) 

128 

129 def getCollectionType(self, name: str) -> CollectionType: 

130 # Docstring inherited from a base class. 

131 return self._registry.getCollectionType(name) 

132 

133 def registerRun(self, name: str, doc: str | None = None) -> bool: 

134 # Docstring inherited from a base class. 

135 return self._registry.registerRun(name, doc) 

136 

137 def removeCollection(self, name: str) -> None: 

138 # Docstring inherited from a base class. 

139 self._registry.removeCollection(name) 

140 

141 def getCollectionChain(self, parent: str) -> Sequence[str]: 

142 # Docstring inherited from a base class. 

143 return self._registry.getCollectionChain(parent) 

144 

145 def setCollectionChain(self, parent: str, children: Any, *, flatten: bool = False) -> None: 

146 # Docstring inherited from a base class. 

147 self._registry.setCollectionChain(parent, children, flatten=flatten) 

148 

149 def getCollectionParentChains(self, collection: str) -> set[str]: 

150 # Docstring inherited from a base class. 

151 return self._registry.getCollectionParentChains(collection) 

152 

153 def getCollectionDocumentation(self, collection: str) -> str | None: 

154 # Docstring inherited from a base class. 

155 return self._registry.getCollectionDocumentation(collection) 

156 

157 def setCollectionDocumentation(self, collection: str, doc: str | None) -> None: 

158 # Docstring inherited from a base class. 

159 self._registry.setCollectionDocumentation(collection, doc) 

160 

161 def getCollectionSummary(self, collection: str) -> CollectionSummary: 

162 # Docstring inherited from a base class. 

163 return self._registry.getCollectionSummary(collection) 

164 

165 def registerDatasetType(self, datasetType: DatasetType) -> bool: 

166 # Docstring inherited from a base class. 

167 return self._registry.registerDatasetType(datasetType) 

168 

169 def removeDatasetType(self, name: str | tuple[str, ...]) -> None: 

170 # Docstring inherited from a base class. 

171 self._registry.removeDatasetType(name) 

172 

173 def getDatasetType(self, name: str) -> DatasetType: 

174 # Docstring inherited from a base class. 

175 return self._registry.getDatasetType(name) 

176 

177 def supportsIdGenerationMode(self, mode: DatasetIdGenEnum) -> bool: 

178 # Docstring inherited from a base class. 

179 return self._registry.supportsIdGenerationMode(mode) 

180 

181 def findDataset( 

182 self, 

183 datasetType: DatasetType | str, 

184 dataId: DataId | None = None, 

185 *, 

186 collections: CollectionArgType | None = None, 

187 timespan: Timespan | None = None, 

188 datastore_records: bool = False, 

189 **kwargs: Any, 

190 ) -> DatasetRef | None: 

191 # Docstring inherited from a base class. 

192 if not isinstance(datasetType, DatasetType): 

193 datasetType = self.getDatasetType(datasetType) 

194 

195 dataId = DataCoordinate.standardize( 

196 dataId, 

197 dimensions=datasetType.dimensions, 

198 universe=self.dimensions, 

199 defaults=self.defaults.dataId, 

200 **kwargs, 

201 ) 

202 

203 with self._butler.query() as query: 

204 resolved_collections = resolve_collections(self._butler, collections) 

205 if not resolved_collections: 

206 if collections is None: 

207 raise NoDefaultCollectionError("No collections provided, and no default collections set") 

208 else: 

209 return None 

210 

211 if datasetType.isCalibration() and timespan is None: 

212 # Filter out calibration collections, because with no timespan 

213 # we have no way of selecting a dataset from them. 

214 collection_info = self._butler.collections.query_info( 

215 resolved_collections, flatten_chains=True 

216 ) 

217 resolved_collections = [ 

218 info.name for info in collection_info if info.type != CollectionType.CALIBRATION 

219 ] 

220 if not resolved_collections: 

221 return None 

222 

223 result = query.datasets(datasetType, resolved_collections, find_first=True).limit(2) 

224 dataset_type_name = result.dataset_type.name 

225 # Search only on the 'required' dimensions for the dataset type. 

226 # Any extra values provided by the user are ignored. 

227 minimal_data_id = DataCoordinate.standardize( 

228 dataId.subset(datasetType.dimensions.required).required, universe=self.dimensions 

229 ) 

230 result = result.where(minimal_data_id) 

231 if ( 

232 datasetType.isCalibration() 

233 and timespan is not None 

234 and (timespan.begin is not None or timespan.end is not None) 

235 ): 

236 timespan_column = query.expression_factory[dataset_type_name].timespan 

237 # The 'logical_or(timespan.is_null)' allows non-calibration 

238 # collections to participate in the search, with the assumption 

239 # that they are valid for any time range. 

240 result = result.where(timespan_column.overlaps(timespan).logical_or(timespan_column.is_null)) 

241 

242 datasets = list(result) 

243 if len(datasets) == 1: 

244 ref = datasets[0] 

245 if dataId.hasRecords(): 

246 ref = ref.expanded(dataId) 

247 # Propagate storage class from user-provided DatasetType, which 

248 # may not match the definition in the database. 

249 ref = ref.overrideStorageClass(datasetType.storageClass_name) 

250 if datastore_records: 

251 ref = self._registry.get_datastore_records(ref) 

252 return ref 

253 elif len(datasets) == 0: 

254 return None 

255 else: 

256 raise CalibrationLookupError( 

257 f"Ambiguous calibration lookup for {datasetType} with timespan {timespan}" 

258 f" in collections {resolved_collections}." 

259 ) 

260 

261 def insertDatasets( 

262 self, 

263 datasetType: DatasetType | str, 

264 dataIds: Iterable[DataId], 

265 run: str | None = None, 

266 expand: bool = True, 

267 idGenerationMode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE, 

268 ) -> list[DatasetRef]: 

269 # Docstring inherited from a base class. 

270 return self._registry.insertDatasets(datasetType, dataIds, run, expand, idGenerationMode) 

271 

272 def _importDatasets(self, datasets: Iterable[DatasetRef], expand: bool = True) -> list[DatasetRef]: 

273 # Docstring inherited from a base class. 

274 return self._registry._importDatasets(datasets, expand) 

275 

276 def getDataset(self, id: DatasetId) -> DatasetRef | None: 

277 # Docstring inherited from a base class. 

278 return self._registry.getDataset(id) 

279 

280 def removeDatasets(self, refs: Iterable[DatasetRef]) -> None: 

281 # Docstring inherited from a base class. 

282 self._registry.removeDatasets(refs) 

283 

284 def associate(self, collection: str, refs: Iterable[DatasetRef]) -> None: 

285 # Docstring inherited from a base class. 

286 self._registry.associate(collection, refs) 

287 

288 def disassociate(self, collection: str, refs: Iterable[DatasetRef]) -> None: 

289 # Docstring inherited from a base class. 

290 self._registry.disassociate(collection, refs) 

291 

292 def certify(self, collection: str, refs: Iterable[DatasetRef], timespan: Timespan) -> None: 

293 # Docstring inherited from a base class. 

294 self._registry.certify(collection, refs, timespan) 

295 

296 def decertify( 

297 self, 

298 collection: str, 

299 datasetType: str | DatasetType, 

300 timespan: Timespan, 

301 *, 

302 dataIds: Iterable[DataId] | None = None, 

303 ) -> None: 

304 # Docstring inherited from a base class. 

305 self._registry.decertify(collection, datasetType, timespan, dataIds=dataIds) 

306 

307 def getDatasetLocations(self, ref: DatasetRef) -> Iterable[str]: 

308 # Docstring inherited from a base class. 

309 return self._registry.getDatasetLocations(ref) 

310 

311 def expandDataId( 

312 self, 

313 dataId: DataId | None = None, 

314 *, 

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

316 records: Mapping[str, DimensionRecord | None] | None = None, 

317 withDefaults: bool = True, 

318 **kwargs: Any, 

319 ) -> DataCoordinate: 

320 # Docstring inherited from a base class. 

321 return self._registry.expandDataId( 

322 dataId, dimensions=dimensions, records=records, withDefaults=withDefaults, **kwargs 

323 ) 

324 

325 def insertDimensionData( 

326 self, 

327 element: DimensionElement | str, 

328 *data: Mapping[str, Any] | DimensionRecord, 

329 conform: bool = True, 

330 replace: bool = False, 

331 skip_existing: bool = False, 

332 ) -> None: 

333 # Docstring inherited from a base class. 

334 self._registry.insertDimensionData( 

335 element, *data, conform=conform, replace=replace, skip_existing=skip_existing 

336 ) 

337 

338 def syncDimensionData( 

339 self, 

340 element: DimensionElement | str, 

341 row: Mapping[str, Any] | DimensionRecord, 

342 conform: bool = True, 

343 update: bool = False, 

344 ) -> bool | dict[str, Any]: 

345 # Docstring inherited from a base class. 

346 return self._registry.syncDimensionData(element, row, conform, update) 

347 

348 def queryDatasetTypes( 

349 self, 

350 expression: Any = ..., 

351 *, 

352 missing: list[str] | None = None, 

353 ) -> Iterable[DatasetType]: 

354 # Docstring inherited from a base class. 

355 return self._registry.queryDatasetTypes(expression, missing=missing) 

356 

357 def queryCollections( 

358 self, 

359 expression: Any = ..., 

360 datasetType: DatasetType | None = None, 

361 collectionTypes: Iterable[CollectionType] | CollectionType = CollectionType.all(), 

362 flattenChains: bool = False, 

363 includeChains: bool | None = None, 

364 ) -> Sequence[str]: 

365 # Docstring inherited from a base class. 

366 return self._registry.queryCollections( 

367 expression, datasetType, collectionTypes, flattenChains, includeChains 

368 ) 

369 

370 @property 

371 def obsCoreTableManager(self) -> ObsCoreTableManager | None: 

372 # Docstring inherited from a base class. 

373 return self._registry.obsCoreTableManager 

374 

375 @property 

376 def storageClasses(self) -> StorageClassFactory: 

377 return self._registry.storageClasses