Coverage for tests/test_fits_output_archive.py: 100%

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 

14from collections.abc import Callable 

15from pathlib import Path 

16from typing import Any, ClassVar 

17 

18import astropy.io.fits 

19import astropy.table 

20import numpy as np 

21import pydantic 

22 

23from lsst.images.fits import FitsInputArchive, FitsOutputArchive 

24from lsst.images.serialization import ArchiveTree, InputArchive 

25 

26 

27class _TinyTree(ArchiveTree): 

28 """Minimal concrete ArchiveTree for low-level archive writes.""" 

29 

30 SCHEMA_NAME: ClassVar[str] = "test_fits_output_archive" 

31 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

32 MIN_READ_VERSION: ClassVar[int] = 1 

33 PUBLIC_TYPE: ClassVar[type] = object 

34 

35 def deserialize( 

36 self, archive: InputArchive[Any], **kwargs: Any 

37 ) -> _TinyTree: # pragma: no cover - never invoked 

38 raise NotImplementedError() 

39 

40 

41class _PointerTarget(pydantic.BaseModel): 

42 """A trivial pointer-target model holding an array reference.""" 

43 

44 data: dict[str, Any] | None = None 

45 

46 

47def _write_archive(body: Callable[[FitsOutputArchive], None], tmp_path: Path) -> list[tuple[str, int | None]]: 

48 """Write an archive, applying ``body`` to it, and return the 

49 ``(EXTNAME, EXTVER)`` pairs of the resulting extension HDUs. 

50 """ 

51 filename = tmp_path / "test.fits" 

52 with FitsOutputArchive.open(filename) as archive: 

53 body(archive) 

54 archive.add_tree(_TinyTree()) 

55 with astropy.io.fits.open(filename) as hdu_list: 

56 return [ 

57 (hdu.header["EXTNAME"], hdu.header.get("EXTVER")) 

58 for hdu in hdu_list[1:] 

59 if hdu.header.get("EXTNAME") not in ("JSON", "INDEX") 

60 ] 

61 

62 

63def test_repeated_direct_names_get_increasing_extver(tmp_path: Path) -> None: 

64 """Verify repeated direct names get increasing EXTVER disambiguation.""" 

65 array = np.zeros((2, 2), dtype=np.float32) 

66 sources = [] 

67 

68 def body(archive: FitsOutputArchive) -> None: 

69 sources.append(archive.add_array(array, name="data").source) 

70 sources.append(archive.add_array(array, name="data").source) 

71 

72 keys = _write_archive(body, tmp_path) 

73 assert sources == ["fits:DATA", "fits:DATA,2"] 

74 assert keys == [("DATA", None), ("DATA", 2)] 

75 

76 

77def test_direct_and_pointer_target_names_do_not_collide(tmp_path: Path) -> None: 

78 """Verify a direct name and a pointer target's nested name do not 

79 collide. 

80 """ 

81 # A direct name and a pointer target's nested name (registered with 

82 # a leading slash because the pointer's nested archive is rooted at 

83 # "") already produce distinct EXTNAMEs, so neither needs EXTVER 

84 # disambiguation. 

85 array = np.zeros((2, 2), dtype=np.float32) 

86 sources = [] 

87 

88 def serializer(archive: FitsOutputArchive): 

89 ref = archive.add_array(array, name="data") 

90 sources.append(ref.source) 

91 return _PointerTarget(data=ref.model_dump()) 

92 

93 def body(archive: FitsOutputArchive): 

94 sources.append(archive.add_array(array, name="data").source) 

95 archive.serialize_pointer("psf", serializer, key="psf-key") # type: ignore[arg-type] 

96 

97 keys = _write_archive(body, tmp_path) 

98 assert sources == ["fits:DATA", "fits:/DATA"] 

99 assert keys == [("DATA", None), ("/DATA", None)] 

100 

101 

102def test_table_read_is_native_byte_order(tmp_path: Path) -> None: 

103 """Verify that tables read back from FITS (always big-endian on disk) 

104 are in native byte order, with scaled and logical columns intact. 

105 """ 

106 table = astropy.table.Table( 

107 { 

108 "f": np.array([1.5, 2.5]), 

109 "i": np.array([1, 2], dtype=np.int32), 

110 # Stored as a signed column with TZERO; naive byte swapping of 

111 # the raw storage loses the offset. 

112 "u": np.array([1, 2**31 + 5], dtype=np.uint32), 

113 "b": np.array([True, False]), 

114 "v": np.arange(6, dtype=np.float32).reshape(2, 3), 

115 } 

116 ) 

117 filename = tmp_path / "table.fits" 

118 with FitsOutputArchive.open(filename) as archive: 

119 model = archive.add_table(table, name="t") 

120 archive.add_tree(_TinyTree()) 

121 with FitsInputArchive.open(filename) as archive: 

122 array = archive.get_structured_array(model) 

123 read_table = archive.get_table(model) 

124 assert array.dtype.isnative 

125 for name in table.colnames: 

126 assert read_table[name].dtype.isnative, name 

127 np.testing.assert_array_equal(array[name], table[name], err_msg=name) 

128 np.testing.assert_array_equal(read_table[name], table[name], err_msg=name) 

129 assert array["u"].dtype == np.uint32 

130 assert array["b"].dtype == np.bool_