Coverage for python/lsst/images/serialization/_tables.py: 59%

1# This file is part of lsst-images. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14__all__ = ( 

15 "TableColumnModel", 

16 "TableModel", 

17 "UnsupportedTableError", 

18) 

19 

20import operator 

21from typing import TYPE_CHECKING 

22 

23import astropy.units 

24import numpy as np 

25import numpy.typing as npt 

26import pydantic 

27 

28from ._asdf_utils import ArrayReferenceModel, InlineArrayModel, Unit 

29from ._common import ArchiveReadError 

30from ._dtypes import NumberType 

31 

32if TYPE_CHECKING: 

33 import astropy.table 

34 

35 

36class UnsupportedTableError(NotImplementedError): 

37 """Exception raised if a table object has column types or structure that 

38 are not supported by this library. 

39 """ 

40 

41 

42class TableColumnModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

43 """Model for a subset of the ASDF table/column schema.""" 

44 

45 data: InlineArrayModel | ArrayReferenceModel 

46 """Column data.""" 

47 

48 name: str 

49 """Name of the column.""" 

50 

51 description: str = pydantic.Field(default="", exclude_if=operator.not_) 

52 """Extended description of the column.""" 

53 

54 unit: Unit | None = pydantic.Field(default=None, exclude_if=operator.not_) 

55 """Units of the column.""" 

56 

57 meta: dict[str, int | float | str | bool | None] = pydantic.Field( 

58 default_factory=dict, exclude_if=operator.not_, description="Free-form metadata for the column." 

59 ) 

60 

61 model_config = pydantic.ConfigDict( 

62 json_schema_extra={ 

63 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

64 "id": "http://stsci.edu/schemas/asdf/table/column-1.1.0", 

65 "tag": "!table/column-1.1.0", 

66 } 

67 ) 

68 

69 @classmethod 

70 def from_record_dtype(cls, dtype: npt.DTypeLike) -> list[TableColumnModel]: 

71 """Extract a list of column definitions from a structured numpy dtype. 

72 

73 Parameters 

74 ---------- 

75 dtype 

76 Object convertible to `numpy.dtype`. 

77 

78 Notes 

79 ----- 

80 This sets the `data` field to an `ArrayReferenceModel` with ``source`` 

81 set to an empty string. This will need to be modified later. 

82 """ 

83 dtype = np.dtype(dtype) 

84 result: list[TableColumnModel] = [] 

85 if dtype.fields is None: 85 ↛ 86line 85 didn't jump to line 86 because the condition on line 85 was never true

86 raise TypeError(f"{dtype} is not a structured dtype.") 

87 for name, (field_dtype, *_) in dtype.fields.items(): 

88 # TODO: support string and variable-length array columns here. 

89 try: 

90 datatype, shape = NumberType.from_numpy_with_shape(field_dtype) 

91 except TypeError: 

92 raise UnsupportedTableError(f"Column type {field_dtype} is not supported.") from None 

93 result.append( 

94 TableColumnModel( 

95 data=ArrayReferenceModel(source="", datatype=datatype, shape=list(shape)), 

96 name=name, 

97 ) 

98 ) 

99 return result 

100 

101 @classmethod 

102 def from_record_array(cls, array: np.ndarray, inline: bool = False) -> list[TableColumnModel]: 

103 """Extract a list of column definitions from a structured numpy array. 

104 

105 Parameters 

106 ---------- 

107 array 

108 A table-like array. 

109 inline 

110 Whether to store the array data directly in the columns. 

111 

112 Notes 

113 ----- 

114 When ``inline=False``, this sets the `data` field to an 

115 `ArrayReferenceModel` with ``source`` set to an empty string. This 

116 will need to be modified later. 

117 """ 

118 if not inline: 

119 return cls.from_record_dtype(array.dtype) 

120 result: list[TableColumnModel] = [] 

121 if array.dtype.fields is None: 

122 raise TypeError(f"{array.dtype} is not a structured dtype.") 

123 for name, (field_dtype, *_) in array.dtype.fields.items(): 

124 # TODO: support string and variable-length array columns here. 

125 try: 

126 datatype, shape = NumberType.from_numpy_with_shape(field_dtype) 

127 except TypeError: 

128 raise UnsupportedTableError(f"Column type {field_dtype} is not supported.") from None 

129 result.append( 

130 TableColumnModel( 

131 data=InlineArrayModel(data=array[name].tolist(), datatype=datatype), 

132 name=name, 

133 ) 

134 ) 

135 return result 

136 

137 @classmethod 

138 def from_table(cls, table: astropy.table.Table, inline: bool = False) -> list[TableColumnModel]: 

139 """Extract column definitions and (optionally) data from an Astropy 

140 table. 

141 

142 Parameters 

143 ---------- 

144 table 

145 Astropy table whose columns are described. 

146 inline 

147 If `True`, store the column data inline instead of as an array 

148 reference. 

149 """ 

150 return [cls.from_column(c, inline=inline) for c in table.columns.values()] 

151 

152 @classmethod 

153 def from_column(cls, column: astropy.table.Column, inline: bool = False) -> TableColumnModel: 

154 """Extract a column definition and (optionally) data from an Astropy 

155 column. 

156 

157 Parameters 

158 ---------- 

159 column 

160 Astropy column to describe. 

161 inline 

162 If `True`, store the column data inline instead of as an array 

163 reference. 

164 

165 Notes 

166 ----- 

167 When ``inline=False`, this sets the `data` field to an 

168 `ArrayReferenceModel` with ``source`` set to an empty string. This 

169 will need to be modified later. 

170 """ 

171 # TODO: support string and variable-length array columns here. 

172 try: 

173 datatype = NumberType.from_numpy(column.dtype) 

174 except TypeError: 

175 raise UnsupportedTableError(f"Column type {column.dtype} is not supported.") from None 

176 

177 data = ( 

178 InlineArrayModel(data=column.tolist(), datatype=datatype) 

179 if inline 

180 else ArrayReferenceModel( 

181 source="", 

182 datatype=datatype, 

183 shape=column.shape[1:], 

184 ) 

185 ) 

186 return TableColumnModel( 

187 data=data, 

188 name=column.name, 

189 unit=astropy.units.Unit(column.unit) if column.unit is not None else None, 

190 meta=column.meta, 

191 description=column.description or "", 

192 ) 

193 

194 def update_table(self, table: astropy.table.Table) -> None: 

195 """Update the unit and description of an astropy column from this 

196 object. 

197 

198 Parameters 

199 ---------- 

200 table 

201 Astropy table whose matching column is updated in place. 

202 """ 

203 astropy_column: astropy.table.Column = table.columns[self.name] 

204 astropy_column.unit = self.unit 

205 astropy_column.description = self.description 

206 if (datatype := NumberType.from_numpy(astropy_column.dtype)) != self.data.datatype: 

207 raise ArchiveReadError( 

208 f"Table column {self.name} has type {datatype}; expected {self.data.datatype}." 

209 ) 

210 if (shape := astropy_column.shape[1:]) != tuple(self.data.shape): 

211 raise ArchiveReadError( 

212 f"Table column {self.name} has shape {shape}; expected {tuple(self.data.shape)}." 

213 ) 

214 

215 

216class TableModel(pydantic.BaseModel): 

217 """Placeholder for an ASDF-like model for referencing or holding binary 

218 tabular data. 

219 """ 

220 

221 columns: list[TableColumnModel] = pydantic.Field( 

222 default_factory=list, description="Definitions of all columns." 

223 ) 

224 meta: dict[str, int | float | str | bool | None] = pydantic.Field( 

225 default_factory=dict, exclude_if=operator.not_, description="Free-form metadata for the table." 

226 )