Coverage for python / lsst / daf / butler / tests / _testRepo.py: 12%
145 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:10 -0800
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/>.
28from __future__ import annotations
30__all__ = [
31 "DatastoreMock",
32 "addDataIdValue",
33 "addDatasetType",
34 "expandUniqueId",
35 "makeTestCollection",
36 "makeTestRepo",
37]
39import random
40from collections.abc import Iterable, Mapping
41from typing import TYPE_CHECKING, Any
42from unittest.mock import MagicMock
44import sqlalchemy
46from lsst.daf.butler import (
47 Butler,
48 Config,
49 DataCoordinate,
50 DatasetRef,
51 DatasetType,
52 Dimension,
53 DimensionUniverse,
54 FileDataset,
55 StorageClass,
56)
58if TYPE_CHECKING:
59 from lsst.daf.butler import DatasetId
62def makeTestRepo(
63 root: str, dataIds: Mapping[str, Iterable] | None = None, *, config: Config | None = None, **kwargs: Any
64) -> Butler:
65 """Create an empty test repository.
67 Parameters
68 ----------
69 root : `str`
70 The location of the root directory for the repository.
71 dataIds : `~collections.abc.Mapping` [`str`, `iterable`], optional
72 A mapping keyed by the dimensions used in the test. Each value is an
73 iterable of names for that dimension (e.g., detector IDs for
74 `"detector"`). Related dimensions (e.g., instruments and detectors) are
75 linked arbitrarily, with values created for implied dimensions only
76 when needed. This parameter is provided for compatibility with old
77 code; newer code should make the repository, then call
78 `~lsst.daf.butler.tests.addDataIdValue`.
79 config : `lsst.daf.butler.Config`, optional
80 A configuration for the repository (for details, see
81 `lsst.daf.butler.Butler.makeRepo`). If omitted, creates a repository
82 with default dataset and storage types, but optimized for speed. The
83 defaults set ``.datastore.cls``, ``.datastore.checksum`` and
84 ``.registry.db``. If a supplied config does not specify these values
85 the internal defaults will be used to ensure that we have a usable
86 configuration.
87 **kwargs
88 Extra arguments to `lsst.daf.butler.Butler.makeRepo`.
90 Returns
91 -------
92 butler : `lsst.daf.butler.Butler`
93 A Butler referring to the new repository. This Butler is provided only
94 for additional setup; to keep test cases isolated, it is highly
95 recommended that each test create its own Butler with a unique
96 run/collection. See `makeTestCollection`.
98 Notes
99 -----
100 This function provides a "quick and dirty" repository for simple unit tests
101 that don't depend on complex data relationships. It is ill-suited for tests
102 where the structure of the data matters. If you need such a dataset, create
103 it directly or use a saved test dataset.
104 """
105 defaults = Config()
106 defaults["datastore", "cls"] = "lsst.daf.butler.datastores.inMemoryDatastore.InMemoryDatastore"
107 defaults["datastore", "checksum"] = False # In case of future changes
108 defaults["registry", "db"] = "sqlite:///<butlerRoot>/gen3.sqlite3"
110 if config:
111 defaults.update(config)
113 if not dataIds:
114 dataIds = {}
116 # Disable config root by default so that our registry override will
117 # not be ignored.
118 # newConfig guards against location-related keywords like outfile
119 newConfig = Butler.makeRepo(root, config=defaults, forceConfigRoot=False, **kwargs)
120 butler = Butler.from_config(newConfig, writeable=True)
121 dimensionRecords = _makeRecords(dataIds, butler.dimensions)
122 for dimension, records in dimensionRecords.items():
123 if butler.dimensions[dimension].has_own_table:
124 butler.registry.insertDimensionData(dimension, *records)
125 return butler
128def makeTestCollection(repo: Butler, uniqueId: str | None = None) -> Butler:
129 """Create a read/write Butler to a fresh collection.
131 Parameters
132 ----------
133 repo : `lsst.daf.butler.Butler`
134 A previously existing Butler to a repository, such as that returned by
135 `~lsst.daf.butler.Butler.makeRepo` or `makeTestRepo`.
136 uniqueId : `str`, optional
137 A collection ID guaranteed by external code to be unique across all
138 calls to ``makeTestCollection`` for the same repository.
140 Returns
141 -------
142 butler : `lsst.daf.butler.Butler`
143 A Butler referring to a new collection in the repository at ``root``.
144 The collection is (almost) guaranteed to be new.
146 Notes
147 -----
148 This function creates a single run collection that does not necessarily
149 conform to any repository conventions. It is only suitable for creating an
150 isolated test area, and not for repositories intended for real data
151 processing or analysis.
152 """
153 if not uniqueId:
154 # Create a "random" collection name
155 # Speed matters more than cryptographic guarantees
156 uniqueId = str(random.randrange(1_000_000_000))
157 collection = "test_" + uniqueId
158 return Butler.from_config(butler=repo, run=collection)
161def _makeRecords(dataIds: Mapping[str, Iterable], universe: DimensionUniverse) -> Mapping[str, Iterable]:
162 """Create cross-linked dimension records from a collection of
163 data ID values.
165 Parameters
166 ----------
167 dataIds : `~collections.abc.Mapping` [`str`, `iterable`]
168 A mapping keyed by the dimensions of interest. Each value is an
169 iterable of names for that dimension (e.g., detector IDs for
170 `"detector"`).
171 universe : lsst.daf.butler.DimensionUniverse
172 Set of all known dimensions and their relationships.
174 Returns
175 -------
176 dataIds : `~collections.abc.Mapping` [`str`, `iterable`]
177 A mapping keyed by the dimensions of interest, giving one
178 `~lsst.daf.butler.DimensionRecord` for each input name. Related
179 dimensions (e.g., instruments and detectors) are linked arbitrarily.
180 """
181 # Create values for all dimensions that are (recursive) required or implied
182 # dependencies of the given ones.
183 complete_data_id_values = {}
184 for dimension_name in universe.conform(dataIds.keys()).names:
185 if dimension_name in dataIds:
186 complete_data_id_values[dimension_name] = list(dataIds[dimension_name])
187 if dimension_name not in complete_data_id_values:
188 complete_data_id_values[dimension_name] = [
189 _makeRandomDataIdValue(universe.dimensions[dimension_name])
190 ]
192 # Start populating dicts that will become DimensionRecords by providing
193 # alternate keys like detector names
194 record_dicts_by_dimension_name: dict[str, list[dict[str, str | int | bytes]]] = {}
195 for name, values in complete_data_id_values.items():
196 record_dicts_by_dimension_name[name] = []
197 dimension_el = universe[name]
198 for value in values:
199 # _fillAllKeys wants Dimension and not DimensionElement.
200 # universe.__getitem__ says it returns DimensionElement but this
201 # really does also seem to be a Dimension here.
202 record_dicts_by_dimension_name[name].append(
203 _fillAllKeys(dimension_el, value) # type: ignore[arg-type]
204 )
206 # Pick cross-relationships arbitrarily
207 for name, record_dicts in record_dicts_by_dimension_name.items():
208 dimension_el = universe[name]
209 for record_dict in record_dicts:
210 for other in dimension_el.dimensions:
211 if other != dimension_el:
212 relation = record_dicts_by_dimension_name[other.name][0]
213 record_dict[other.name] = relation[other.primaryKey.name]
215 return {
216 dimension: [universe[dimension].RecordClass(**record_dict) for record_dict in record_dicts]
217 for dimension, record_dicts in record_dicts_by_dimension_name.items()
218 }
221def _fillAllKeys(dimension: Dimension, value: str | int) -> dict[str, str | int | bytes]:
222 """Create an arbitrary mapping of all required keys for a given dimension
223 that do not refer to other dimensions.
225 Parameters
226 ----------
227 dimension : `lsst.daf.butler.Dimension`
228 The dimension for which to generate a set of keys (e.g., detector).
229 value
230 The value assigned to ``dimension`` (e.g., detector ID).
232 Returns
233 -------
234 expandedValue : `dict` [`str`]
235 A mapping of dimension keys to values. ``dimension's`` primary key
236 maps to ``value``, but all other mappings (e.g., detector name)
237 are arbitrary.
238 """
239 expandedValue: dict[str, str | int | bytes] = {}
240 for key in dimension.uniqueKeys:
241 if key.nbytes:
242 # For `bytes` fields, we want something that casts at least `str`
243 # and `int` values to bytes and yields b'' when called with no
244 # arguments (as in the except block below). Unfortunately, the
245 # `bytes` type itself fails for both `str` and `int`, but this
246 # lambda does what we need. This particularly important for the
247 # skymap dimensions' bytes 'hash' field, which has a unique
248 # constraint; without this, all skymaps would get a hash of b''
249 # and end up conflicting.
250 castType = lambda *args: str(*args).encode() # noqa: E731
251 else:
252 castType = key.dtype().python_type
253 try:
254 castValue = castType(value)
255 except TypeError:
256 castValue = castType()
257 expandedValue[key.name] = castValue
258 for key in dimension.metadata:
259 if not key.nullable:
260 expandedValue[key.name] = key.dtype().python_type(value)
261 return expandedValue
264def _makeRandomDataIdValue(dimension: Dimension) -> int | str:
265 """Generate a random value of the appropriate type for a data ID key.
267 Parameters
268 ----------
269 dimension : `Dimension`
270 Dimension the value corresponds to.
272 Returns
273 -------
274 value : `int` or `str`
275 Random value.
276 """
277 if dimension.primaryKey.getPythonType() is str:
278 return str(random.randrange(1000))
279 else:
280 return random.randrange(1000)
283def expandUniqueId(butler: Butler, partialId: Mapping[str, Any]) -> DataCoordinate:
284 """Return a complete data ID matching some criterion.
286 Parameters
287 ----------
288 butler : `lsst.daf.butler.Butler`
289 The repository to query.
290 partialId : `~collections.abc.Mapping` [`str`]
291 A mapping of known dimensions and values.
293 Returns
294 -------
295 dataId : `lsst.daf.butler.DataCoordinate`
296 The unique data ID that matches ``partialId``.
298 Raises
299 ------
300 ValueError
301 Raised if ``partialId`` does not uniquely identify a data ID.
303 Notes
304 -----
305 This method will only work correctly if all dimensions attached to the
306 target dimension (eg., "physical_filter" for "visit") are known to the
307 repository, even if they're not needed to identify a dataset. This function
308 is only suitable for certain kinds of test repositories, and not for
309 repositories intended for real data processing or analysis.
311 Examples
312 --------
313 .. code-block:: py
315 >>> butler = makeTestRepo(
316 "testdir", {"instrument": ["notACam"], "detector": [1]})
317 >>> expandUniqueId(butler, {"detector": 1})
318 DataCoordinate({instrument, detector}, ('notACam', 1))
319 """
320 # The example is *not* a doctest because it requires dangerous I/O
321 registry = butler.registry
322 dimensions = registry.dimensions.conform(partialId.keys()).required
324 query = " AND ".join(f"{dimension} = {value!r}" for dimension, value in partialId.items())
326 # Much of the purpose of this function is to do something we explicitly
327 # reject most of the time: query for a governor dimension (e.g. instrument)
328 # given something that depends on it (e.g. visit), hence check=False.
329 dataId = list(registry.queryDataIds(dimensions, where=query, check=False))
330 if len(dataId) == 1:
331 return dataId[0]
332 else:
333 raise ValueError(f"Found {len(dataId)} matches for {partialId}, expected 1.")
336def _findOrInventDataIdValue(
337 butler: Butler, data_id: dict[str, str | int], dimension: Dimension
338) -> tuple[str | int, bool]:
339 """Look up an arbitrary value for a dimension that is consistent with a
340 partial data ID that does not specify that dimension, or invent one if no
341 such value exists.
343 Parameters
344 ----------
345 butler : `Butler`
346 Butler to use to look up data ID values.
347 data_id : `dict` [ `str`, `str` or `int` ]
348 Dictionary of possibly-related data ID values.
349 dimension : `Dimension`
350 Dimension to obtain a value for.
352 Returns
353 -------
354 value : `int` or `str`
355 Value for this dimension.
356 invented : `bool`
357 `True` if the value had to be invented, `False` if a compatible value
358 already existed.
359 """
360 # No values given by caller for this dimension. See if any exist
361 # in the registry that are consistent with the values of dimensions
362 # we do have:
363 match_data_id = {key: data_id[key] for key in data_id.keys() & dimension.dimensions.names}
364 matches = list(butler.registry.queryDimensionRecords(dimension, dataId=match_data_id).limit(1))
365 if not matches:
366 # Nothing in the registry matches: invent a data ID value
367 # with the right type (actual value does not matter).
368 # We may or may not actually make a record with this; that's
369 # easier to check later.
370 dimension_value = _makeRandomDataIdValue(dimension)
371 return dimension_value, True
372 else:
373 # A record does exist in the registry. Use its data ID value.
374 dim_value = matches[0].dataId[dimension.name]
375 assert dim_value is not None
376 return dim_value, False
379def _makeDimensionRecordDict(data_id: dict[str, str | int], dimension: Dimension) -> dict[str, Any]:
380 """Create a dictionary that can be used to build a `DimensionRecord` that
381 is consistent with the given data ID.
383 Parameters
384 ----------
385 data_id : `dict` [ `str`, `str` or `int` ]
386 Dictionary that contains values for at least all of
387 ``dimension.dimensions.names`` (the main dimension, its recursive
388 required dependencies, and its non-recursive implied dependencies).
389 dimension : `Dimension`
390 Dimension to build a record dictionary for.
392 Returns
393 -------
394 record_dict : `dict` [ `str`, `object` ]
395 Dictionary that can be passed as ``**kwargs`` to this dimensions
396 record class constructor.
397 """
398 # Add the primary key field for this dimension.
399 record_dict: dict[str, Any] = {dimension.primaryKey.name: data_id[dimension.name]}
400 # Define secondary keys (e.g., detector name given detector id)
401 record_dict.update(_fillAllKeys(dimension, data_id[dimension.name]))
402 # Set the foreign key values for any related dimensions that should
403 # appear in the record.
404 for related_dimension in dimension.dimensions:
405 if related_dimension.name != dimension.name:
406 record_dict[related_dimension.name] = data_id[related_dimension.name]
407 return record_dict
410def addDataIdValue(butler: Butler, dimension: str, value: str | int, **related: str | int) -> None:
411 """Add the records that back a new data ID to a repository.
413 Parameters
414 ----------
415 butler : `lsst.daf.butler.Butler`
416 The repository to update.
417 dimension : `str`
418 The name of the dimension to gain a new value.
419 value : `str` or `int`
420 The value to register for the dimension.
421 **related : `typing.Any`
422 Any existing dimensions to be linked to ``value``.
424 Notes
425 -----
426 Related dimensions (e.g., the instrument associated with a detector) may be
427 specified using ``related``, which requires a value for those dimensions to
428 have been added to the repository already (generally with a previous call
429 to `addDataIdValue`. Any dependencies of the given dimension that are not
430 included in ``related`` will be linked to existing values arbitrarily, and
431 (for implied dependencies only) created and also inserted into the registry
432 if they do not exist. Values for required dimensions and those given in
433 ``related`` are never created.
435 Because this function creates filler data, it is only suitable for test
436 repositories. It should not be used for repositories intended for real data
437 processing or analysis, which have known dimension values.
439 Examples
440 --------
441 See the guide on :ref:`using-butler-in-tests-make-repo` for usage examples.
442 """
443 # Example is not doctest, because it's probably unsafe to create even an
444 # in-memory butler in that environment.
445 try:
446 full_dimension = butler.dimensions[dimension]
447 except KeyError as e:
448 raise ValueError from e
449 # Bad keys ignored by registry code
450 extra_keys = related.keys() - full_dimension.minimal_group.names
451 if extra_keys:
452 raise ValueError(
453 f"Unexpected keywords {extra_keys} not found in {full_dimension.minimal_group.names}"
454 )
456 # Assemble a dictionary data ID holding the given primary dimension value
457 # and all of the related ones.
458 data_id: dict[str, int | str] = {dimension: value}
459 data_id.update(related)
461 # Compute the set of all dimensions that these recursively depend on.
462 all_dimensions = butler.dimensions.conform(data_id.keys())
464 # Create dicts that will become DimensionRecords for all of these data IDs.
465 # This iteration is guaranteed to be in topological order, so we can count
466 # on new data ID values being invented before they are needed.
467 record_dicts_by_dimension: dict[Dimension, dict[str, Any]] = {}
468 for dimension_name in all_dimensions.names:
469 dimension_obj = butler.dimensions.dimensions[dimension_name]
470 dimension_value = data_id.get(dimension_name)
471 if dimension_value is None:
472 data_id[dimension_name], invented = _findOrInventDataIdValue(butler, data_id, dimension_obj)
473 if not invented:
474 # No need to make a new record; one already exists.
475 continue
476 if dimension_name in related:
477 # Caller passed in a value of this dimension explicitly, but it
478 # isn't the primary dimension they asked to have a record created
479 # for. That means they expect this record to already exist.
480 continue
481 if dimension_name != dimension and dimension_name in all_dimensions.required:
482 # We also don't want to automatically create new dimension records
483 # for required dimensions (except for the main dimension the caller
484 # asked for); those are also asserted by the caller to already
485 # exist.
486 continue
487 if not dimension_obj.has_own_table:
488 # Don't need to bother generating full records for dimensions whose
489 # records are not actually stored.
490 continue
491 record_dicts_by_dimension[dimension_obj] = _makeDimensionRecordDict(data_id, dimension_obj)
493 # Sync those dimension record dictionaries with the database.
494 for dimension_obj, record_dict in record_dicts_by_dimension.items():
495 record = dimension_obj.RecordClass(**record_dict)
496 try:
497 butler.registry.syncDimensionData(dimension_obj, record)
498 except sqlalchemy.exc.IntegrityError as e:
499 raise RuntimeError(
500 "Could not create data ID value. Automatic relationship generation "
501 "may have failed; try adding keywords to assign a specific instrument, "
502 "physical_filter, etc. based on the nested exception message."
503 ) from e
506def addDatasetType(butler: Butler, name: str, dimensions: set[str], storageClass: str) -> DatasetType:
507 """Add a new dataset type to a repository.
509 Parameters
510 ----------
511 butler : `lsst.daf.butler.Butler`
512 The repository to update.
513 name : `str`
514 The name of the dataset type.
515 dimensions : `set` [`str`]
516 The dimensions of the new dataset type.
517 storageClass : `str`
518 The storage class the dataset will use.
520 Returns
521 -------
522 datasetType : `lsst.daf.butler.DatasetType`
523 The new type.
525 Raises
526 ------
527 ValueError
528 Raised if the dimensions or storage class is invalid.
530 Notes
531 -----
532 Dataset types are shared across all collections in a repository, so this
533 function does not need to be run for each collection.
534 """
535 try:
536 datasetType = DatasetType(name, dimensions, storageClass, universe=butler.dimensions)
537 butler.registry.registerDatasetType(datasetType)
538 return datasetType
539 except KeyError as e:
540 raise ValueError from e
543class DatastoreMock:
544 """Mocks a butler datastore.
546 Has functions that mock the datastore in a butler. Provides an `apply`
547 function to replace the relevent butler datastore functions with the mock
548 functions.
549 """
551 @staticmethod
552 def apply(butler: Butler) -> None:
553 """Apply datastore mocks to a butler.
555 Parameters
556 ----------
557 butler : `~lsst.daf.butler.Butler`
558 Butler to be modified.
559 """
560 butler._datastore.export = DatastoreMock._mock_export # type: ignore
561 butler._datastore.get = DatastoreMock._mock_get # type: ignore
562 butler._datastore.ingest = MagicMock() # type: ignore
564 @staticmethod
565 def _mock_export(
566 refs: Iterable[DatasetRef], *, directory: str | None = None, transfer: str | None = None
567 ) -> Iterable[FileDataset]:
568 """Mock of `Datastore.export` that satisfies the requirement that
569 the refs passed in are included in the `FileDataset` objects
570 returned.
572 This can be used to construct a `Datastore` mock that can be used
573 in repository export via::
575 datastore = unittest.mock.Mock(spec=Datastore)
576 datastore.export = DatastoreMock._mock_export
578 """
579 for ref in refs:
580 yield FileDataset(
581 refs=[ref], path="mock/path", formatter="lsst.daf.butler.formatters.json.JsonFormatter"
582 )
584 @staticmethod
585 def _mock_get(
586 ref: DatasetRef,
587 parameters: Mapping[str, Any] | None = None,
588 storageClass: StorageClass | str | None = None,
589 ) -> tuple[DatasetId, Mapping[str, Any] | None]:
590 """Mock of `Datastore.get` that just returns the integer dataset ID
591 value and parameters it was given.
592 """
593 return (ref.id, parameters)