Coverage for tests/test_simpleButler.py: 11%

255 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-21 09:55 +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 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 json 

25import os 

26import re 

27import tempfile 

28import unittest 

29from typing import Any 

30 

31try: 

32 import numpy as np 

33except ImportError: 

34 np = None 

35 

36import astropy.time 

37from lsst.daf.butler import Butler, ButlerConfig, CollectionType, DatasetId, DatasetRef, DatasetType, Timespan 

38from lsst.daf.butler.registry import RegistryConfig, RegistryDefaults, _RegistryFactory 

39from lsst.daf.butler.tests import DatastoreMock 

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

41 

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

43 

44 

45class SimpleButlerTestCase(unittest.TestCase): 

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

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

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

49 """ 

50 

51 datasetsManager = ( 

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

53 ) 

54 datasetsImportFile = "datasets-uuid.yaml" 

55 

56 def setUp(self): 

57 self.root = makeTestTempDir(TESTDIR) 

58 

59 def tearDown(self): 

60 removeTestTempDir(self.root) 

61 

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

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

64 config = ButlerConfig() 

65 

66 # make separate temporary directory for registry of this instance 

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

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

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

70 config["root"] = self.root 

71 

72 # have to make a registry first 

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

74 _RegistryFactory(registryConfig).create_from_config() 

75 

76 butler = Butler(config, **kwargs) 

77 DatastoreMock.apply(butler) 

78 return butler 

79 

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

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

82 other repository. 

83 

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

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

86 returns unchanged ref. 

87 """ 

88 return ref 

89 

90 def testReadBackwardsCompatibility(self): 

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

92 and commit to the daf_butler git repo. 

93 

94 Notes 

95 ----- 

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

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

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

99 the export format required for CALIBRATION collections to land. 

100 """ 

101 butler = self.makeButler(writeable=True) 

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

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

104 # the above does not raise. 

105 self.assertGreaterEqual( 

106 {record.id for record in butler.registry.queryDimensionRecords("detector", instrument="HSC")}, 

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

108 ) 

109 self.assertGreaterEqual( 

110 { 

111 (record.id, record.physical_filter) 

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

113 }, 

114 { 

115 (27136, "HSC-Z"), 

116 (11694, "HSC-G"), 

117 (23910, "HSC-R"), 

118 (11720, "HSC-Y"), 

119 (23900, "HSC-R"), 

120 (22646, "HSC-Y"), 

121 (1248, "HSC-I"), 

122 (19680, "HSC-I"), 

123 (1240, "HSC-I"), 

124 (424, "HSC-Y"), 

125 (19658, "HSC-I"), 

126 (344, "HSC-Y"), 

127 (1218, "HSC-R"), 

128 (1190, "HSC-Z"), 

129 (23718, "HSC-R"), 

130 (11700, "HSC-G"), 

131 (26036, "HSC-G"), 

132 (23872, "HSC-R"), 

133 (1170, "HSC-Z"), 

134 (1876, "HSC-Y"), 

135 }, 

136 ) 

137 

138 def testDatasetTransfers(self): 

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

140 back in again. 

141 """ 

142 # Import data to play with. 

143 butler1 = self.makeButler(writeable=True) 

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

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

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

147 # Export all datasets. 

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

149 exporter.saveDatasets(butler1.registry.queryDatasets(..., collections=...)) 

150 # Import it all again. 

151 butler2 = self.makeButler(writeable=True) 

152 butler2.import_(filename=file.name) 

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

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

155 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets1)) 

156 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets2)) 

157 self.assertCountEqual( 

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

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

160 ) 

161 

162 def testImportTwice(self): 

163 """Test exporting dimension records and datasets from a repo and then 

164 importing them all back in again twice. 

165 """ 

166 # Import data to play with. 

167 butler1 = self.makeButler(writeable=True) 

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

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

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

171 # Export all datasets. 

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

173 exporter.saveDatasets(butler1.registry.queryDatasets(..., collections=...)) 

174 butler2 = self.makeButler(writeable=True) 

175 # Import it once. 

176 butler2.import_(filename=file.name) 

177 # Import it again 

178 butler2.import_(filename=file.name) 

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

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

181 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets1)) 

182 self.assertTrue(all(isinstance(ref.id, DatasetId) for ref in datasets2)) 

183 self.assertCountEqual( 

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

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

186 ) 

187 

188 def testCollectionTransfers(self): 

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

190 # Populate a registry with some datasets. 

191 butler1 = self.makeButler(writeable=True) 

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

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

194 registry1 = butler1.registry 

195 # Add some more collections. 

196 registry1.registerRun("run1") 

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

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

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

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

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

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

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

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

205 registry1.associate("tag1", flats1) 

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

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

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

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

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

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

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

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

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

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

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

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

218 

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

220 # Export all collections, and some datasets. 

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

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

223 # intentionally not topological order. 

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

225 exporter.saveCollection(collection) 

226 exporter.saveDatasets(flats1) 

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

228 # Import them into a new registry. 

229 butler2 = self.makeButler(writeable=True) 

230 butler2.import_(filename=file.name) 

231 registry2 = butler2.registry 

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

233 # themselves. 

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

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

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

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

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

239 self.assertEqual( 

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

241 ["tag1", "run1", "chain2"], 

242 ) 

243 self.assertEqual( 

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

245 ["calibration1", "run1"], 

246 ) 

247 # Check that tag collection contents are the same. 

248 self.maxDiff = None 

249 self.assertCountEqual( 

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

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

252 ) 

253 # Check that calibration collection contents are the same. 

254 self.assertCountEqual( 

255 [ 

256 (self.comparableRef(assoc.ref), assoc.timespan) 

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

258 ], 

259 [ 

260 (self.comparableRef(assoc.ref), assoc.timespan) 

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

262 ], 

263 ) 

264 

265 def testButlerGet(self): 

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

267 # Import data to play with. 

268 butler = self.makeButler(writeable=True) 

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

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

271 

272 # Find the DatasetRef for a flat 

273 coll = "imported_g" 

274 flat2g = butler.registry.findDataset( 

275 "flat", instrument="Cam1", detector=2, physical_filter="Cam1-G", collections=coll 

276 ) 

277 

278 # Create a numpy integer to check that works fine 

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

280 

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

282 # arguments 

283 # Note that instrument.class_name does not work 

284 variants = ( 

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

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

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

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

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

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

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

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

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

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

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

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

297 ( 

298 {"detector.name_in_raft": "b", "detector.raft": "A"}, 

299 {"instrument": "Cam1", "physical_filter": "Cam1-G"}, 

300 ), 

301 ( 

302 { 

303 "detector.name_in_raft": "b", 

304 "detector.raft": "A", 

305 "instrument": "Cam1", 

306 "physical_filter": "Cam1-G", 

307 }, 

308 {}, 

309 ), 

310 # Duplicate (but valid) information. 

311 (None, {"instrument": "Cam1", "detector": 2, "raft": "A", "physical_filter": "Cam1-G"}), 

312 ({"detector": 2}, {"instrument": "Cam1", "raft": "A", "physical_filter": "Cam1-G"}), 

313 ({"raft": "A"}, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

314 ({"raft": "A"}, {"instrument": "Cam1", "detector": "Ab", "physical_filter": "Cam1-G"}), 

315 ) 

316 

317 for dataId, kwds in variants: 

318 try: 

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

320 except Exception as e: 

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

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

323 

324 # Check that bad combinations raise. 

325 variants = ( 

326 # Inconsistent detector information. 

327 (None, {"instrument": "Cam1", "detector": 2, "raft": "B", "physical_filter": "Cam1-G"}), 

328 ({"detector": 2}, {"instrument": "Cam1", "raft": "B", "physical_filter": "Cam1-G"}), 

329 ({"detector": 12}, {"instrument": "Cam1", "raft": "B", "physical_filter": "Cam1-G"}), 

330 ({"raft": "B"}, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

331 ({"raft": "B"}, {"instrument": "Cam1", "detector": "Ab", "physical_filter": "Cam1-G"}), 

332 # Under-specified. 

333 ({"raft": "B"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

334 # Spurious kwargs. 

335 (None, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G", "x": "y"}), 

336 ({"x": "y"}, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

337 ) 

338 for dataId, kwds in variants: 

339 with self.assertRaises((ValueError, LookupError)): 

340 butler.get("flat", dataId=dataId, collections=coll, **kwds) 

341 

342 def testGetCalibration(self): 

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

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

345 extra dimensions with temporal information. 

346 """ 

347 # Import data to play with. 

348 butler = self.makeButler(writeable=True) 

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

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

351 # Certify some biases into a CALIBRATION collection. 

352 registry = butler.registry 

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

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

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

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

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

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

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

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

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

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

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

364 # Insert some exposure dimension data. 

365 registry.insertDimensionData( 

366 "exposure", 

367 { 

368 "instrument": "Cam1", 

369 "id": 3, 

370 "obs_id": "three", 

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

372 "physical_filter": "Cam1-G", 

373 "day_obs": 20201114, 

374 "seq_num": 55, 

375 }, 

376 { 

377 "instrument": "Cam1", 

378 "id": 4, 

379 "obs_id": "four", 

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

381 "physical_filter": "Cam1-G", 

382 "day_obs": 20211114, 

383 "seq_num": 42, 

384 }, 

385 ) 

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

387 bias2a_id, _ = butler.get( 

388 "bias", {"instrument": "Cam1", "exposure": 3, "detector": 2}, collections="calibs" 

389 ) 

390 self.assertEqual(bias2a_id, bias2a.id) 

391 bias3b_id, _ = butler.get( 

392 "bias", {"instrument": "Cam1", "exposure": 4, "detector": 3}, collections="calibs" 

393 ) 

394 self.assertEqual(bias3b_id, bias3b.id) 

395 

396 # Get using the kwarg form 

397 bias3b_id, _ = butler.get("bias", instrument="Cam1", exposure=4, detector=3, collections="calibs") 

398 self.assertEqual(bias3b_id, bias3b.id) 

399 

400 # Do it again but using the record information 

401 bias2a_id, _ = butler.get( 

402 "bias", 

403 {"instrument": "Cam1", "exposure.obs_id": "three", "detector.full_name": "Ab"}, 

404 collections="calibs", 

405 ) 

406 self.assertEqual(bias2a_id, bias2a.id) 

407 bias3b_id, _ = butler.get( 

408 "bias", 

409 {"exposure.obs_id": "four", "detector.full_name": "Ba"}, 

410 collections="calibs", 

411 instrument="Cam1", 

412 ) 

413 self.assertEqual(bias3b_id, bias3b.id) 

414 

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

416 # the primary. 

417 bias3b_id, _ = butler.get( 

418 "bias", {"exposure": "four", "detector": "Ba"}, collections="calibs", instrument="Cam1" 

419 ) 

420 self.assertEqual(bias3b_id, bias3b.id) 

421 

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

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

424 bias3b_id, _ = butler.get( 

425 "bias", exposure="four", detector="Ba", collections="calibs", instrument="Cam1" 

426 ) 

427 self.assertEqual(bias3b_id, bias3b.id) 

428 

429 # Now with implied record columns 

430 bias3b_id, _ = butler.get( 

431 "bias", 

432 day_obs=20211114, 

433 seq_num=42, 

434 raft="B", 

435 name_in_raft="a", 

436 collections="calibs", 

437 instrument="Cam1", 

438 ) 

439 self.assertEqual(bias3b_id, bias3b.id) 

440 

441 # Allow a fully-specified dataId and unnecessary extra information 

442 # that comes from the record. 

443 bias3b_id, _ = butler.get( 

444 "bias", 

445 dataId=dict( 

446 exposure=4, 

447 day_obs=20211114, 

448 seq_num=42, 

449 detector=3, 

450 instrument="Cam1", 

451 ), 

452 collections="calibs", 

453 ) 

454 self.assertEqual(bias3b_id, bias3b.id) 

455 

456 # Extra but inconsistent record values are a problem. 

457 with self.assertRaises(ValueError): 

458 bias3b_id, _ = butler.get( 

459 "bias", 

460 exposure=3, 

461 day_obs=20211114, 

462 seq_num=42, 

463 detector=3, 

464 collections="calibs", 

465 instrument="Cam1", 

466 ) 

467 

468 # Ensure that spurious kwargs cause an exception. 

469 with self.assertRaises(ValueError): 

470 butler.get( 

471 "bias", 

472 {"exposure.obs_id": "four", "immediate": True, "detector.full_name": "Ba"}, 

473 collections="calibs", 

474 instrument="Cam1", 

475 ) 

476 

477 with self.assertRaises(ValueError): 

478 butler.get( 

479 "bias", 

480 day_obs=20211114, 

481 seq_num=42, 

482 raft="B", 

483 name_in_raft="a", 

484 collections="calibs", 

485 instrument="Cam1", 

486 immediate=True, 

487 ) 

488 

489 def testRegistryDefaults(self): 

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

491 constructing a butler. 

492 

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

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

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

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

497 database backend at all. 

498 """ 

499 butler = self.makeButler(writeable=True) 

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

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

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

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

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

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

506 # input collections. 

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

508 # Use findDataset without collections or instrument. 

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

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

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

512 # it works at all. 

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

514 self.assertEqual(ref.id, dataset_id) 

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

516 # in the WHERE expression. 

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

518 self.assertEqual({ref}, queried_refs_1) 

519 queried_refs_2 = set( 

520 butler.registry.queryDatasets("flat", where="detector=2 AND physical_filter='Cam1-G'") 

521 ) 

522 self.assertEqual({ref}, queried_refs_2) 

523 # Query for data IDs with a dataset constraint. 

524 queried_data_ids = set( 

525 butler.registry.queryDataIds( 

526 {"instrument", "detector", "physical_filter"}, 

527 datasets={"flat"}, 

528 detector=2, 

529 physical_filter="Cam1-G", 

530 ) 

531 ) 

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

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

534 # the `imported_g` collection. 

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

536 camera = DatasetType( 

537 "camera", 

538 dimensions=butler.dimensions["instrument"].graph, 

539 storageClass="Camera", 

540 ) 

541 butler.registry.registerDatasetType(camera) 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

556 

557 # Check that repr() does not fail. 

558 defaults = RegistryDefaults(collections=["imported_g"], run="test") 

559 r = repr(defaults) 

560 self.assertIn("collections=('imported_g',)", r) 

561 self.assertIn("run='test'", r) 

562 

563 defaults = RegistryDefaults(run="test", instrument="DummyCam", skypix="pix") 

564 r = repr(defaults) 

565 self.assertIn("skypix='pix'", r) 

566 self.assertIn("instrument='DummyCam'", r) 

567 

568 def testJson(self): 

569 """Test JSON serialization mediated by registry.""" 

570 butler = self.makeButler(writeable=True) 

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

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

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

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

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

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

577 # input collections. 

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

579 # Use findDataset without collections or instrument. 

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

581 

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

583 # and check that it can be reconstructed properly 

584 

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

586 compRef = ref.makeComponentRef("wcs") 

587 

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

589 for minimal in (False, True): 

590 json_str = test_item.to_json(minimal=minimal) 

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

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

593 

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

595 if not minimal: 

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

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

598 

599 def testJsonDimensionRecordsAndHtmlRepresentation(self): 

600 # Dimension Records 

601 butler = self.makeButler(writeable=True) 

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

603 

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

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

606 for r in records: 

607 for minimal in (True, False): 

608 json_str = r.to_json(minimal=minimal) 

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

610 self.assertEqual(r_json, r) 

611 # check with direct method 

612 simple = r.to_simple() 

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

614 self.assertEqual(simple, fromDirect) 

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

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

617 

618 # check the html representation of records 

619 r_html = r._repr_html_() 

620 self.assertTrue(isinstance(r_html, str)) 

621 self.assertIn(dimension, r_html) 

622 

623 def testWildcardQueries(self): 

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

625 # Import data to play with. 

626 butler = self.makeButler(writeable=True) 

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

628 

629 # Create some collections 

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

631 for collection in created: 

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

633 

634 collections = butler.registry.queryCollections() 

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

636 

637 expressions = ( 

638 ("collection", {"collection"}), 

639 (..., created), 

640 ("*", created), 

641 (("collection", "*"), created), 

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

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

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

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

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

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

648 ) 

649 for expression, expected in expressions: 

650 result = butler.registry.queryCollections(expression) 

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

652 

653 

654class SimpleButlerMixedUUIDTestCase(SimpleButlerTestCase): 

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

656 loads datasets from YAML file with integer IDs. 

657 """ 

658 

659 datasetsImportFile = "datasets.yaml" 

660 

661 

662if __name__ == "__main__": 

663 unittest.main()