Coverage for python/lsst/pipe/base/_dataset_handle.py: 19%

77 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-04-22 02:19 -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 

22 

23__all__ = ["InMemoryDatasetHandle"] 

24 

25import dataclasses 

26from typing import Any, Optional 

27 

28from frozendict import frozendict 

29from lsst.daf.butler import DataCoordinate, DimensionUniverse, StorageClass, StorageClassFactory 

30 

31 

32# Use an empty dataID as a default. 

33def _default_dataId() -> DataCoordinate: 

34 return DataCoordinate.makeEmpty(DimensionUniverse()) 

35 

36 

37@dataclasses.dataclass(frozen=True, init=False) 

38class InMemoryDatasetHandle: 

39 """An in-memory version of a `~lsst.daf.butler.DeferredDatasetHandle`. 

40 

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 """ 

45 

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) 

73 

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 

82 

83 This handle may be used multiple times, possibly with different 

84 parameters. 

85 

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. 

102 

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. 

111 

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 

120 

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 = {} 

129 

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 

137 

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 

143 

144 # Parameters for derived components are applied against the 

145 # composite. 

146 if component in thisStorageClass.derivedComponents: 

147 thisStorageClass.validateParameters(parameters) 

148 

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 

155 

156 readStorageClass = thisStorageClass.derivedComponents[component] 

157 else: 

158 if component: 

159 readStorageClass = thisStorageClass.components[component] 

160 else: 

161 readStorageClass = thisStorageClass 

162 readStorageClass.validateParameters(mergedParameters) 

163 

164 if component: 

165 inMemoryDataset = thisStorageClass.delegate().getComponent(inMemoryDataset, component) 

166 

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 

180 

181 def _getStorageClass(self) -> StorageClass: 

182 """Return the relevant storage class. 

183 

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. 

189 

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) 

198 

199 # Need to match python type. 

200 pytype = type(self.inMemoryDataset) 

201 return factory.findStorageClass(pytype) 

202 

203 inMemoryDataset: Any 

204 """The object to store in this dataset handle for later retrieval. 

205 """ 

206 

207 dataId: DataCoordinate | frozendict # type:ignore 

208 """The `~lsst.daf.butler.DataCoordinate` associated with this dataset 

209 handle. 

210 """ 

211 

212 storageClass: Optional[str] = None 

213 """The name of the `~lsst.daf.butler.StorageClass` associated with this 

214 dataset. 

215 

216 If `None`, the storage class will be looked up from the factory. 

217 """ 

218 

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 """