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

16 "TableCellReferenceModel", 

17 "TableReferenceModel", 

18 "UnsupportedTableError", 

19) 

20 

21import operator 

22from typing import TYPE_CHECKING, ClassVar, Literal 

23 

24import astropy.units 

25import numpy as np 

26import numpy.typing as npt 

27import pydantic 

28 

29from ._asdf_utils import Unit 

30from ._common import ArchiveReadError 

31from ._dtypes import NumberType 

32 

33if TYPE_CHECKING: 

34 import astropy.table 

35 

36 

37class UnsupportedTableError(NotImplementedError): 

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

39 are not supported by this library. 

40 """ 

41 

42 

43class ColumnDefinitionModel(pydantic.BaseModel): 

44 """A model that describes a column in a table.""" 

45 

46 name: str 

47 """Name of the column.""" 

48 

49 datatype: NumberType 

50 """Type of the column.""" 

51 

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

53 """Units of the column.""" 

54 

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

56 """Extended description of the column.""" 

57 

58 shape: tuple[int, ...] = pydantic.Field(default=(), exclude_if=operator.not_) 

59 """Shape of a single cell in of this column. 

60 

61 An empty `tuple` is used to represent a scalar column. 

62 """ 

63 

64 is_variable_length: bool = pydantic.Field(default=False, exclude_if=operator.not_) 

65 """Whether this column is a variable-length array.""" 

66 

67 @classmethod 

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

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

70 

71 Parameters 

72 ---------- 

73 dtype 

74 Object convertible to `numpy.dtype`. 

75 """ 

76 dtype = np.dtype(dtype) 

77 result: list[ColumnDefinitionModel] = [] 

78 if dtype.fields is None: 

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

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

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

82 try: 

83 datatype, shape = NumberType.from_numpy_with_shape(field_dtype) 

84 except TypeError: 

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

86 result.append(ColumnDefinitionModel.model_construct(name=name, datatype=datatype, shape=shape)) 

87 return result 

88 

89 @classmethod 

90 def from_table(cls, table: astropy.table.Table) -> list[ColumnDefinitionModel]: 

91 """Extract column definitions from an Astropy table.""" 

92 return [cls.from_column(c) for c in table.columns.values()] 

93 

94 @classmethod 

95 def from_column(cls, column: astropy.table.Column) -> ColumnDefinitionModel: 

96 """Extract a column definition from an Astropy column.""" 

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

98 try: 

99 datatype = NumberType.from_numpy(column.dtype) 

100 except TypeError: 

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

102 return ColumnDefinitionModel( 

103 name=column.name, 

104 datatype=datatype, 

105 shape=column.shape[1:], 

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

107 description=column.description or "", 

108 ) 

109 

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

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

112 object. 

113 """ 

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

115 astropy_column.unit = self.unit 

116 astropy_column.description = self.description 

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

118 raise ArchiveReadError(f"Table column {self.name} has type {datatype}; expected {self.datatype}.") 

119 if (shape := astropy_column.shape[1:]) != self.shape: 

120 raise ArchiveReadError(f"Table column {self.name} has shape {shape}; expected {self.shape}.") 

121 

122 

123class TableReferenceModel(pydantic.BaseModel): 

124 """Placeholder for an ASDF-like model for referencing binary tabular 

125 data. 

126 """ 

127 

128 source: str | int 

129 """Reference to the table data. 

130 

131 This is analogous to the ASDF ``ndarray`` field of the same name, i.e 

132 for a FITS binary table, use "fits:EXTNAME[,EXTVER]" or "fits:INDEX" 

133 (zero-indexed) to identify the HDU. 

134 """ 

135 

136 columns: list[ColumnDefinitionModel] = pydantic.Field(default_factory=list) 

137 """Definitions of all columns.""" 

138 

139 source_is_table: ClassVar[Literal[True]] = True 

140 

141 

142class TableCellReferenceModel(pydantic.BaseModel): 

143 """A model that acts as a pointer to data in a table cell.""" 

144 

145 model_config = pydantic.ConfigDict(frozen=True) 

146 

147 source: str | int 

148 """Identifier for the table as a whole. 

149 

150 This is analogous to the ASDF ``ndarray`` field of the same name, i.e 

151 for a FITS binary table, use "fits:EXTNAME[,EXTVER]" or "fits:INDEX" 

152 (zero-indexed) to identify the HDU. 

153 """ 

154 

155 column: str 

156 """Name of the column.""" 

157 

158 row: int 

159 """Row of the cell (zero-indexed).""" 

160 

161 source_is_table: ClassVar[Literal[True]] = True