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

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

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 ._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 ColumnDefinitionModel(pydantic.BaseModel): 

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

44 

45 name: str 

46 """Name of the column.""" 

47 

48 datatype: NumberType 

49 """Type of the column.""" 

50 

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

52 """Units of the column.""" 

53 

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

55 """Extended description of the column.""" 

56 

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

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

59 

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

61 """ 

62 

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

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

65 

66 @classmethod 

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

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

69 

70 Parameters 

71 ---------- 

72 dtype 

73 Object convertible to `numpy.dtype`. 

74 """ 

75 dtype = np.dtype(dtype) 

76 result: list[ColumnDefinitionModel] = [] 

77 if dtype.fields is None: 

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

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

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

81 try: 

82 datatype, shape = NumberType.from_numpy_with_shape(field_dtype) 

83 except TypeError: 

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

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

86 return result 

87 

88 def update_from_table(self, table: astropy.table.Table) -> None: 

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

90 table. 

91 """ 

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

93 self.unit = astropy.units.Unit(astropy_column.unit) if astropy_column.unit is not None else None 

94 self.description = astropy_column.description 

95 

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

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

98 object. 

99 """ 

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

101 astropy_column.unit = self.unit 

102 astropy_column.description = self.description 

103 

104 

105class TableModel(pydantic.BaseModel): 

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

107 data. 

108 """ 

109 

110 source: str | int 

111 """Reference to the table data. 

112 

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

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

115 (zero-indexed) to identify the HDU. 

116 """ 

117 

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

119 """Definitions of all columns.""" 

120 

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

122 

123 

124class TableCellReferenceModel(pydantic.BaseModel): 

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

126 

127 model_config = pydantic.ConfigDict(frozen=True) 

128 

129 source: str | int 

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

131 

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

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

134 (zero-indexed) to identify the HDU. 

135 """ 

136 

137 column: str 

138 """Name of the column.""" 

139 

140 row: int 

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

142 

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