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

111 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-13 09:58 +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__ = ("Registry",) 

31 

32import contextlib 

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

34from typing import TYPE_CHECKING, Any 

35 

36from ._dataset_association import DatasetAssociation 

37from ._dataset_ref import DatasetId, DatasetIdGenEnum, DatasetRef 

38from ._dataset_type import DatasetType 

39from ._named import NameLookupMapping 

40from ._timespan import Timespan 

41from .dimensions import ( 

42 DataCoordinate, 

43 DataId, 

44 Dimension, 

45 DimensionElement, 

46 DimensionGraph, 

47 DimensionGroup, 

48 DimensionRecord, 

49 DimensionUniverse, 

50) 

51from .registry import Registry 

52from .registry._collection_summary import CollectionSummary 

53from .registry._collection_type import CollectionType 

54from .registry._defaults import RegistryDefaults 

55from .registry.queries import DataCoordinateQueryResults, DatasetQueryResults, DimensionRecordQueryResults 

56from .utils import _DefaultMarker, _Marker 

57 

58if TYPE_CHECKING: 

59 from .direct_butler import DirectButler 

60 from .registry._registry import CollectionArgType 

61 from .registry.interfaces import ObsCoreTableManager 

62 

63 

64class RegistryShim(Registry): 

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

66 

67 Parameters 

68 ---------- 

69 butler : `DirectButler` 

70 Data butler instance. 

71 

72 Notes 

73 ----- 

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

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

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

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

78 """ 

79 

80 def __init__(self, butler: DirectButler): 

81 self._butler = butler 

82 self._registry = butler._registry 

83 

84 def isWriteable(self) -> bool: 

85 # Docstring inherited from a base class. 

86 return self._registry.isWriteable() 

87 

88 @property 

89 def dimensions(self) -> DimensionUniverse: 

90 # Docstring inherited from a base class. 

91 return self._registry.dimensions 

92 

93 @property 

94 def defaults(self) -> RegistryDefaults: 

95 # Docstring inherited from a base class. 

96 return self._registry.defaults 

97 

98 @defaults.setter 

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

100 # Docstring inherited from a base class. 

101 self._registry.defaults = value 

102 

103 def refresh(self) -> None: 

104 # Docstring inherited from a base class. 

105 self._registry.refresh() 

106 

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

108 # Docstring inherited from a base class. 

109 return self._butler._caching_context() 

110 

111 @contextlib.contextmanager 

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

113 # Docstring inherited from a base class. 

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

115 yield 

116 

117 def resetConnectionPool(self) -> None: 

118 # Docstring inherited from a base class. 

119 self._registry.resetConnectionPool() 

120 

121 def registerCollection( 

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

123 ) -> bool: 

124 # Docstring inherited from a base class. 

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

126 

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

128 # Docstring inherited from a base class. 

129 return self._registry.getCollectionType(name) 

130 

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

132 # Docstring inherited from a base class. 

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

134 

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

136 # Docstring inherited from a base class. 

137 self._registry.removeCollection(name) 

138 

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

140 # Docstring inherited from a base class. 

141 return self._registry.getCollectionChain(parent) 

142 

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

144 # Docstring inherited from a base class. 

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

146 

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

148 # Docstring inherited from a base class. 

149 return self._registry.getCollectionParentChains(collection) 

150 

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

152 # Docstring inherited from a base class. 

153 return self._registry.getCollectionDocumentation(collection) 

154 

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

156 # Docstring inherited from a base class. 

157 self._registry.setCollectionDocumentation(collection, doc) 

158 

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

160 # Docstring inherited from a base class. 

161 return self._registry.getCollectionSummary(collection) 

162 

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

164 # Docstring inherited from a base class. 

165 return self._registry.registerDatasetType(datasetType) 

166 

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

168 # Docstring inherited from a base class. 

169 self._registry.removeDatasetType(name) 

170 

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

172 # Docstring inherited from a base class. 

173 return self._registry.getDatasetType(name) 

174 

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

176 # Docstring inherited from a base class. 

177 return self._registry.supportsIdGenerationMode(mode) 

178 

179 def findDataset( 

180 self, 

181 datasetType: DatasetType | str, 

182 dataId: DataId | None = None, 

183 *, 

184 collections: CollectionArgType | None = None, 

185 timespan: Timespan | None = None, 

186 **kwargs: Any, 

187 ) -> DatasetRef | None: 

188 # Docstring inherited from a base class. 

189 return self._registry.findDataset( 

190 datasetType, dataId, collections=collections, timespan=timespan, **kwargs 

191 ) 

192 

193 def insertDatasets( 

194 self, 

195 datasetType: DatasetType | str, 

196 dataIds: Iterable[DataId], 

197 run: str | None = None, 

198 expand: bool = True, 

199 idGenerationMode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE, 

200 ) -> list[DatasetRef]: 

201 # Docstring inherited from a base class. 

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

203 

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

205 # Docstring inherited from a base class. 

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

207 

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

209 # Docstring inherited from a base class. 

210 return self._registry.getDataset(id) 

211 

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

213 # Docstring inherited from a base class. 

214 self._registry.removeDatasets(refs) 

215 

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

217 # Docstring inherited from a base class. 

218 self._registry.associate(collection, refs) 

219 

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

221 # Docstring inherited from a base class. 

222 self._registry.disassociate(collection, refs) 

223 

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

225 # Docstring inherited from a base class. 

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

227 

228 def decertify( 

229 self, 

230 collection: str, 

231 datasetType: str | DatasetType, 

232 timespan: Timespan, 

233 *, 

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

235 ) -> None: 

236 # Docstring inherited from a base class. 

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

238 

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

240 # Docstring inherited from a base class. 

241 return self._registry.getDatasetLocations(ref) 

242 

243 def expandDataId( 

244 self, 

245 dataId: DataId | None = None, 

246 *, 

247 dimensions: Iterable[str] | DimensionGroup | DimensionGraph | None = None, 

248 graph: DimensionGraph | None = None, 

249 records: NameLookupMapping[DimensionElement, DimensionRecord | None] | None = None, 

250 withDefaults: bool = True, 

251 **kwargs: Any, 

252 ) -> DataCoordinate: 

253 # Docstring inherited from a base class. 

254 return self._registry.expandDataId( 

255 dataId, dimensions=dimensions, graph=graph, records=records, withDefaults=withDefaults, **kwargs 

256 ) 

257 

258 def insertDimensionData( 

259 self, 

260 element: DimensionElement | str, 

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

262 conform: bool = True, 

263 replace: bool = False, 

264 skip_existing: bool = False, 

265 ) -> None: 

266 # Docstring inherited from a base class. 

267 self._registry.insertDimensionData( 

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

269 ) 

270 

271 def syncDimensionData( 

272 self, 

273 element: DimensionElement | str, 

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

275 conform: bool = True, 

276 update: bool = False, 

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

278 # Docstring inherited from a base class. 

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

280 

281 def queryDatasetTypes( 

282 self, 

283 expression: Any = ..., 

284 *, 

285 components: bool | _Marker = _DefaultMarker, 

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

287 ) -> Iterable[DatasetType]: 

288 # Docstring inherited from a base class. 

289 return self._registry.queryDatasetTypes(expression, components=components, missing=missing) 

290 

291 def queryCollections( 

292 self, 

293 expression: Any = ..., 

294 datasetType: DatasetType | None = None, 

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

296 flattenChains: bool = False, 

297 includeChains: bool | None = None, 

298 ) -> Sequence[str]: 

299 # Docstring inherited from a base class. 

300 return self._registry.queryCollections( 

301 expression, datasetType, collectionTypes, flattenChains, includeChains 

302 ) 

303 

304 def queryDatasets( 

305 self, 

306 datasetType: Any, 

307 *, 

308 collections: CollectionArgType | None = None, 

309 dimensions: Iterable[Dimension | str] | None = None, 

310 dataId: DataId | None = None, 

311 where: str = "", 

312 findFirst: bool = False, 

313 components: bool | _Marker = _DefaultMarker, 

314 bind: Mapping[str, Any] | None = None, 

315 check: bool = True, 

316 **kwargs: Any, 

317 ) -> DatasetQueryResults: 

318 # Docstring inherited from a base class. 

319 return self._registry.queryDatasets( 

320 datasetType, 

321 collections=collections, 

322 dimensions=dimensions, 

323 dataId=dataId, 

324 where=where, 

325 findFirst=findFirst, 

326 components=components, 

327 bind=bind, 

328 check=check, 

329 **kwargs, 

330 ) 

331 

332 def queryDataIds( 

333 self, 

334 # TODO: Drop Dimension support on DM-41326. 

335 dimensions: DimensionGroup | Iterable[Dimension | str] | Dimension | str, 

336 *, 

337 dataId: DataId | None = None, 

338 datasets: Any = None, 

339 collections: CollectionArgType | None = None, 

340 where: str = "", 

341 components: bool | _Marker = _DefaultMarker, 

342 bind: Mapping[str, Any] | None = None, 

343 check: bool = True, 

344 **kwargs: Any, 

345 ) -> DataCoordinateQueryResults: 

346 # Docstring inherited from a base class. 

347 return self._registry.queryDataIds( 

348 dimensions, 

349 dataId=dataId, 

350 datasets=datasets, 

351 collections=collections, 

352 where=where, 

353 components=components, 

354 bind=bind, 

355 check=check, 

356 **kwargs, 

357 ) 

358 

359 def queryDimensionRecords( 

360 self, 

361 element: DimensionElement | str, 

362 *, 

363 dataId: DataId | None = None, 

364 datasets: Any = None, 

365 collections: CollectionArgType | None = None, 

366 where: str = "", 

367 components: bool | _Marker = _DefaultMarker, 

368 bind: Mapping[str, Any] | None = None, 

369 check: bool = True, 

370 **kwargs: Any, 

371 ) -> DimensionRecordQueryResults: 

372 # Docstring inherited from a base class. 

373 return self._registry.queryDimensionRecords( 

374 element, 

375 dataId=dataId, 

376 datasets=datasets, 

377 collections=collections, 

378 where=where, 

379 components=components, 

380 bind=bind, 

381 check=check, 

382 **kwargs, 

383 ) 

384 

385 def queryDatasetAssociations( 

386 self, 

387 datasetType: str | DatasetType, 

388 collections: CollectionArgType | None = ..., 

389 *, 

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

391 flattenChains: bool = False, 

392 ) -> Iterator[DatasetAssociation]: 

393 # Docstring inherited from a base class. 

394 return self._registry.queryDatasetAssociations( 

395 datasetType, 

396 collections, 

397 collectionTypes=collectionTypes, 

398 flattenChains=flattenChains, 

399 ) 

400 

401 @property 

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

403 # Docstring inherited from a base class. 

404 return self._registry.obsCoreTableManager