Coverage for python / lsst / daf / butler / tests / _datasetsHelper.py: 38%
64 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-18 08:43 +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/>.
28from __future__ import annotations
30__all__ = (
31 "BadNoWriteFormatter",
32 "BadWriteFormatter",
33 "DatasetTestHelper",
34 "DatastoreTestHelper",
35 "MultiDetectorFormatter",
36)
38import os
39from collections.abc import Iterable, Mapping
40from typing import TYPE_CHECKING, Any
42from lsst.daf.butler import DataCoordinate, DatasetRef, DatasetType, DimensionGroup, StorageClass
43from lsst.daf.butler.datastore import Datastore
44from lsst.daf.butler.formatters.yaml import YamlFormatter
45from lsst.resources import ResourcePath
47if TYPE_CHECKING:
48 from lsst.daf.butler import Config, DatasetId
49 from lsst.daf.butler.datastore.cache_manager import AbstractDatastoreCacheManager
52class DatasetTestHelper:
53 """Helper methods for Datasets."""
55 def makeDatasetRef(
56 self,
57 datasetTypeName: str,
58 dimensions: DimensionGroup | Iterable[str],
59 storageClass: StorageClass | str,
60 dataId: DataCoordinate | Mapping[str, Any],
61 *,
62 id: DatasetId | None = None,
63 run: str | None = None,
64 conform: bool = True,
65 ) -> DatasetRef:
66 """Make a DatasetType and wrap it in a DatasetRef for a test.
68 Parameters
69 ----------
70 datasetTypeName : `str`
71 The name of the dataset type.
72 dimensions : `DimensionGroup` or `~collections.abc.Iterable` of `str`
73 The dimensions to use for this dataset type.
74 storageClass : `StorageClass` or `str`
75 The relevant storage class.
76 dataId : `DataCoordinate` or `~collections.abc.Mapping`
77 The data ID of this ref.
78 id : `DatasetId` or `None`, optional
79 The Id of this ref. Will be assigned automatically.
80 run : `str` or `None`, optional
81 The run for this ref. Will be assigned a default value if `None`.
82 conform : `bool`, optional
83 Whther to force the dataID to be checked for conformity with
84 the provided dimensions.
86 Returns
87 -------
88 ref : `DatasetRef`
89 The new ref.
90 """
91 return self._makeDatasetRef(
92 datasetTypeName,
93 dimensions,
94 storageClass,
95 dataId,
96 id=id,
97 run=run,
98 conform=conform,
99 )
101 def _makeDatasetRef(
102 self,
103 datasetTypeName: str,
104 dimensions: DimensionGroup | Iterable[str],
105 storageClass: StorageClass | str,
106 dataId: DataCoordinate | Mapping,
107 *,
108 id: DatasetId | None = None,
109 run: str | None = None,
110 conform: bool = True,
111 ) -> DatasetRef:
112 # helper for makeDatasetRef
114 # Pretend we have a parent if this looks like a composite
115 compositeName, componentName = DatasetType.splitDatasetTypeName(datasetTypeName)
116 parentStorageClass = StorageClass("component") if componentName else None
118 datasetType = DatasetType(
119 datasetTypeName, dimensions, storageClass, parentStorageClass=parentStorageClass
120 )
122 if run is None:
123 run = "dummy"
124 if not isinstance(dataId, DataCoordinate):
125 dataId = DataCoordinate.standardize(dataId, dimensions=datasetType.dimensions)
126 return DatasetRef(datasetType, dataId, id=id, run=run, conform=conform)
129class DatastoreTestHelper:
130 """Helper methods for Datastore tests."""
132 root: str | None
133 config: Config
134 datastoreType: type[Datastore]
135 configFile: str
137 def setUpDatastoreTests(self, registryClass: type, configClass: type[Config]) -> None:
138 """Shared setUp code for all Datastore tests.
140 Parameters
141 ----------
142 registryClass : `type`
143 Type of registry to use.
144 configClass : `type`
145 Type of config to use.
146 """
147 self.registry = registryClass()
148 self.config = configClass(self.configFile)
150 # Some subclasses override the working root directory
151 if self.root is not None:
152 self.datastoreType.setConfigRoot(self.root, self.config, self.config.copy())
154 def makeDatastore(self, sub: str | None = None) -> Datastore:
155 """Make a new Datastore instance of the appropriate type.
157 Parameters
158 ----------
159 sub : `str`, optional
160 If not None, the returned Datastore will be distinct from any
161 Datastore constructed with a different value of ``sub``. For
162 PosixDatastore, for example, the converse is also true, and ``sub``
163 is used as a subdirectory to form the new root.
165 Returns
166 -------
167 datastore : `Datastore`
168 Datastore constructed by this routine using the supplied
169 optional subdirectory if supported.
170 """
171 config = self.config.copy()
172 if sub is not None and self.root is not None:
173 self.datastoreType.setConfigRoot(os.path.join(self.root, sub), config, self.config)
174 if sub is not None:
175 # Ensure that each datastore gets its own registry
176 registryClass = type(self.registry)
177 registry = registryClass()
178 else:
179 registry = self.registry
180 return Datastore.fromConfig(config=config, bridgeManager=registry.getDatastoreBridgeManager())
183class BadWriteFormatter(YamlFormatter):
184 """A formatter that never works but does leave a file behind."""
186 can_read_from_uri = False
187 can_read_from_local_file = False
188 can_read_from_stream = False
190 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
191 return NotImplemented
193 def write_direct(
194 self,
195 in_memory_dataset: Any,
196 uri: ResourcePath,
197 cache_manager: AbstractDatastoreCacheManager | None = None,
198 ) -> bool:
199 """Write empty file and immediately fail.
201 Parameters
202 ----------
203 in_memory_dataset : `typing.Any`
204 The Python object to serialize.
205 uri : `lsst.resources.ResourcePath`
206 The location to write the content.
207 cache_manager : `AbstractDatastoreCacheManager`
208 Cache manager. Unused.
210 Raises
211 ------
212 RuntimeError
213 Raised every time specifically for testing this scenario.
214 """
215 uri.write(b"")
216 raise RuntimeError("Did not succeed in writing file.")
219class BadNoWriteFormatter(BadWriteFormatter):
220 """A formatter that always fails without writing anything."""
222 def write_direct(
223 self,
224 in_memory_dataset: Any,
225 uri: ResourcePath,
226 cache_manager: AbstractDatastoreCacheManager | None = None,
227 ) -> bool:
228 raise RuntimeError("Did not writing anything at all")
231class MultiDetectorFormatter(YamlFormatter):
232 """A formatter that requires a detector to be specified in the dataID."""
234 can_read_from_uri = True
236 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
237 if self.data_id is None:
238 raise RuntimeError("This formatter requires a dataId")
239 if "detector" not in self.data_id:
240 raise RuntimeError("This formatter requires detector to be present in dataId")
242 key = f"detector{self.data_id['detector']}"
244 data = super().read_from_uri(uri, component)
245 if key not in data:
246 raise RuntimeError(f"Could not find '{key}' in data file.")
248 return data[key]