Coverage for python / lsst / pipe / base / tests / mocks / _storage_class.py: 42%

183 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-13 01:30 -0800

1# This file is part of pipe_base. 

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 "ConvertedUnmockedDataset", 

32 "MockDataset", 

33 "MockDatasetQuantum", 

34 "MockStorageClass", 

35 "MockStorageClassDelegate", 

36 "get_mock_name", 

37 "get_original_name", 

38 "is_mock_name", 

39) 

40 

41import sys 

42import uuid 

43from collections.abc import Callable, Iterable, Mapping 

44from typing import Any, cast 

45 

46import pydantic 

47 

48from lsst.daf.butler import ( 

49 DataIdValue, 

50 DatasetComponent, 

51 DatasetRef, 

52 DatasetType, 

53 Formatter, 

54 FormatterFactory, 

55 FormatterV2, 

56 LookupKey, 

57 SerializedDatasetType, 

58 StorageClass, 

59 StorageClassDelegate, 

60 StorageClassFactory, 

61) 

62from lsst.daf.butler.formatters.json import JsonFormatter 

63from lsst.utils.introspection import get_full_type_name 

64 

65_NAME_PREFIX: str = "_mock_" 

66 

67 

68def get_mock_name(original: str) -> str: 

69 """Return the name of the mock storage class, dataset type, or task label 

70 for the given original name. 

71 

72 Parameters 

73 ---------- 

74 original : `str` 

75 Original name. 

76 

77 Returns 

78 ------- 

79 name : `str` 

80 The name of the mocked version. 

81 """ 

82 return _NAME_PREFIX + original 

83 

84 

85def get_original_name(mock: str) -> str: 

86 """Return the name of the original storage class, dataset type, or task 

87 label that corresponds to the given mock name. 

88 

89 Parameters 

90 ---------- 

91 mock : `str` 

92 The mocked name. 

93 

94 Returns 

95 ------- 

96 original : `str` 

97 The original name. 

98 """ 

99 assert mock.startswith(_NAME_PREFIX) 

100 return mock.removeprefix(_NAME_PREFIX) 

101 

102 

103def is_mock_name(name: str | None) -> bool: 

104 """Return whether the given name is that of a mock storage class, dataset 

105 type, or task label. 

106 

107 Parameters 

108 ---------- 

109 name : `str` or `None` 

110 The given name to check. 

111 

112 Returns 

113 ------- 

114 is_mock : `bool` 

115 Whether the name is for a mock or not. 

116 """ 

117 return name is not None and name.startswith(_NAME_PREFIX) 

118 

119 

120# Tests for this module are in the ci_middleware package, where we have easy 

121# access to complex real storage classes (and their pytypes) to test against. 

122 

123 

124class MockDataset(pydantic.BaseModel): 

125 """The in-memory dataset type used by `MockStorageClass`.""" 

126 

127 dataset_id: uuid.UUID | None 

128 """Universal unique identifier for this dataset.""" 

129 

130 dataset_type: SerializedDatasetType 

131 """Butler dataset type or this dataset. 

132 

133 See the documentation for ``data_id`` for why this is a 

134 `~lsst.daf.butler.SerializedDatasetType` instead of a "real" one. 

135 """ 

136 

137 data_id: dict[str, DataIdValue] 

138 """Butler data ID for this dataset. 

139 

140 This is a `~lsst.daf.butler.SerializedDataCoordinate` instead of a "real" 

141 one for two reasons: 

142 

143 - the mock dataset may need to be read from disk in a context in which a 

144 `~lsst.daf.butler.DimensionUniverse` is unavailable; 

145 - we don't want the complexity of having a separate 

146 ``SerializedMockDataCoordinate``. 

147 """ 

148 

149 run: str | None 

150 """`~lsst.daf.butler.CollectionType.RUN` collection this dataset belongs 

151 to. 

152 """ 

153 

154 quantum: MockDatasetQuantum | None = None 

155 """Description of the quantum that produced this dataset. 

156 """ 

157 

158 output_connection_name: str | None = None 

159 """The name of the PipelineTask output connection that produced this 

160 dataset. 

161 """ 

162 

163 converted_from: MockDataset | None = None 

164 """Another `MockDataset` that underwent a storage class conversion to 

165 produce this one. 

166 """ 

167 

168 parent: MockDataset | None = None 

169 """Another `MockDataset` from which a component was extract to form this 

170 one. 

171 """ 

172 

173 parameters: dict[str, str] | None = None 

174 """`repr` of all parameters applied when reading this dataset.""" 

175 

176 int_value: int | None = None 

177 """An arbitrary integer value stored in the mock dataset.""" 

178 

179 str_value: int | None = None 

180 """An arbitrary string value stored in the mock dataset.""" 

181 

182 @property 

183 def storage_class(self) -> str: 

184 return cast(str, self.dataset_type.storageClass) 

185 

186 def make_derived(self, **kwargs: Any) -> MockDataset: 

187 """Return a new MockDataset that represents applying some storage class 

188 operation to this one. 

189 

190 Parameters 

191 ---------- 

192 **kwargs : `~typing.Any` 

193 Keyword arguments are fields of `MockDataset` or 

194 `~lsst.daf.butler.SerializedDatasetType` to override in the result. 

195 

196 Returns 

197 ------- 

198 derived : `MockDataset` 

199 The newly-mocked dataset. 

200 """ 

201 dataset_type_updates = { 

202 k: kwargs.pop(k) for k in list(kwargs) if k in SerializedDatasetType.model_fields 

203 } 

204 kwargs.setdefault("dataset_type", self.dataset_type.model_copy(update=dataset_type_updates)) 

205 # Fields below are those that should not be propagated to the derived 

206 # dataset, because they're not about the intrinsic on-disk thing. 

207 kwargs.setdefault("converted_from", None) 

208 kwargs.setdefault("parent", None) 

209 kwargs.setdefault("parameters", None) 

210 # Also use setdefault on the ref in case caller wants to override that 

211 # directly, but this is expected to be rare enough that it's not worth 

212 # it to try to optimize out the work above to make derived_ref. 

213 return self.model_copy(update=kwargs) 

214 

215 # Work around the fact that Sphinx chokes on Pydantic docstring formatting, 

216 # when we inherit those docstrings in our public classes. 

217 if "sphinx" in sys.modules: 

218 

219 def copy(self, *args: Any, **kwargs: Any) -> Any: 

220 """See `pydantic.BaseModel.copy`.""" 

221 return super().copy(*args, **kwargs) 

222 

223 def model_dump(self, *args: Any, **kwargs: Any) -> Any: 

224 """See `pydantic.BaseModel.model_dump`.""" 

225 return super().model_dump(*args, **kwargs) 

226 

227 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any: 

228 """See `pydantic.BaseModel.model_dump_json`.""" 

229 return super().model_dump(*args, **kwargs) 

230 

231 def model_copy(self, *args: Any, **kwargs: Any) -> Any: 

232 """See `pydantic.BaseModel.model_copy`.""" 

233 return super().model_copy(*args, **kwargs) 

234 

235 @classmethod 

236 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[override] 

237 """See `pydantic.BaseModel.model_construct`.""" 

238 return super().model_construct(*args, **kwargs) 

239 

240 @classmethod 

241 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any: 

242 """See `pydantic.BaseModel.model_json_schema`.""" 

243 return super().model_json_schema(*args, **kwargs) 

244 

245 @classmethod 

246 def model_validate(cls, *args: Any, **kwargs: Any) -> Any: 

247 """See `pydantic.BaseModel.model_validate`.""" 

248 return super().model_validate(*args, **kwargs) 

249 

250 @classmethod 

251 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any: 

252 """See `pydantic.BaseModel.model_validate_json`.""" 

253 return super().model_validate_json(*args, **kwargs) 

254 

255 @classmethod 

256 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any: 

257 """See `pydantic.BaseModel.model_validate_strings`.""" 

258 return super().model_validate_strings(*args, **kwargs) 

259 

260 

261class ConvertedUnmockedDataset(pydantic.BaseModel): 

262 """A marker class that represents a conversion from a regular in-memory 

263 dataset to a mock storage class. 

264 """ 

265 

266 original_type: str 

267 """The full Python type of the original unmocked in-memory dataset.""" 

268 

269 # Work around the fact that Sphinx chokes on Pydantic docstring formatting, 

270 # when we inherit those docstrings in our public classes. 

271 if "sphinx" in sys.modules: 

272 

273 def copy(self, *args: Any, **kwargs: Any) -> Any: 

274 """See `pydantic.BaseModel.copy`.""" 

275 return super().copy(*args, **kwargs) 

276 

277 def model_dump(self, *args: Any, **kwargs: Any) -> Any: 

278 """See `pydantic.BaseModel.model_dump`.""" 

279 return super().model_dump(*args, **kwargs) 

280 

281 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any: 

282 """See `pydantic.BaseModel.model_dump_json`.""" 

283 return super().model_dump(*args, **kwargs) 

284 

285 def model_copy(self, *args: Any, **kwargs: Any) -> Any: 

286 """See `pydantic.BaseModel.model_copy`.""" 

287 return super().model_copy(*args, **kwargs) 

288 

289 @classmethod 

290 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override] 

291 """See `pydantic.BaseModel.model_construct`.""" 

292 return super().model_construct(*args, **kwargs) 

293 

294 @classmethod 

295 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any: 

296 """See `pydantic.BaseModel.model_json_schema`.""" 

297 return super().model_json_schema(*args, **kwargs) 

298 

299 @classmethod 

300 def model_validate(cls, *args: Any, **kwargs: Any) -> Any: 

301 """See `pydantic.BaseModel.model_validate`.""" 

302 return super().model_validate(*args, **kwargs) 

303 

304 @classmethod 

305 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any: 

306 """See `pydantic.BaseModel.model_validate_json`.""" 

307 return super().model_validate_json(*args, **kwargs) 

308 

309 @classmethod 

310 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any: 

311 """See `pydantic.BaseModel.model_validate_strings`.""" 

312 return super().model_validate_strings(*args, **kwargs) 

313 

314 

315class MockDatasetQuantum(pydantic.BaseModel): 

316 """Description of the quantum that produced a mock dataset. 

317 

318 This is also used to represent task-init operations for init-output mock 

319 datasets. 

320 """ 

321 

322 task_label: str 

323 """Label of the producing PipelineTask in its pipeline.""" 

324 

325 data_id: dict[str, DataIdValue] 

326 """Data ID for the quantum.""" 

327 

328 inputs: dict[str, list[MockDataset | ConvertedUnmockedDataset]] 

329 """Mock datasets provided as input to the quantum. 

330 

331 Keys are task-internal connection names, not dataset type names. 

332 """ 

333 

334 # Work around the fact that Sphinx chokes on Pydantic docstring formatting, 

335 # when we inherit those docstrings in our public classes. 

336 if "sphinx" in sys.modules: 

337 

338 def copy(self, *args: Any, **kwargs: Any) -> Any: 

339 """See `pydantic.BaseModel.copy`.""" 

340 return super().copy(*args, **kwargs) 

341 

342 def model_dump(self, *args: Any, **kwargs: Any) -> Any: 

343 """See `pydantic.BaseModel.model_dump`.""" 

344 return super().model_dump(*args, **kwargs) 

345 

346 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any: 

347 """See `pydantic.BaseModel.model_dump_json`.""" 

348 return super().model_dump(*args, **kwargs) 

349 

350 def model_copy(self, *args: Any, **kwargs: Any) -> Any: 

351 """See `pydantic.BaseModel.model_copy`.""" 

352 return super().model_copy(*args, **kwargs) 

353 

354 @classmethod 

355 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override] 

356 """See `pydantic.BaseModel.model_construct`.""" 

357 return super().model_construct(*args, **kwargs) 

358 

359 @classmethod 

360 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any: 

361 """See `pydantic.BaseModel.model_json_schema`.""" 

362 return super().model_json_schema(*args, **kwargs) 

363 

364 @classmethod 

365 def model_validate(cls, *args: Any, **kwargs: Any) -> Any: 

366 """See `pydantic.BaseModel.model_validate`.""" 

367 return super().model_validate(*args, **kwargs) 

368 

369 @classmethod 

370 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any: 

371 """See `pydantic.BaseModel.model_validate_json`.""" 

372 return super().model_validate_json(*args, **kwargs) 

373 

374 @classmethod 

375 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any: 

376 """See `pydantic.BaseModel.model_validate_strings`.""" 

377 return super().model_validate_strings(*args, **kwargs) 

378 

379 

380MockDataset.model_rebuild() 

381 

382 

383class MockStorageClassDelegate(StorageClassDelegate): 

384 """Implementation of the StorageClassDelegate interface for mock datasets. 

385 

386 This class does not implement assembly and disassembly just because it's 

387 not needed right now. That could be added in the future with some 

388 additional tracking attributes in `MockDataset`. 

389 """ 

390 

391 def assemble(self, components: dict[str, Any], pytype: type | None = None) -> MockDataset: 

392 # Docstring inherited. 

393 raise NotImplementedError("Mock storage classes do not implement assembly.") 

394 

395 def getComponent(self, composite: Any, componentName: str) -> Any: 

396 # Docstring inherited. 

397 assert isinstance(composite, MockDataset), ( 

398 f"MockStorageClassDelegate given a non-mock dataset {composite!r}." 

399 ) 

400 return composite.make_derived( 

401 name=f"{composite.dataset_type.name}.{componentName}", 

402 storageClass=self.storageClass.allComponents()[componentName].name, 

403 parentStorageClass=self.storageClass.name, 

404 parent=composite, 

405 ) 

406 

407 def disassemble( 

408 self, composite: Any, subset: Iterable | None = None, override: Any | None = None 

409 ) -> dict[str, DatasetComponent]: 

410 # Docstring inherited. 

411 raise NotImplementedError("Mock storage classes do not implement disassembly.") 

412 

413 def handleParameters(self, inMemoryDataset: Any, parameters: Mapping[str, Any] | None = None) -> Any: 

414 # Docstring inherited. 

415 assert isinstance(inMemoryDataset, MockDataset), ( 

416 f"MockStorageClassDelegate given a non-mock dataset {inMemoryDataset!r}." 

417 ) 

418 if not parameters: 

419 return inMemoryDataset 

420 return inMemoryDataset.make_derived(parameters={k: repr(v) for k, v in parameters.items()}) 

421 

422 

423class MockStorageClass(StorageClass): 

424 """A reimplementation of `lsst.daf.butler.StorageClass` for mock datasets. 

425 

426 Parameters 

427 ---------- 

428 original : `~lsst.daf.butler.StorageClass` 

429 The original storage class. 

430 factory : `~lsst.daf.butler.StorageClassFactory` or `None`, optional 

431 Storage class factory to use. If `None` the default factory is used. 

432 

433 Notes 

434 ----- 

435 Each `MockStorageClass` instance corresponds to a real "original" storage 

436 class, with components and conversions that are mocks of the original's 

437 components and conversions. The ``pytype`` for all `MockStorageClass` 

438 instances is `MockDataset`. 

439 """ 

440 

441 def __init__(self, original: StorageClass, factory: StorageClassFactory | None = None): 

442 name = get_mock_name(original.name) 

443 if factory is None: 

444 factory = StorageClassFactory() 

445 super().__init__( 

446 name=name, 

447 pytype=MockDataset, 

448 components={ 

449 k: self.get_or_register_mock(v.name, factory) for k, v in original.components.items() 

450 }, 

451 derivedComponents={ 

452 k: self.get_or_register_mock(v.name, factory) for k, v in original.derivedComponents.items() 

453 }, 

454 parameters=frozenset(original.parameters), 

455 delegate=get_full_type_name(MockStorageClassDelegate), 

456 # Conversions work differently for mock storage classes, since they 

457 # all have the same pytype: we use the original storage class being 

458 # mocked to see if we can convert, then just make a new MockDataset 

459 # that points back to the original. 

460 converters={}, 

461 ) 

462 self.original = original 

463 # Make certain no one tries to use the converters. 

464 self._converters = None # type: ignore 

465 

466 def _get_converters_by_type(self) -> dict[type, Callable[[Any], Any]]: 

467 # Docstring inherited. 

468 raise NotImplementedError("MockStorageClass does not use converters.") 

469 

470 @classmethod 

471 def get_or_register_mock( 

472 cls, original: str, factory: StorageClassFactory | None = None 

473 ) -> MockStorageClass: 

474 """Return a mock storage class for the given original storage class, 

475 creating and registering it if necessary. 

476 

477 Parameters 

478 ---------- 

479 original : `str` 

480 Name of the original storage class to be mocked. 

481 factory : `~lsst.daf.butler.StorageClassFactory`, optional 

482 Storage class factory singleton instance. 

483 

484 Returns 

485 ------- 

486 mock : `MockStorageClass` 

487 New storage class that mocks ``original``. 

488 """ 

489 name = get_mock_name(original) 

490 if factory is None: 

491 factory = StorageClassFactory() 

492 if name in factory: 

493 return cast(MockStorageClass, factory.getStorageClass(name)) 

494 else: 

495 result = cls(factory.getStorageClass(original), factory) 

496 factory.registerStorageClass(result) 

497 return result 

498 

499 def allComponents(self) -> Mapping[str, MockStorageClass]: 

500 # Docstring inherited. 

501 return cast(Mapping[str, MockStorageClass], super().allComponents()) 

502 

503 @property 

504 def components(self) -> Mapping[str, MockStorageClass]: 

505 # Docstring inherited. 

506 return cast(Mapping[str, MockStorageClass], super().components) 

507 

508 @property 

509 def derivedComponents(self) -> Mapping[str, MockStorageClass]: 

510 # Docstring inherited. 

511 return cast(Mapping[str, MockStorageClass], super().derivedComponents) 

512 

513 def can_convert(self, other: StorageClass) -> bool: 

514 # Docstring inherited. 

515 if not isinstance(other, MockStorageClass): 

516 # Allow conversions from an original type (and others compatible 

517 # with it) to a mock, to allow for cases where an upstream task 

518 # did not use a mock to write something but the downstream one is 

519 # trying to us a mock to read it. 

520 return self.original.can_convert(other) 

521 return self.original.can_convert(other.original) 

522 

523 def coerce_type(self, incorrect: Any) -> Any: 

524 # Docstring inherited. 

525 if not isinstance(incorrect, MockDataset): 

526 if isinstance(incorrect, ConvertedUnmockedDataset): 

527 return incorrect 

528 return ConvertedUnmockedDataset(original_type=get_full_type_name(incorrect)) 

529 factory = StorageClassFactory() 

530 other_storage_class = factory.getStorageClass(incorrect.storage_class) 

531 assert isinstance(other_storage_class, MockStorageClass), "Should not get a MockDataset otherwise." 

532 if other_storage_class.name == self.name: 

533 return incorrect 

534 if not self.can_convert(other_storage_class): 

535 raise TypeError( 

536 f"Mocked storage class {self.original.name!r} cannot convert from " 

537 f"{other_storage_class.original.name!r}." 

538 ) 

539 return incorrect.make_derived(storageClass=self.name, converted_from=incorrect) 

540 

541 @staticmethod 

542 def mock_dataset_type(original_type: DatasetType) -> DatasetType: 

543 """Replace a dataset type with a version that uses a mock storage class 

544 and name. 

545 

546 Parameters 

547 ---------- 

548 original_type : `lsst.daf.butler.DatasetType` 

549 Original dataset type to be mocked. 

550 

551 Returns 

552 ------- 

553 mock_type : `lsst.daf.butler.DatasetType` 

554 A mock version of the dataset type, with name and storage class 

555 changed and everything else unchanged. 

556 """ 

557 mock_storage_class = MockStorageClass.get_or_register_mock(original_type.storageClass_name) 

558 mock_parent_storage_class = None 

559 if original_type.parentStorageClass is not None: 

560 mock_parent_storage_class = MockStorageClass.get_or_register_mock( 

561 original_type.parentStorageClass.name 

562 ) 

563 return DatasetType( 

564 get_mock_name(original_type.name), 

565 original_type.dimensions, 

566 mock_storage_class, 

567 isCalibration=original_type.isCalibration(), 

568 parentStorageClass=mock_parent_storage_class, 

569 ) 

570 

571 @staticmethod 

572 def mock_dataset_refs(original_refs: Iterable[DatasetRef]) -> list[DatasetRef]: 

573 """Replace dataset references with versions that uses a mock storage 

574 class and dataset type name. 

575 

576 Parameters 

577 ---------- 

578 original_refs : `~collections.abc.Iterable` [ \ 

579 `lsst.daf.butler.DatasetRef` ] 

580 Original dataset references to be mocked. 

581 

582 Returns 

583 ------- 

584 mock_refs : `list` [ `lsst.daf.butler.DatasetRef` ] 

585 Mocked version of the dataset references, with dataset type name 

586 and storage class changed and everything else unchanged. 

587 """ 

588 original_refs = list(original_refs) 

589 if not original_refs: 

590 return original_refs 

591 dataset_type = MockStorageClass.mock_dataset_type(original_refs[0].datasetType) 

592 return [ 

593 DatasetRef(dataset_type, original_ref.dataId, run=original_ref.run, id=original_ref.id) 

594 for original_ref in original_refs 

595 ] 

596 

597 @staticmethod 

598 def unmock_dataset_type(mock_type: DatasetType) -> DatasetType: 

599 """Replace a mock dataset type with the original one it was created 

600 from. 

601 

602 Parameters 

603 ---------- 

604 mock_type : `lsst.daf.butler.DatasetType` 

605 A dataset type with a mocked name and storage class. 

606 

607 Returns 

608 ------- 

609 original_type : `lsst.daf.butler.DatasetType` 

610 The original dataset type. 

611 """ 

612 storage_class = mock_type.storageClass 

613 parent_storage_class = mock_type.parentStorageClass 

614 if isinstance(storage_class, MockStorageClass): 

615 storage_class = storage_class.original 

616 if parent_storage_class is not None and isinstance(parent_storage_class, MockStorageClass): 

617 parent_storage_class = parent_storage_class.original 

618 return DatasetType( 

619 get_original_name(mock_type.name), 

620 mock_type.dimensions, 

621 storage_class, 

622 isCalibration=mock_type.isCalibration(), 

623 parentStorageClass=parent_storage_class, 

624 ) 

625 

626 @staticmethod 

627 def unmock_dataset_refs(mock_refs: Iterable[DatasetRef]) -> list[DatasetRef]: 

628 """Replace dataset references with versions that do not use a mock 

629 storage class and dataset type name. 

630 

631 Parameters 

632 ---------- 

633 mock_refs : `~collections.abc.Iterable` [ \ 

634 `lsst.daf.butler.DatasetRef` ] 

635 Dataset references that use a mocked dataset type name and storage 

636 class. 

637 

638 Returns 

639 ------- 

640 original_refs : `list` [ `lsst.daf.butler.DatasetRef` ] 

641 The original dataset references. 

642 """ 

643 mock_refs = list(mock_refs) 

644 if not mock_refs: 

645 return mock_refs 

646 dataset_type = MockStorageClass.unmock_dataset_type(mock_refs[0].datasetType) 

647 return [ 

648 DatasetRef(dataset_type, mock_ref.dataId, run=mock_ref.run, id=mock_ref.id) 

649 for mock_ref in mock_refs 

650 ] 

651 

652 

653def _monkeypatch_daf_butler() -> None: 

654 """Replace methods in daf_butler's StorageClassFactory and FormatterFactory 

655 classes to automatically recognize mock storage classes. 

656 

657 This monkey-patching is executed when the `lsst.pipe.base.tests.mocks` 

658 package is imported, and it affects all butler instances created before or 

659 after that imported. 

660 """ 

661 original_get_storage_class = StorageClassFactory.getStorageClass 

662 

663 def new_get_storage_class(self: StorageClassFactory, storageClassName: str) -> StorageClass: 

664 try: 

665 return original_get_storage_class(self, storageClassName) 

666 except KeyError: 

667 if is_mock_name(storageClassName): 

668 return MockStorageClass.get_or_register_mock(get_original_name(storageClassName)) 

669 raise 

670 

671 StorageClassFactory.getStorageClass = new_get_storage_class # type: ignore 

672 

673 del new_get_storage_class 

674 

675 original_get_formatter_class_with_match = FormatterFactory.getFormatterClassWithMatch 

676 

677 def new_get_formatter_class_with_match( 

678 self: FormatterFactory, entity: Any 

679 ) -> tuple[LookupKey, type[Formatter | FormatterV2], dict[str, Any]]: 

680 try: 

681 return original_get_formatter_class_with_match(self, entity) 

682 except KeyError: 

683 lookup_keys = (LookupKey(name=entity),) if isinstance(entity, str) else entity._lookupNames() 

684 for key in lookup_keys: 

685 # This matches mock dataset type names before mock storage 

686 # classes, and it would even match some regular dataset types 

687 # that are automatic connections (logs, configs, metadata) of 

688 # mocked tasks. The latter would be a problem, except that 

689 # those should have already matched in the try block above. 

690 if is_mock_name(key.name): 

691 return (key, JsonFormatter, {}) 

692 raise 

693 

694 FormatterFactory.getFormatterClassWithMatch = new_get_formatter_class_with_match # type: ignore 

695 

696 del new_get_formatter_class_with_match 

697 

698 original_get_formatter_with_match = FormatterFactory.getFormatterWithMatch 

699 

700 def new_get_formatter_with_match( 

701 self: FormatterFactory, entity: Any, *args: Any, **kwargs: Any 

702 ) -> tuple[LookupKey, Formatter | FormatterV2]: 

703 try: 

704 return original_get_formatter_with_match(self, entity, *args, **kwargs) 

705 except KeyError: 

706 lookup_keys = (LookupKey(name=entity),) if isinstance(entity, str) else entity._lookupNames() 

707 for key in lookup_keys: 

708 if is_mock_name(key.name): 

709 return (key, JsonFormatter(*args, **kwargs)) 

710 raise 

711 

712 FormatterFactory.getFormatterWithMatch = new_get_formatter_with_match # type: ignore 

713 

714 del new_get_formatter_with_match 

715 

716 

717_monkeypatch_daf_butler()