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

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 ._common import ArchiveTree 

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(ArchiveTree): 

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 def update_from_table(self, table: astropy.table.Table) -> None: 

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

91 table. 

92 """ 

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

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

95 self.description = astropy_column.description 

96 

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

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

99 object. 

100 """ 

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

102 astropy_column.unit = self.unit 

103 astropy_column.description = self.description 

104 

105 

106class TableModel(ArchiveTree): 

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

108 data. 

109 """ 

110 

111 source: str | int 

112 """Reference to the table data. 

113 

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

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

116 (zero-indexed) to identify the HDU. 

117 """ 

118 

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

120 """Definitions of all columns.""" 

121 

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

123 

124 

125class TableCellReferenceModel(ArchiveTree): 

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

127 

128 model_config = pydantic.ConfigDict(frozen=True) 

129 

130 source: str | int 

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

132 

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

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

135 (zero-indexed) to identify the HDU. 

136 """ 

137 

138 column: str 

139 """Name of the column.""" 

140 

141 row: int 

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

143 

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