Coverage for tests/test_simpleButler.py: 15%

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

276 statements  

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 program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24import os 

25import tempfile 

26from typing import Any 

27import unittest 

28import uuid 

29import re 

30import json 

31 

32try: 

33 import numpy as np 

34except ImportError: 

35 np = None 

36 

37import astropy.time 

38 

39from lsst.daf.butler import ( 

40 Butler, 

41 ButlerConfig, 

42 CollectionType, 

43 DatasetRef, 

44 DatasetType, 

45 Registry, 

46 Timespan, 

47) 

48from lsst.daf.butler.registry import RegistryConfig, RegistryDefaults, ConflictingDefinitionError 

49from lsst.daf.butler.tests import DatastoreMock 

50from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir 

51 

52 

53TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

54 

55 

56class SimpleButlerTestCase(unittest.TestCase): 

57 """Tests for butler (including import/export functionality) that should not 

58 depend on the Registry Database backend or Datastore implementation, and 

59 can instead utilize an in-memory SQLite Registry and a mocked Datastore. 

60 """ 

61 

62 datasetsManager = \ 

63 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManager" 

64 datasetsImportFile = "datasets.yaml" 

65 datasetsIdType = int 

66 

67 def setUp(self): 

68 self.root = makeTestTempDir(TESTDIR) 

69 

70 def tearDown(self): 

71 removeTestTempDir(self.root) 

72 

73 def makeButler(self, **kwargs: Any) -> Butler: 

74 """Return new Butler instance on each call. 

75 """ 

76 config = ButlerConfig() 

77 

78 # make separate temporary directory for registry of this instance 

79 tmpdir = tempfile.mkdtemp(dir=self.root) 

80 config["registry", "db"] = f"sqlite:///{tmpdir}/gen3.sqlite3" 

81 config["registry", "managers", "datasets"] = self.datasetsManager 

82 config["root"] = self.root 

83 

84 # have to make a registry first 

85 registryConfig = RegistryConfig(config.get("registry")) 

86 Registry.createFromConfig(registryConfig) 

87 

88 butler = Butler(config, **kwargs) 

89 DatastoreMock.apply(butler) 

90 return butler 

91 

92 def comparableRef(self, ref: DatasetRef) -> DatasetRef: 

93 """Return a DatasetRef that can be compared to a DatasetRef from 

94 other repository. 

95 

96 For repositories that do not support round-trip of ID values this 

97 method returns unresolved DatasetRef, for round-trip-safe repos it 

98 returns unchanged ref. 

99 """ 

100 return ref if self.datasetsIdType is uuid.UUID else ref.unresolved() 

101 

102 def testReadBackwardsCompatibility(self): 

103 """Test that we can read an export file written by a previous version 

104 and commit to the daf_butler git repo. 

105 

106 Notes 

107 ----- 

108 At present this export file includes only dimension data, not datasets, 

109 which greatly limits the usefulness of this test. We should address 

110 this at some point, but I think it's best to wait for the changes to 

111 the export format required for CALIBRATION collections to land. 

112 """ 

113 butler = self.makeButler(writeable=True) 

114 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

115 # Spot-check a few things, but the most important test is just that 

116 # the above does not raise. 

117 self.assertGreaterEqual( 

118 set(record.id for record in butler.registry.queryDimensionRecords("detector", instrument="HSC")), 

119 set(range(104)), # should have all science CCDs; may have some focus ones. 

120 ) 

121 self.assertGreaterEqual( 

122 { 

123 (record.id, record.physical_filter) 

124 for record in butler.registry.queryDimensionRecords("visit", instrument="HSC") 

125 }, 

126 { 

127 (27136, 'HSC-Z'), 

128 (11694, 'HSC-G'), 

129 (23910, 'HSC-R'), 

130 (11720, 'HSC-Y'), 

131 (23900, 'HSC-R'), 

132 (22646, 'HSC-Y'), 

133 (1248, 'HSC-I'), 

134 (19680, 'HSC-I'), 

135 (1240, 'HSC-I'), 

136 (424, 'HSC-Y'), 

137 (19658, 'HSC-I'), 

138 (344, 'HSC-Y'), 

139 (1218, 'HSC-R'), 

140 (1190, 'HSC-Z'), 

141 (23718, 'HSC-R'), 

142 (11700, 'HSC-G'), 

143 (26036, 'HSC-G'), 

144 (23872, 'HSC-R'), 

145 (1170, 'HSC-Z'), 

146 (1876, 'HSC-Y'), 

147 } 

148 ) 

149 

150 def testDatasetTransfers(self): 

151 """Test exporting all datasets from a repo and then importing them all 

152 back in again. 

153 """ 

154 # Import data to play with. 

155 butler1 = self.makeButler(writeable=True) 

156 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

157 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

158 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

159 # Export all datasets. 

160 with butler1.export(filename=file.name) as exporter: 

161 exporter.saveDatasets( 

162 butler1.registry.queryDatasets(..., collections=...) 

163 ) 

164 # Import it all again. 

165 butler2 = self.makeButler(writeable=True) 

166 butler2.import_(filename=file.name) 

167 datasets1 = list(butler1.registry.queryDatasets(..., collections=...)) 

168 datasets2 = list(butler2.registry.queryDatasets(..., collections=...)) 

169 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets1)) 

170 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets2)) 

171 self.assertCountEqual( 

172 [self.comparableRef(ref) for ref in datasets1], 

173 [self.comparableRef(ref) for ref in datasets2], 

174 ) 

175 

176 def testComponentExport(self): 

177 """Test exporting component datasets and then importing them. 

178 

179 This test intentionally does not depend on whether just the component 

180 is exported and then imported vs. the full composite dataset, because 

181 I don't want it to assume more than it needs to about the 

182 implementation. 

183 """ 

184 # Import data to play with. 

185 butler1 = self.makeButler(writeable=True) 

186 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

187 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

188 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

189 # Export all datasets. 

190 with butler1.export(filename=file.name) as exporter: 

191 exporter.saveDatasets( 

192 butler1.registry.queryDatasets("flat.psf", collections=...) 

193 ) 

194 # Import it all again. 

195 butler2 = self.makeButler(writeable=True) 

196 butler2.import_(filename=file.name) 

197 datasets1 = list(butler1.registry.queryDatasets("flat.psf", collections=...)) 

198 datasets2 = list(butler2.registry.queryDatasets("flat.psf", collections=...)) 

199 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets1)) 

200 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets2)) 

201 self.assertCountEqual( 

202 [self.comparableRef(ref) for ref in datasets1], 

203 [self.comparableRef(ref) for ref in datasets2], 

204 ) 

205 

206 def testDatasetImportTwice(self): 

207 """Test exporting all datasets from a repo and then importing them all 

208 back in again twice. 

209 """ 

210 if self.datasetsIdType is not uuid.UUID: 

211 self.skipTest("This test can only work for UUIDs") 

212 # Import data to play with. 

213 butler1 = self.makeButler(writeable=True) 

214 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

215 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

216 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml", delete=False) as file: 

217 # Export all datasets. 

218 with butler1.export(filename=file.name) as exporter: 

219 exporter.saveDatasets( 

220 butler1.registry.queryDatasets(..., collections=...) 

221 ) 

222 butler2 = self.makeButler(writeable=True) 

223 # Import it once. 

224 butler2.import_(filename=file.name) 

225 # Import it again, but ignore all dimensions 

226 dimensions = set( 

227 dimension.name for dimension in butler2.registry.dimensions.getStaticDimensions()) 

228 butler2.import_(filename=file.name, skip_dimensions=dimensions) 

229 datasets1 = list(butler1.registry.queryDatasets(..., collections=...)) 

230 datasets2 = list(butler2.registry.queryDatasets(..., collections=...)) 

231 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets1)) 

232 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets2)) 

233 self.assertCountEqual( 

234 [self.comparableRef(ref) for ref in datasets1], 

235 [self.comparableRef(ref) for ref in datasets2], 

236 ) 

237 

238 def testDatasetImportReuseIds(self): 

239 """Test for import that should preserve dataset IDs. 

240 

241 This test assumes that dataset IDs in datasets YAML are different from 

242 what auto-incremental insert would produce. 

243 """ 

244 if self.datasetsIdType is not int: 

245 self.skipTest("This test can only work for UUIDs") 

246 # Import data to play with. 

247 butler = self.makeButler(writeable=True) 

248 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

249 filename = os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile) 

250 butler.import_(filename=filename, reuseIds=True) 

251 datasets = list(butler.registry.queryDatasets(..., collections=...)) 

252 self.assertTrue(all(isinstance(ref.id, self.datasetsIdType) for ref in datasets)) 

253 # IDs are copied from YAML, list needs to be updated if file contents 

254 # is changed. 

255 self.assertCountEqual( 

256 [ref.id for ref in datasets], 

257 [1001, 1002, 1003, 1010, 1020, 1030, 2001, 2002, 2003, 2010, 2020, 2030, 2040], 

258 ) 

259 

260 # Try once again, it will raise 

261 with self.assertRaises(ConflictingDefinitionError): 

262 butler.import_(filename=filename, reuseIds=True) 

263 

264 def testCollectionTransfers(self): 

265 """Test exporting and then importing collections of various types. 

266 """ 

267 # Populate a registry with some datasets. 

268 butler1 = self.makeButler(writeable=True) 

269 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

270 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

271 registry1 = butler1.registry 

272 # Add some more collections. 

273 registry1.registerRun("run1") 

274 registry1.registerCollection("tag1", CollectionType.TAGGED) 

275 registry1.registerCollection("calibration1", CollectionType.CALIBRATION) 

276 registry1.registerCollection("chain1", CollectionType.CHAINED) 

277 registry1.registerCollection("chain2", CollectionType.CHAINED) 

278 registry1.setCollectionChain("chain1", ["tag1", "run1", "chain2"]) 

279 registry1.setCollectionChain("chain2", ["calibration1", "run1"]) 

280 # Associate some datasets into the TAGGED and CALIBRATION collections. 

281 flats1 = list(registry1.queryDatasets("flat", collections=...)) 

282 registry1.associate("tag1", flats1) 

283 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai") 

284 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai") 

285 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai") 

286 bias1a = registry1.findDataset("bias", instrument="Cam1", detector=1, collections="imported_g") 

287 bias2a = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

288 bias3a = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

289 bias2b = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

290 bias3b = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

291 registry1.certify("calibration1", [bias2a, bias3a], Timespan(t1, t2)) 

292 registry1.certify("calibration1", [bias2b], Timespan(t2, None)) 

293 registry1.certify("calibration1", [bias3b], Timespan(t2, t3)) 

294 registry1.certify("calibration1", [bias1a], Timespan.makeEmpty()) 

295 

296 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

297 # Export all collections, and some datasets. 

298 with butler1.export(filename=file.name) as exporter: 

299 # Sort results to put chain1 before chain2, which is 

300 # intentionally not topological order. 

301 for collection in sorted(registry1.queryCollections()): 

302 exporter.saveCollection(collection) 

303 exporter.saveDatasets(flats1) 

304 exporter.saveDatasets([bias1a, bias2a, bias2b, bias3a, bias3b]) 

305 # Import them into a new registry. 

306 butler2 = self.makeButler(writeable=True) 

307 butler2.import_(filename=file.name) 

308 registry2 = butler2.registry 

309 # Check that it all round-tripped, starting with the collections 

310 # themselves. 

311 self.assertIs(registry2.getCollectionType("run1"), CollectionType.RUN) 

312 self.assertIs(registry2.getCollectionType("tag1"), CollectionType.TAGGED) 

313 self.assertIs(registry2.getCollectionType("calibration1"), CollectionType.CALIBRATION) 

314 self.assertIs(registry2.getCollectionType("chain1"), CollectionType.CHAINED) 

315 self.assertIs(registry2.getCollectionType("chain2"), CollectionType.CHAINED) 

316 self.assertEqual( 

317 list(registry2.getCollectionChain("chain1")), 

318 ["tag1", "run1", "chain2"], 

319 ) 

320 self.assertEqual( 

321 list(registry2.getCollectionChain("chain2")), 

322 ["calibration1", "run1"], 

323 ) 

324 # Check that tag collection contents are the same. 

325 self.maxDiff = None 

326 self.assertCountEqual( 

327 [self.comparableRef(ref) for ref in registry1.queryDatasets(..., collections="tag1")], 

328 [self.comparableRef(ref) for ref in registry2.queryDatasets(..., collections="tag1")], 

329 ) 

330 # Check that calibration collection contents are the same. 

331 self.assertCountEqual( 

332 [(self.comparableRef(assoc.ref), assoc.timespan) 

333 for assoc in registry1.queryDatasetAssociations("bias", collections="calibration1")], 

334 [(self.comparableRef(assoc.ref), assoc.timespan) 

335 for assoc in registry2.queryDatasetAssociations("bias", collections="calibration1")], 

336 ) 

337 

338 def testButlerGet(self): 

339 """Test that butler.get can work with different variants.""" 

340 

341 # Import data to play with. 

342 butler = self.makeButler(writeable=True) 

343 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

344 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

345 

346 # Find the DatasetRef for a flat 

347 coll = "imported_g" 

348 flat2g = butler.registry.findDataset("flat", instrument="Cam1", detector=2, physical_filter="Cam1-G", 

349 collections=coll) 

350 

351 # Create a numpy integer to check that works fine 

352 detector_np = np.int64(2) if np else 2 

353 print(type(detector_np)) 

354 

355 # Try to get it using different variations of dataId + keyword 

356 # arguments 

357 # Note that instrument.class_name does not work 

358 variants = ( 

359 (None, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

360 (None, {"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}), 

361 ({"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}, {}), 

362 ({"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}, {}), 

363 ({"instrument": "Cam1", "detector": 2}, {"physical_filter": "Cam1-G"}), 

364 ({"detector.full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

365 ({"full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

366 (None, {"full_name": "Ab", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

367 (None, {"detector": "Ab", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

368 ({"name_in_raft": "b", "raft": "A"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

369 ({"name_in_raft": "b"}, {"raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

370 (None, {"name_in_raft": "b", "raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

371 ({"detector.name_in_raft": "b", "detector.raft": "A"}, 

372 {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

373 ({"detector.name_in_raft": "b", "detector.raft": "A", 

374 "instrument": "Cam1", "physical_filter": "Cam1-G"}, {}), 

375 ) 

376 

377 for dataId, kwds in variants: 

378 try: 

379 flat_id, _ = butler.get("flat", dataId=dataId, collections=coll, **kwds) 

380 except Exception as e: 

381 raise type(e)(f"{str(e)}: dataId={dataId}, kwds={kwds}") from e 

382 self.assertEqual(flat_id, flat2g.id, msg=f"DataId: {dataId}, kwds: {kwds}") 

383 

384 def testGetCalibration(self): 

385 """Test that `Butler.get` can be used to fetch from 

386 `~CollectionType.CALIBRATION` collections if the data ID includes 

387 extra dimensions with temporal information. 

388 """ 

389 # Import data to play with. 

390 butler = self.makeButler(writeable=True) 

391 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

392 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

393 # Certify some biases into a CALIBRATION collection. 

394 registry = butler.registry 

395 registry.registerCollection("calibs", CollectionType.CALIBRATION) 

396 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai") 

397 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai") 

398 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai") 

399 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

400 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

401 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

402 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

403 registry.certify("calibs", [bias2a, bias3a], Timespan(t1, t2)) 

404 registry.certify("calibs", [bias2b], Timespan(t2, None)) 

405 registry.certify("calibs", [bias3b], Timespan(t2, t3)) 

406 # Insert some exposure dimension data. 

407 registry.insertDimensionData( 

408 "exposure", 

409 { 

410 "instrument": "Cam1", 

411 "id": 3, 

412 "obs_id": "three", 

413 "timespan": Timespan(t1, t2), 

414 "physical_filter": "Cam1-G", 

415 "day_obs": 20201114, 

416 "seq_num": 55, 

417 }, 

418 { 

419 "instrument": "Cam1", 

420 "id": 4, 

421 "obs_id": "four", 

422 "timespan": Timespan(t2, t3), 

423 "physical_filter": "Cam1-G", 

424 "day_obs": 20211114, 

425 "seq_num": 42, 

426 }, 

427 ) 

428 # Get some biases from raw-like data IDs. 

429 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 3, "detector": 2}, 

430 collections="calibs") 

431 self.assertEqual(bias2a_id, bias2a.id) 

432 bias3b_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 4, "detector": 3}, 

433 collections="calibs") 

434 self.assertEqual(bias3b_id, bias3b.id) 

435 

436 # Get using the kwarg form 

437 bias3b_id, _ = butler.get("bias", 

438 instrument="Cam1", exposure=4, detector=3, 

439 collections="calibs") 

440 self.assertEqual(bias3b_id, bias3b.id) 

441 

442 # Do it again but using the record information 

443 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure.obs_id": "three", 

444 "detector.full_name": "Ab"}, 

445 collections="calibs") 

446 self.assertEqual(bias2a_id, bias2a.id) 

447 bias3b_id, _ = butler.get("bias", {"exposure.obs_id": "four", 

448 "detector.full_name": "Ba"}, 

449 collections="calibs", instrument="Cam1") 

450 self.assertEqual(bias3b_id, bias3b.id) 

451 

452 # And again but this time using the alternate value rather than 

453 # the primary. 

454 bias3b_id, _ = butler.get("bias", {"exposure": "four", 

455 "detector": "Ba"}, 

456 collections="calibs", instrument="Cam1") 

457 self.assertEqual(bias3b_id, bias3b.id) 

458 

459 # And again but this time using the alternate value rather than 

460 # the primary and do it in the keyword arguments. 

461 bias3b_id, _ = butler.get("bias", 

462 exposure="four", detector="Ba", 

463 collections="calibs", instrument="Cam1") 

464 self.assertEqual(bias3b_id, bias3b.id) 

465 

466 # Now with implied record columns 

467 bias3b_id, _ = butler.get("bias", day_obs=20211114, seq_num=42, 

468 raft="B", name_in_raft="a", 

469 collections="calibs", instrument="Cam1") 

470 self.assertEqual(bias3b_id, bias3b.id) 

471 

472 # Ensure that spurious kwargs cause an exception. 

473 with self.assertRaises(ValueError): 

474 butler.get("bias", {"exposure.obs_id": "four", "immediate": True, 

475 "detector.full_name": "Ba"}, 

476 collections="calibs", instrument="Cam1") 

477 

478 with self.assertRaises(ValueError): 

479 butler.get("bias", day_obs=20211114, seq_num=42, 

480 raft="B", name_in_raft="a", 

481 collections="calibs", instrument="Cam1", immediate=True) 

482 

483 def testRegistryDefaults(self): 

484 """Test that we can default the collections and some data ID keys when 

485 constructing a butler. 

486 

487 Many tests that use default run already exist in ``test_butler.py``, so 

488 that isn't tested here. And while most of this functionality is 

489 implemented in `Registry`, we test it here instead of 

490 ``daf/butler/tests/registry.py`` because it shouldn't depend on the 

491 database backend at all. 

492 """ 

493 butler = self.makeButler(writeable=True) 

494 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

495 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

496 # Need to actually set defaults later, not at construction, because 

497 # we need to import the instrument before we can use it as a default. 

498 # Don't set a default instrument value for data IDs, because 'Cam1' 

499 # should be inferred by virtue of that being the only value in the 

500 # input collections. 

501 butler.registry.defaults = RegistryDefaults(collections=["imported_g"]) 

502 # Use findDataset without collections or instrument. 

503 ref = butler.registry.findDataset("flat", detector=2, physical_filter="Cam1-G") 

504 # Do the same with Butler.get; this should ultimately invoke a lot of 

505 # the same code, so it's a bit circular, but mostly we're checking that 

506 # it works at all. 

507 dataset_id, _ = butler.get("flat", detector=2, physical_filter="Cam1-G") 

508 self.assertEqual(ref.id, dataset_id) 

509 # Query for datasets. Test defaulting the data ID in both kwargs and 

510 # in the WHERE expression. 

511 queried_refs_1 = set(butler.registry.queryDatasets("flat", detector=2, physical_filter="Cam1-G")) 

512 self.assertEqual({ref}, queried_refs_1) 

513 queried_refs_2 = set(butler.registry.queryDatasets("flat", 

514 where="detector=2 AND physical_filter='Cam1-G'")) 

515 self.assertEqual({ref}, queried_refs_2) 

516 # Query for data IDs with a dataset constraint. 

517 queried_data_ids = set(butler.registry.queryDataIds({"instrument", "detector", "physical_filter"}, 

518 datasets={"flat"}, 

519 detector=2, physical_filter="Cam1-G")) 

520 self.assertEqual({ref.dataId}, queried_data_ids) 

521 # Add another instrument to the repo, and a dataset that uses it to 

522 # the `imported_g` collection. 

523 butler.registry.insertDimensionData("instrument", {"name": "Cam2"}) 

524 camera = DatasetType( 

525 "camera", 

526 dimensions=butler.registry.dimensions["instrument"].graph, 

527 storageClass="Camera", 

528 ) 

529 butler.registry.registerDatasetType(camera) 

530 butler.registry.insertDatasets(camera, [{"instrument": "Cam2"}], run="imported_g") 

531 # Initialize a new butler with `imported_g` as its default run. 

532 # This should not have a default instrument, because there are two. 

533 # Pass run instead of collections; this should set both. 

534 butler2 = Butler(butler=butler, run="imported_g") 

535 self.assertEqual(list(butler2.registry.defaults.collections), ["imported_g"]) 

536 self.assertEqual(butler2.registry.defaults.run, "imported_g") 

537 self.assertFalse(butler2.registry.defaults.dataId) 

538 # Initialize a new butler with an instrument default explicitly given. 

539 # Set collections instead of run, which should then be None. 

540 butler3 = Butler(butler=butler, collections=["imported_g"], instrument="Cam2") 

541 self.assertEqual(list(butler3.registry.defaults.collections), ["imported_g"]) 

542 self.assertIsNone(butler3.registry.defaults.run, None) 

543 self.assertEqual(butler3.registry.defaults.dataId.byName(), {"instrument": "Cam2"}) 

544 

545 def testJson(self): 

546 """Test JSON serialization mediated by registry. 

547 """ 

548 butler = self.makeButler(writeable=True) 

549 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

550 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", self.datasetsImportFile)) 

551 # Need to actually set defaults later, not at construction, because 

552 # we need to import the instrument before we can use it as a default. 

553 # Don't set a default instrument value for data IDs, because 'Cam1' 

554 # should be inferred by virtue of that being the only value in the 

555 # input collections. 

556 butler.registry.defaults = RegistryDefaults(collections=["imported_g"]) 

557 # Use findDataset without collections or instrument. 

558 ref = butler.registry.findDataset("flat", detector=2, physical_filter="Cam1-G") 

559 

560 # Transform the ref and dataset type to and from JSON 

561 # and check that it can be reconstructed properly 

562 

563 # Do it with the ref and a component ref in minimal and standard form 

564 compRef = ref.makeComponentRef("wcs") 

565 

566 for test_item in (ref, ref.datasetType, compRef, compRef.datasetType): 

567 for minimal in (False, True): 

568 json_str = test_item.to_json(minimal=minimal) 

569 from_json = type(test_item).from_json(json_str, registry=butler.registry) 

570 self.assertEqual(from_json, test_item, msg=f"From JSON '{json_str}' using registry") 

571 

572 # for minimal=False case also do a test without registry 

573 if not minimal: 

574 from_json = type(test_item).from_json(json_str, universe=butler.registry.dimensions) 

575 self.assertEqual(from_json, test_item, msg=f"From JSON '{json_str}' using universe") 

576 

577 def testJsonDimensionRecords(self): 

578 # Dimension Records 

579 butler = self.makeButler(writeable=True) 

580 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

581 

582 for dimension in ("detector", "visit"): 

583 records = butler.registry.queryDimensionRecords(dimension, instrument="HSC") 

584 for r in records: 

585 for minimal in (True, False): 

586 json_str = r.to_json(minimal=minimal) 

587 r_json = type(r).from_json(json_str, registry=butler.registry) 

588 self.assertEqual(r_json, r) 

589 # check with direct method 

590 simple = r.to_simple() 

591 fromDirect = type(simple).direct(**json.loads(json_str)) 

592 self.assertEqual(simple, fromDirect) 

593 # Also check equality of each of the components as dicts 

594 self.assertEqual(r_json.toDict(), r.toDict()) 

595 

596 def testWildcardQueries(self): 

597 """Test that different collection type queries work.""" 

598 

599 # Import data to play with. 

600 butler = self.makeButler(writeable=True) 

601 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

602 

603 # Create some collections 

604 created = {"collection", "u/user/test", "coll3"} 

605 for collection in created: 

606 butler.registry.registerCollection(collection, type=CollectionType.RUN) 

607 

608 collections = butler.registry.queryCollections() 

609 self.assertEqual(set(collections), created) 

610 

611 expressions = ( 

612 ("collection", {"collection"}), 

613 (..., created), 

614 ("*", created), 

615 (("collection", "*"), created), 

616 ("u/*", {"u/user/test"}), 

617 (re.compile("u.*"), {"u/user/test"}), 

618 (re.compile(".*oll.*"), {"collection", "coll3"}), 

619 ("*oll*", {"collection", "coll3"}), 

620 ((re.compile(r".*\d$"), "u/user/test"), {"coll3", "u/user/test"}), 

621 ("*[0-9]", {"coll3"}), 

622 ) 

623 for expression, expected in expressions: 

624 result = butler.registry.queryCollections(expression) 

625 self.assertEqual(set(result), expected) 

626 

627 

628class SimpleButlerUUIDTestCase(SimpleButlerTestCase): 

629 """Same as SimpleButlerTestCase but uses UUID-based datasets manager and 

630 loads datasets from YAML file with UUIDs. 

631 """ 

632 

633 datasetsManager = \ 

634 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID" 

635 datasetsImportFile = "datasets-uuid.yaml" 

636 datasetsIdType = uuid.UUID 

637 

638 

639class SimpleButlerMixedUUIDTestCase(SimpleButlerTestCase): 

640 """Same as SimpleButlerTestCase but uses UUID-based datasets manager and 

641 loads datasets from YAML file with integer IDs. 

642 """ 

643 

644 datasetsManager = \ 

645 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID" 

646 datasetsImportFile = "datasets.yaml" 

647 datasetsIdType = uuid.UUID 

648 

649 

650if __name__ == "__main__": 650 ↛ 651line 650 didn't jump to line 651, because the condition on line 650 was never true

651 unittest.main()