Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/PIL/FitsImagePlugin.py: 58%

1# 

2# The Python Imaging Library 

3# $Id$ 

4# 

5# FITS file handling 

6# 

7# Copyright (c) 1998-2003 by Fredrik Lundh 

8# 

9# See the README file for information on usage and redistribution. 

10# 

11from __future__ import annotations 

12 

13__lazy_modules__ = {"gzip", "math", "zlib"} 

14 

15import gzip 

16import math 

17import zlib 

18 

19from . import Image, ImageFile 

20 

21 

22def _accept(prefix: bytes) -> bool: 

23 return prefix.startswith(b"SIMPLE") 

24 

25 

26class FitsImageFile(ImageFile.ImageFile): 

27 format = "FITS" 

28 format_description = "FITS" 

29 

30 def _open(self) -> None: 

31 assert self.fp is not None 

32 

33 headers: dict[bytes, bytes] = {} 

34 header_in_progress = False 

35 decoder_name = "" 

36 while True: 

37 header = self.fp.read(80) 

38 if not header: 

39 msg = "Truncated FITS file" 

40 raise OSError(msg) 

41 keyword = header[:8].strip() 

42 if keyword in (b"SIMPLE", b"XTENSION"): 

43 header_in_progress = True 

44 elif headers and not header_in_progress: 

45 # This is now a data unit 

46 break 

47 elif keyword == b"END": 

48 # Seek to the end of the header unit 

49 self.fp.seek(math.ceil(self.fp.tell() / 2880) * 2880) 

50 if not decoder_name: 

51 decoder_name, offset, args = self._parse_headers(headers) 

52 

53 header_in_progress = False 

54 continue 

55 

56 if decoder_name: 

57 # Keep going to read past the headers 

58 continue 

59 

60 value = header[8:].split(b"/")[0].strip() 

61 if value.startswith(b"="): 

62 value = value[1:].strip() 

63 if not headers and (not _accept(keyword) or value != b"T"): 

64 msg = "Not a FITS file" 

65 raise SyntaxError(msg) 

66 headers[keyword] = value 

67 

68 if not decoder_name: 

69 msg = "No image data" 

70 raise ValueError(msg) 

71 

72 offset += self.fp.tell() - 80 

73 self.tile = [ImageFile._Tile(decoder_name, None, offset, args)] 

74 

75 def _get_size( 

76 self, headers: dict[bytes, bytes], prefix: bytes 

77 ) -> tuple[int, int] | None: 

78 naxis = int(headers[prefix + b"NAXIS"]) 

79 if naxis == 0: 

80 return None 

81 

82 if naxis == 1: 

83 return 1, int(headers[prefix + b"NAXIS1"]) 

84 else: 

85 return int(headers[prefix + b"NAXIS1"]), int(headers[prefix + b"NAXIS2"]) 

86 

87 def _parse_headers( 

88 self, headers: dict[bytes, bytes] 

89 ) -> tuple[str, int, tuple[str | int, ...]]: 

90 prefix = b"" 

91 decoder_name = "raw" 

92 offset = 0 

93 if ( 

94 headers.get(b"XTENSION") == b"'BINTABLE'" 

95 and headers.get(b"ZIMAGE") == b"T" 

96 and headers[b"ZCMPTYPE"] == b"'GZIP_1 '" 

97 ): 

98 no_prefix_size = self._get_size(headers, prefix) or (0, 0) 

99 number_of_bits = int(headers[b"BITPIX"]) 

100 offset = no_prefix_size[0] * no_prefix_size[1] * (number_of_bits // 8) 

101 

102 prefix = b"Z" 

103 decoder_name = "fits_gzip" 

104 

105 size = self._get_size(headers, prefix) 

106 if not size: 

107 return "", 0, () 

108 

109 self._size = size 

110 

111 number_of_bits = int(headers[prefix + b"BITPIX"]) 

112 if number_of_bits == 8: 

113 self._mode = rawmode = "L" 

114 elif number_of_bits == 16: 

115 self._mode = "I;16" 

116 rawmode = "I;16B" 

117 elif number_of_bits == 32: 

118 self._mode = "I" 

119 rawmode = "I;32B" 

120 elif number_of_bits == -32 or (decoder_name == "raw" and number_of_bits == -64): 

121 self._mode = "F" 

122 rawmode = f"F;{number_of_bits * -1}BF" 

123 else: 

124 msg = "Unsupported number of bits" 

125 raise OSError(msg) 

126 

127 args: tuple[str | int, ...] 

128 if decoder_name == "raw": 

129 args = (rawmode, 0, -1) 

130 else: 

131 args = (number_of_bits,) 

132 return decoder_name, offset, args 

133 

134 

135class FitsGzipDecoder(ImageFile.PyDecoder): 

136 _pulls_fd = True 

137 

138 def decode(self, buffer: Image.DecoderInput) -> tuple[int, int]: 

139 assert self.fd is not None 

140 with gzip.open(self.fd) as fp: 

141 try: 

142 value = fp.read(self.state.xsize * self.state.ysize * 4) 

143 except EOFError as e: 

144 msg = "Truncated FITS gzip data" 

145 raise OSError(msg) from e 

146 except zlib.error as e: 

147 msg = "Corrupt FITS gzip data" 

148 raise OSError(msg) from e 

149 

150 rows = [] 

151 offset = 0 

152 number_of_bytes = abs(self.args[0]) // 8 

153 for y in range(self.state.ysize): 

154 row = bytearray() 

155 for x in range(self.state.xsize): 

156 row += value[ 

157 offset + 4 - 1 : offset + (4 - number_of_bytes) - 1 : -1 

158 ] 

159 offset += 4 

160 rows.append(row) 

161 self.set_as_raw(bytes([pixel for row in rows[::-1] for pixel in row])) 

162 return -1, 0 

163 

164 

165# -------------------------------------------------------------------- 

166# Registry 

167 

168Image.register_open(FitsImageFile.format, FitsImageFile, _accept) 

169Image.register_decoder("fits_gzip", FitsGzipDecoder) 

170 

171Image.register_extensions(FitsImageFile.format, [".fit", ".fits"])