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

110 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-12-06 10:53 +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 

56 

57if TYPE_CHECKING: 

58 from .direct_butler import DirectButler 

59 from .registry._registry import CollectionArgType 

60 from .registry.interfaces import ObsCoreTableManager 

61 

62 

63class RegistryShim(Registry): 

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

65 

66 Parameters 

67 ---------- 

68 butler : `DirectButler` 

69 Data butler instance. 

70 

71 Notes 

72 ----- 

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

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

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

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

77 """ 

78 

79 def __init__(self, butler: DirectButler): 

80 self._butler = butler 

81 self._registry = butler._registry 

82 

83 def isWriteable(self) -> bool: 

84 # Docstring inherited from a base class. 

85 return self._registry.isWriteable() 

86 

87 @property 

88 def dimensions(self) -> DimensionUniverse: 

89 # Docstring inherited from a base class. 

90 return self._registry.dimensions 

91 

92 @property 

93 def defaults(self) -> RegistryDefaults: 

94 # Docstring inherited from a base class. 

95 return self._registry.defaults 

96 

97 @defaults.setter 

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

99 # Docstring inherited from a base class. 

100 self._registry.defaults = value 

101 

102 def refresh(self) -> None: 

103 # Docstring inherited from a base class. 

104 self._registry.refresh() 

105 

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

107 # Docstring inherited from a base class. 

108 return self._butler._caching_context() 

109 

110 @contextlib.contextmanager 

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

112 # Docstring inherited from a base class. 

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

114 yield 

115 

116 def resetConnectionPool(self) -> None: 

117 # Docstring inherited from a base class. 

118 self._registry.resetConnectionPool() 

119 

120 def registerCollection( 

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

122 ) -> bool: 

123 # Docstring inherited from a base class. 

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

125 

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

127 # Docstring inherited from a base class. 

128 return self._registry.getCollectionType(name) 

129 

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

131 # Docstring inherited from a base class. 

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

133 

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

135 # Docstring inherited from a base class. 

136 self._registry.removeCollection(name) 

137 

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

139 # Docstring inherited from a base class. 

140 return self._registry.getCollectionChain(parent) 

141 

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

143 # Docstring inherited from a base class. 

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

145 

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

147 # Docstring inherited from a base class. 

148 return self._registry.getCollectionParentChains(collection) 

149 

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

151 # Docstring inherited from a base class. 

152 return self._registry.getCollectionDocumentation(collection) 

153 

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

155 # Docstring inherited from a base class. 

156 self._registry.setCollectionDocumentation(collection, doc) 

157 

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

159 # Docstring inherited from a base class. 

160 return self._registry.getCollectionSummary(collection) 

161 

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

163 # Docstring inherited from a base class. 

164 return self._registry.registerDatasetType(datasetType) 

165 

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

167 # Docstring inherited from a base class. 

168 self._registry.removeDatasetType(name) 

169 

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

171 # Docstring inherited from a base class. 

172 return self._registry.getDatasetType(name) 

173 

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

175 # Docstring inherited from a base class. 

176 return self._registry.supportsIdGenerationMode(mode) 

177 

178 def findDataset( 

179 self, 

180 datasetType: DatasetType | str, 

181 dataId: DataId | None = None, 

182 *, 

183 collections: CollectionArgType | None = None, 

184 timespan: Timespan | None = None, 

185 **kwargs: Any, 

186 ) -> DatasetRef | None: 

187 # Docstring inherited from a base class. 

188 return self._registry.findDataset( 

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

190 ) 

191 

192 def insertDatasets( 

193 self, 

194 datasetType: DatasetType | str, 

195 dataIds: Iterable[DataId], 

196 run: str | None = None, 

197 expand: bool = True, 

198 idGenerationMode: DatasetIdGenEnum = DatasetIdGenEnum.UNIQUE, 

199 ) -> list[DatasetRef]: 

200 # Docstring inherited from a base class. 

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

202 

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

204 # Docstring inherited from a base class. 

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

206 

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

208 # Docstring inherited from a base class. 

209 return self._registry.getDataset(id) 

210 

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

212 # Docstring inherited from a base class. 

213 self._registry.removeDatasets(refs) 

214 

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

216 # Docstring inherited from a base class. 

217 self._registry.associate(collection, refs) 

218 

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

220 # Docstring inherited from a base class. 

221 self._registry.disassociate(collection, refs) 

222 

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

224 # Docstring inherited from a base class. 

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

226 

227 def decertify( 

228 self, 

229 collection: str, 

230 datasetType: str | DatasetType, 

231 timespan: Timespan, 

232 *, 

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

234 ) -> None: 

235 # Docstring inherited from a base class. 

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

237 

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

239 # Docstring inherited from a base class. 

240 return self._registry.getDatasetLocations(ref) 

241 

242 def expandDataId( 

243 self, 

244 dataId: DataId | None = None, 

245 *, 

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

247 graph: DimensionGraph | None = None, 

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

249 withDefaults: bool = True, 

250 **kwargs: Any, 

251 ) -> DataCoordinate: 

252 # Docstring inherited from a base class. 

253 return self._registry.expandDataId( 

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

255 ) 

256 

257 def insertDimensionData( 

258 self, 

259 element: DimensionElement | str, 

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

261 conform: bool = True, 

262 replace: bool = False, 

263 skip_existing: bool = False, 

264 ) -> None: 

265 # Docstring inherited from a base class. 

266 self._registry.insertDimensionData( 

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

268 ) 

269 

270 def syncDimensionData( 

271 self, 

272 element: DimensionElement | str, 

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

274 conform: bool = True, 

275 update: bool = False, 

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

277 # Docstring inherited from a base class. 

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

279 

280 def queryDatasetTypes( 

281 self, 

282 expression: Any = ..., 

283 *, 

284 components: bool | None = None, 

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

286 ) -> Iterable[DatasetType]: 

287 # Docstring inherited from a base class. 

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

289 

290 def queryCollections( 

291 self, 

292 expression: Any = ..., 

293 datasetType: DatasetType | None = None, 

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

295 flattenChains: bool = False, 

296 includeChains: bool | None = None, 

297 ) -> Sequence[str]: 

298 # Docstring inherited from a base class. 

299 return self._registry.queryCollections( 

300 expression, datasetType, collectionTypes, flattenChains, includeChains 

301 ) 

302 

303 def queryDatasets( 

304 self, 

305 datasetType: Any, 

306 *, 

307 collections: CollectionArgType | None = None, 

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

309 dataId: DataId | None = None, 

310 where: str = "", 

311 findFirst: bool = False, 

312 components: bool | None = None, 

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

314 check: bool = True, 

315 **kwargs: Any, 

316 ) -> DatasetQueryResults: 

317 # Docstring inherited from a base class. 

318 return self._registry.queryDatasets( 

319 datasetType, 

320 collections=collections, 

321 dimensions=dimensions, 

322 dataId=dataId, 

323 where=where, 

324 findFirst=findFirst, 

325 components=components, 

326 bind=bind, 

327 check=check, 

328 **kwargs, 

329 ) 

330 

331 def queryDataIds( 

332 self, 

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

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

335 *, 

336 dataId: DataId | None = None, 

337 datasets: Any = None, 

338 collections: CollectionArgType | None = None, 

339 where: str = "", 

340 components: bool | None = None, 

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

342 check: bool = True, 

343 **kwargs: Any, 

344 ) -> DataCoordinateQueryResults: 

345 # Docstring inherited from a base class. 

346 return self._registry.queryDataIds( 

347 dimensions, 

348 dataId=dataId, 

349 datasets=datasets, 

350 collections=collections, 

351 where=where, 

352 components=components, 

353 bind=bind, 

354 check=check, 

355 **kwargs, 

356 ) 

357 

358 def queryDimensionRecords( 

359 self, 

360 element: DimensionElement | str, 

361 *, 

362 dataId: DataId | None = None, 

363 datasets: Any = None, 

364 collections: CollectionArgType | None = None, 

365 where: str = "", 

366 components: bool | None = None, 

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

368 check: bool = True, 

369 **kwargs: Any, 

370 ) -> DimensionRecordQueryResults: 

371 # Docstring inherited from a base class. 

372 return self._registry.queryDimensionRecords( 

373 element, 

374 dataId=dataId, 

375 datasets=datasets, 

376 collections=collections, 

377 where=where, 

378 components=components, 

379 bind=bind, 

380 check=check, 

381 **kwargs, 

382 ) 

383 

384 def queryDatasetAssociations( 

385 self, 

386 datasetType: str | DatasetType, 

387 collections: CollectionArgType | None = ..., 

388 *, 

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

390 flattenChains: bool = False, 

391 ) -> Iterator[DatasetAssociation]: 

392 # Docstring inherited from a base class. 

393 return self._registry.queryDatasetAssociations( 

394 datasetType, 

395 collections, 

396 collectionTypes=collectionTypes, 

397 flattenChains=flattenChains, 

398 ) 

399 

400 @property 

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

402 # Docstring inherited from a base class. 

403 return self._registry.obsCoreTableManager