Coverage for tests/test_simpleButler.py: 11%

255 statements  

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

30import json 

31import os 

32import re 

33import tempfile 

34import unittest 

35from typing import Any 

36 

37try: 

38 import numpy as np 

39except ImportError: 

40 np = None 

41 

42import astropy.time 

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

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

45from lsst.daf.butler.tests import DatastoreMock 

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

47 

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

49 

50 

51class SimpleButlerTestCase(unittest.TestCase): 

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

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

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

55 """ 

56 

57 datasetsManager = ( 

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

59 ) 

60 datasetsImportFile = "datasets-uuid.yaml" 

61 

62 def setUp(self): 

63 self.root = makeTestTempDir(TESTDIR) 

64 

65 def tearDown(self): 

66 removeTestTempDir(self.root) 

67 

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

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

70 config = ButlerConfig() 

71 

72 # make separate temporary directory for registry of this instance 

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

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

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

76 config["root"] = self.root 

77 

78 # have to make a registry first 

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

80 _RegistryFactory(registryConfig).create_from_config() 

81 

82 butler = Butler.from_config(config, **kwargs) 

83 DatastoreMock.apply(butler) 

84 return butler 

85 

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

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

88 other repository. 

89 

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

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

92 returns unchanged ref. 

93 """ 

94 return ref 

95 

96 def testReadBackwardsCompatibility(self): 

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

98 and commit to the daf_butler git repo. 

99 

100 Notes 

101 ----- 

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

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

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

105 the export format required for CALIBRATION collections to land. 

106 """ 

107 butler = self.makeButler(writeable=True) 

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

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

110 # the above does not raise. 

111 self.assertGreaterEqual( 

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

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

114 ) 

115 self.assertGreaterEqual( 

116 { 

117 (record.id, record.physical_filter) 

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

119 }, 

120 { 

121 (27136, "HSC-Z"), 

122 (11694, "HSC-G"), 

123 (23910, "HSC-R"), 

124 (11720, "HSC-Y"), 

125 (23900, "HSC-R"), 

126 (22646, "HSC-Y"), 

127 (1248, "HSC-I"), 

128 (19680, "HSC-I"), 

129 (1240, "HSC-I"), 

130 (424, "HSC-Y"), 

131 (19658, "HSC-I"), 

132 (344, "HSC-Y"), 

133 (1218, "HSC-R"), 

134 (1190, "HSC-Z"), 

135 (23718, "HSC-R"), 

136 (11700, "HSC-G"), 

137 (26036, "HSC-G"), 

138 (23872, "HSC-R"), 

139 (1170, "HSC-Z"), 

140 (1876, "HSC-Y"), 

141 }, 

142 ) 

143 

144 def testDatasetTransfers(self): 

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

146 back in again. 

147 """ 

148 # Import data to play with. 

149 butler1 = self.makeButler(writeable=True) 

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

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

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

153 # Export all datasets. 

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

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

156 # Import it all again. 

157 butler2 = self.makeButler(writeable=True) 

158 butler2.import_(filename=file.name) 

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

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

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

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

163 self.assertCountEqual( 

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

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

166 ) 

167 

168 def testImportTwice(self): 

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

170 importing them all back in again twice. 

171 """ 

172 # Import data to play with. 

173 butler1 = self.makeButler(writeable=True) 

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

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

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

177 # Export all datasets. 

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

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

180 butler2 = self.makeButler(writeable=True) 

181 # Import it once. 

182 butler2.import_(filename=file.name) 

183 # Import it again 

184 butler2.import_(filename=file.name) 

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

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

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

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

189 self.assertCountEqual( 

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

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

192 ) 

193 

194 def testCollectionTransfers(self): 

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

196 # Populate a registry with some datasets. 

197 butler1 = self.makeButler(writeable=True) 

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

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

200 registry1 = butler1.registry 

201 # Add some more collections. 

202 registry1.registerRun("run1") 

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

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

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

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

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

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

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

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

211 registry1.associate("tag1", flats1) 

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

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

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

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

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

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

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

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

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

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

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

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

224 

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

226 # Export all collections, and some datasets. 

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

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

229 # intentionally not topological order. 

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

231 exporter.saveCollection(collection) 

232 exporter.saveDatasets(flats1) 

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

234 # Import them into a new registry. 

235 butler2 = self.makeButler(writeable=True) 

236 butler2.import_(filename=file.name) 

237 registry2 = butler2.registry 

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

239 # themselves. 

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

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

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

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

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

245 self.assertEqual( 

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

247 ["tag1", "run1", "chain2"], 

248 ) 

249 self.assertEqual( 

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

251 ["calibration1", "run1"], 

252 ) 

253 # Check that tag collection contents are the same. 

254 self.maxDiff = None 

255 self.assertCountEqual( 

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

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

258 ) 

259 # Check that calibration collection contents are the same. 

260 self.assertCountEqual( 

261 [ 

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

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

264 ], 

265 [ 

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

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

268 ], 

269 ) 

270 

271 def testButlerGet(self): 

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

273 # Import data to play with. 

274 butler = self.makeButler(writeable=True) 

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

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

277 

278 # Find the DatasetRef for a flat 

279 coll = "imported_g" 

280 flat2g = butler.find_dataset( 

281 "flat", instrument="Cam1", full_name="Ab", physical_filter="Cam1-G", collections=coll 

282 ) 

283 

284 # Create a numpy integer to check that works fine 

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

286 

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

288 # arguments 

289 # Note that instrument.class_name does not work 

290 variants = ( 

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

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

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

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

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

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

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

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

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

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

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

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

303 ( 

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

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

306 ), 

307 ( 

308 { 

309 "detector.name_in_raft": "b", 

310 "detector.raft": "A", 

311 "instrument": "Cam1", 

312 "physical_filter": "Cam1-G", 

313 }, 

314 {}, 

315 ), 

316 # Duplicate (but valid) information. 

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

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

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

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

321 ) 

322 

323 for dataId, kwds in variants: 

324 try: 

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

326 except Exception as e: 

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

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

329 

330 # Check that bad combinations raise. 

331 variants = ( 

332 # Inconsistent detector information. 

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

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

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

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

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

338 # Under-specified. 

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

340 # Spurious kwargs. 

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

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

343 ) 

344 for dataId, kwds in variants: 

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

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

347 

348 def testGetCalibration(self): 

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

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

351 extra dimensions with temporal information. 

352 """ 

353 # Import data to play with. 

354 butler = self.makeButler(writeable=True) 

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

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

357 # Certify some biases into a CALIBRATION collection. 

358 registry = butler.registry 

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

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

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

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

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

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

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

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

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

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

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

370 # Insert some exposure dimension data. 

371 registry.insertDimensionData( 

372 "exposure", 

373 { 

374 "instrument": "Cam1", 

375 "id": 3, 

376 "obs_id": "three", 

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

378 "physical_filter": "Cam1-G", 

379 "day_obs": 20201114, 

380 "seq_num": 55, 

381 }, 

382 { 

383 "instrument": "Cam1", 

384 "id": 4, 

385 "obs_id": "four", 

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

387 "physical_filter": "Cam1-G", 

388 "day_obs": 20211114, 

389 "seq_num": 42, 

390 }, 

391 ) 

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

393 bias2a_id, _ = butler.get( 

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

395 ) 

396 self.assertEqual(bias2a_id, bias2a.id) 

397 bias3b_id, _ = butler.get( 

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

399 ) 

400 self.assertEqual(bias3b_id, bias3b.id) 

401 

402 # Get using the kwarg form 

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

404 self.assertEqual(bias3b_id, bias3b.id) 

405 

406 # Do it again but using the record information 

407 bias2a_id, _ = butler.get( 

408 "bias", 

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

410 collections="calibs", 

411 ) 

412 self.assertEqual(bias2a_id, bias2a.id) 

413 bias3b_id, _ = butler.get( 

414 "bias", 

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

416 collections="calibs", 

417 instrument="Cam1", 

418 ) 

419 self.assertEqual(bias3b_id, bias3b.id) 

420 

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

422 # the primary. 

423 bias3b_id, _ = butler.get( 

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

425 ) 

426 self.assertEqual(bias3b_id, bias3b.id) 

427 

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

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

430 bias3b_id, _ = butler.get( 

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

432 ) 

433 self.assertEqual(bias3b_id, bias3b.id) 

434 

435 # Now with implied record columns 

436 bias3b_id, _ = butler.get( 

437 "bias", 

438 day_obs=20211114, 

439 seq_num=42, 

440 raft="B", 

441 name_in_raft="a", 

442 collections="calibs", 

443 instrument="Cam1", 

444 ) 

445 self.assertEqual(bias3b_id, bias3b.id) 

446 

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

448 # that comes from the record. 

449 bias3b_id, _ = butler.get( 

450 "bias", 

451 dataId=dict( 

452 exposure=4, 

453 day_obs=20211114, 

454 seq_num=42, 

455 detector=3, 

456 instrument="Cam1", 

457 ), 

458 collections="calibs", 

459 ) 

460 self.assertEqual(bias3b_id, bias3b.id) 

461 

462 # Extra but inconsistent record values are a problem. 

463 with self.assertRaises(ValueError): 

464 bias3b_id, _ = butler.get( 

465 "bias", 

466 exposure=3, 

467 day_obs=20211114, 

468 seq_num=42, 

469 detector=3, 

470 collections="calibs", 

471 instrument="Cam1", 

472 ) 

473 

474 # Ensure that spurious kwargs cause an exception. 

475 with self.assertRaises(ValueError): 

476 butler.get( 

477 "bias", 

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

479 collections="calibs", 

480 instrument="Cam1", 

481 ) 

482 

483 with self.assertRaises(ValueError): 

484 butler.get( 

485 "bias", 

486 day_obs=20211114, 

487 seq_num=42, 

488 raft="B", 

489 name_in_raft="a", 

490 collections="calibs", 

491 instrument="Cam1", 

492 immediate=True, 

493 ) 

494 

495 def testRegistryDefaults(self): 

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

497 constructing a butler. 

498 

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

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

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

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

503 database backend at all. 

504 """ 

505 butler = self.makeButler(writeable=True) 

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

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

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

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

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

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

512 # input collections. 

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

514 # Use findDataset without collections or instrument. 

515 ref = butler.find_dataset("flat", detector=2, physical_filter="Cam1-G") 

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

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

518 # it works at all. 

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

520 self.assertEqual(ref.id, dataset_id) 

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

522 # in the WHERE expression. 

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

524 self.assertEqual({ref}, queried_refs_1) 

525 queried_refs_2 = set( 

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

527 ) 

528 self.assertEqual({ref}, queried_refs_2) 

529 # Query for data IDs with a dataset constraint. 

530 queried_data_ids = set( 

531 butler.registry.queryDataIds( 

532 {"instrument", "detector", "physical_filter"}, 

533 datasets={"flat"}, 

534 detector=2, 

535 physical_filter="Cam1-G", 

536 ) 

537 ) 

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

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

540 # the `imported_g` collection. 

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

542 camera = DatasetType( 

543 "camera", 

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

545 storageClass="Camera", 

546 ) 

547 butler.registry.registerDatasetType(camera) 

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

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

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

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

552 butler2 = Butler.from_config(butler=butler, run="imported_g") 

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

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

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

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

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

558 butler3 = Butler.from_config(butler=butler, collections=["imported_g"], instrument="Cam2") 

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

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

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

562 

563 # Check that repr() does not fail. 

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

565 r = repr(defaults) 

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

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

568 

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

570 r = repr(defaults) 

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

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

573 

574 def testJson(self): 

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

576 butler = self.makeButler(writeable=True) 

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

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

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

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

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

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

583 # input collections. 

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

585 # Use findDataset without collections or instrument. 

586 ref = butler.find_dataset("flat", detector=2, physical_filter="Cam1-G") 

587 

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

589 # and check that it can be reconstructed properly 

590 

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

592 compRef = ref.makeComponentRef("wcs") 

593 

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

595 for minimal in (False, True): 

596 json_str = test_item.to_json(minimal=minimal) 

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

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

599 

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

601 if not minimal: 

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

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

604 

605 def testJsonDimensionRecordsAndHtmlRepresentation(self): 

606 # Dimension Records 

607 butler = self.makeButler(writeable=True) 

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

609 

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

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

612 for r in records: 

613 for minimal in (True, False): 

614 json_str = r.to_json(minimal=minimal) 

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

616 self.assertEqual(r_json, r) 

617 # check with direct method 

618 simple = r.to_simple() 

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

620 self.assertEqual(simple, fromDirect) 

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

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

623 

624 # check the html representation of records 

625 r_html = r._repr_html_() 

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

627 self.assertIn(dimension, r_html) 

628 

629 def testWildcardQueries(self): 

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

631 # Import data to play with. 

632 butler = self.makeButler(writeable=True) 

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

634 

635 # Create some collections 

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

637 for collection in created: 

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

639 

640 collections = butler.registry.queryCollections() 

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

642 

643 expressions = ( 

644 ("collection", {"collection"}), 

645 (..., created), 

646 ("*", created), 

647 (("collection", "*"), created), 

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

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

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

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

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

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

654 ) 

655 for expression, expected in expressions: 

656 result = butler.registry.queryCollections(expression) 

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

658 

659 

660class SimpleButlerMixedUUIDTestCase(SimpleButlerTestCase): 

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

662 loads datasets from YAML file with integer IDs. 

663 """ 

664 

665 datasetsImportFile = "datasets.yaml" 

666 

667 

668if __name__ == "__main__": 

669 unittest.main()