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

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: 

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 return [cls.from_column(c, inline=inline) for c in table.columns.values()] 

143 

144 @classmethod 

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

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

147 column. 

148 

149 Notes 

150 ----- 

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

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

153 will need to be modified later. 

154 """ 

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

156 try: 

157 datatype = NumberType.from_numpy(column.dtype) 

158 except TypeError: 

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

160 

161 data = ( 

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

163 if inline 

164 else ArrayReferenceModel( 

165 source="", 

166 datatype=datatype, 

167 shape=column.shape[1:], 

168 ) 

169 ) 

170 return TableColumnModel( 

171 data=data, 

172 name=column.name, 

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

174 meta=column.meta, 

175 description=column.description or "", 

176 ) 

177 

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

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

180 object. 

181 """ 

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

183 astropy_column.unit = self.unit 

184 astropy_column.description = self.description 

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

186 raise ArchiveReadError( 

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

188 ) 

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

190 raise ArchiveReadError( 

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

192 ) 

193 

194 

195class TableModel(pydantic.BaseModel): 

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

197 tabular data. 

198 """ 

199 

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

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

202 ) 

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

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

205 )