Coverage for python/lsst/pipe/base/_dataset_handle.py: 19%
77 statements
« prev ^ index » next coverage.py v7.2.3, created at 2023-04-19 04:01 -0700
« prev ^ index » next coverage.py v7.2.3, created at 2023-04-19 04:01 -0700
1# This file is part of pipe_base.
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/>.
21from __future__ import annotations
23__all__ = ["InMemoryDatasetHandle"]
25import dataclasses
26from typing import Any, Optional
28from frozendict import frozendict
29from lsst.daf.butler import DataCoordinate, DimensionUniverse, StorageClass, StorageClassFactory
32# Use an empty dataID as a default.
33def _default_dataId() -> DataCoordinate:
34 return DataCoordinate.makeEmpty(DimensionUniverse())
37@dataclasses.dataclass(frozen=True, init=False)
38class InMemoryDatasetHandle:
39 """An in-memory version of a `~lsst.daf.butler.DeferredDatasetHandle`.
41 If ``dataId`` is not specified, a default empty dataId will be constructed.
42 If ``kwargs`` are provided without specifying a ``dataId``, those
43 parameters will be converted into a dataId-like entity.
44 """
46 def __init__(
47 self,
48 inMemoryDataset: Any,
49 *,
50 storageClass: StorageClass | None = None,
51 parameters: dict[str, Any] | None = None,
52 dataId: dict[str, Any] | DataCoordinate | None = None,
53 **kwargs: Any,
54 ):
55 object.__setattr__(self, "inMemoryDataset", inMemoryDataset)
56 object.__setattr__(self, "storageClass", storageClass)
57 object.__setattr__(self, "parameters", parameters)
58 # Need to be able to construct a dataId from kwargs for convenience.
59 # This will not be a full DataCoordinate.
60 if dataId is None:
61 if kwargs:
62 dataId = frozendict(kwargs) # type: ignore
63 else:
64 dataId = DataCoordinate.makeEmpty(DimensionUniverse())
65 elif kwargs:
66 if isinstance(dataId, DataCoordinate):
67 dataId = DataCoordinate.standardize(kwargs, defaults=dataId, universe=dataId.universe)
68 else:
69 new = dict(dataId)
70 new.update(kwargs)
71 dataId = frozendict(new) # type: ignore
72 object.__setattr__(self, "dataId", dataId)
74 def get(
75 self,
76 *,
77 component: Optional[str] = None,
78 parameters: Optional[dict] = None,
79 storageClass: str | StorageClass | None = None,
80 ) -> Any:
81 """Retrieves the dataset pointed to by this handle
83 This handle may be used multiple times, possibly with different
84 parameters.
86 Parameters
87 ----------
88 component : `str` or None
89 If the deferred object is a component dataset type, this parameter
90 may specify the name of the component to use in the get operation.
91 parameters : `dict` or None
92 The parameters argument will be passed to the butler get method.
93 It defaults to None. If the value is not None, this dict will
94 be merged with the parameters dict used to construct the
95 `DeferredDatasetHandle` class.
96 storageClass : `StorageClass` or `str`, optional
97 The storage class to be used to override the Python type
98 returned by this method. By default the returned type matches
99 the type stored. Specifying a read `StorageClass` can force a
100 different type to be returned.
101 This type must be compatible with the original type.
103 Returns
104 -------
105 return : `object`
106 The dataset pointed to by this handle. This is the actual object
107 that was initially stored and not a copy. Modifying this object
108 will modify the stored object. If the stored object is `None` this
109 method always returns `None` regardless of any component request or
110 parameters.
112 Raises
113 ------
114 KeyError
115 Raised if a component or parameters are used but no storage
116 class can be found.
117 """
118 if self.inMemoryDataset is None:
119 return None
121 if self.parameters is not None:
122 mergedParameters = self.parameters.copy()
123 if parameters is not None:
124 mergedParameters.update(parameters)
125 elif parameters is not None:
126 mergedParameters = parameters
127 else:
128 mergedParameters = {}
130 returnStorageClass: StorageClass | None = None
131 if storageClass:
132 if isinstance(storageClass, str):
133 factory = StorageClassFactory()
134 returnStorageClass = factory.getStorageClass(storageClass)
135 else:
136 returnStorageClass = storageClass
138 if component or mergedParameters:
139 # This requires a storage class look up to locate the delegate
140 # class.
141 thisStorageClass = self._getStorageClass()
142 inMemoryDataset = self.inMemoryDataset
144 # Parameters for derived components are applied against the
145 # composite.
146 if component in thisStorageClass.derivedComponents:
147 thisStorageClass.validateParameters(parameters)
149 # Process the parameters (hoping this never modified the
150 # original object).
151 inMemoryDataset = thisStorageClass.delegate().handleParameters(
152 inMemoryDataset, mergedParameters
153 )
154 mergedParameters = {} # They have now been used
156 readStorageClass = thisStorageClass.derivedComponents[component]
157 else:
158 if component:
159 readStorageClass = thisStorageClass.components[component]
160 else:
161 readStorageClass = thisStorageClass
162 readStorageClass.validateParameters(mergedParameters)
164 if component:
165 inMemoryDataset = thisStorageClass.delegate().getComponent(inMemoryDataset, component)
167 if mergedParameters:
168 inMemoryDataset = readStorageClass.delegate().handleParameters(
169 inMemoryDataset, mergedParameters
170 )
171 if returnStorageClass:
172 return returnStorageClass.coerce_type(inMemoryDataset)
173 return inMemoryDataset
174 else:
175 # If there are no parameters or component requests the object
176 # can be returned as is, but possibly with conversion.
177 if returnStorageClass:
178 return returnStorageClass.coerce_type(self.inMemoryDataset)
179 return self.inMemoryDataset
181 def _getStorageClass(self) -> StorageClass:
182 """Return the relevant storage class.
184 Returns
185 -------
186 storageClass : `StorageClass`
187 The storage class associated with this handle, or one derived
188 from the python type of the stored object.
190 Raises
191 ------
192 KeyError
193 Raised if the storage class could not be found.
194 """
195 factory = StorageClassFactory()
196 if self.storageClass:
197 return factory.getStorageClass(self.storageClass)
199 # Need to match python type.
200 pytype = type(self.inMemoryDataset)
201 return factory.findStorageClass(pytype)
203 inMemoryDataset: Any
204 """The object to store in this dataset handle for later retrieval.
205 """
207 dataId: DataCoordinate | frozendict # type:ignore
208 """The `~lsst.daf.butler.DataCoordinate` associated with this dataset
209 handle.
210 """
212 storageClass: Optional[str] = None
213 """The name of the `~lsst.daf.butler.StorageClass` associated with this
214 dataset.
216 If `None`, the storage class will be looked up from the factory.
217 """
219 parameters: Optional[dict] = None
220 """Optional parameters that may be used to specify a subset of the dataset
221 to be loaded (`dict` or `None`).
222 """