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

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

170 statements  

1from __future__ import annotations 

2 

3import os 

4from io import BytesIO 

5from typing import IO 

6 

7from . import ExifTags, Image, ImageFile 

8 

9try: 

10 from . import _avif 

11 

12 SUPPORTED = True 

13except ImportError: 

14 SUPPORTED = False 

15 

16# Decoder options as module globals, until there is a way to pass parameters 

17# to Image.open (see https://github.com/python-pillow/Pillow/issues/569) 

18DECODE_CODEC_CHOICE = "auto" 

19DEFAULT_MAX_THREADS = 0 

20 

21 

22def get_codec_version(codec_name: str) -> str | None: 

23 versions = _avif.codec_versions() 

24 for version in versions.split(", "): 

25 if version.split(" [")[0] == codec_name: 

26 return version.split(":")[-1].split(" ")[0] 

27 return None 

28 

29 

30def _accept(prefix: bytes) -> bool | str: 

31 if prefix[4:8] != b"ftyp": 

32 return False 

33 major_brand = prefix[8:12] 

34 if major_brand in ( 

35 # coding brands 

36 b"avif", 

37 b"avis", 

38 # We accept files with AVIF container brands; we can't yet know if 

39 # the ftyp box has the correct compatible brands, but if it doesn't 

40 # then the plugin will raise a SyntaxError which Pillow will catch 

41 # before moving on to the next plugin that accepts the file. 

42 # 

43 # Also, because this file might not actually be an AVIF file, we 

44 # don't raise an error if AVIF support isn't properly compiled. 

45 b"mif1", 

46 b"msf1", 

47 ): 

48 if not SUPPORTED: 

49 return ( 

50 "image file could not be identified because AVIF support not installed" 

51 ) 

52 return True 

53 return False 

54 

55 

56def _get_default_max_threads() -> int: 

57 if DEFAULT_MAX_THREADS: 

58 return DEFAULT_MAX_THREADS 

59 if hasattr(os, "sched_getaffinity"): 

60 return len(os.sched_getaffinity(0)) 

61 else: 

62 return os.cpu_count() or 1 

63 

64 

65class AvifImageFile(ImageFile.ImageFile): 

66 format = "AVIF" 

67 format_description = "AVIF image" 

68 __frame = -1 

69 

70 def _open(self) -> None: 

71 if not SUPPORTED: 

72 msg = "image file could not be opened because AVIF support not installed" 

73 raise SyntaxError(msg) 

74 

75 if DECODE_CODEC_CHOICE != "auto" and not _avif.decoder_codec_available( 

76 DECODE_CODEC_CHOICE 

77 ): 

78 msg = "Invalid opening codec" 

79 raise ValueError(msg) 

80 

81 assert self.fp is not None 

82 self._decoder = _avif.AvifDecoder( 

83 self.fp.read(), 

84 DECODE_CODEC_CHOICE, 

85 _get_default_max_threads(), 

86 ) 

87 

88 # Get info from decoder 

89 self._size, self.n_frames, self._mode, icc, exif, exif_orientation, xmp = ( 

90 self._decoder.get_info() 

91 ) 

92 self.is_animated = self.n_frames > 1 

93 

94 if icc: 

95 self.info["icc_profile"] = icc 

96 if xmp: 

97 self.info["xmp"] = xmp 

98 

99 if exif_orientation != 1 or exif: 

100 exif_data = Image.Exif() 

101 if exif: 

102 exif_data.load(exif) 

103 original_orientation = exif_data.get(ExifTags.Base.Orientation, 1) 

104 else: 

105 original_orientation = 1 

106 if exif_orientation != original_orientation: 

107 exif_data[ExifTags.Base.Orientation] = exif_orientation 

108 exif = exif_data.tobytes() 

109 if exif: 

110 self.info["exif"] = exif 

111 self.seek(0) 

112 

113 def seek(self, frame: int) -> None: 

114 if not self._seek_check(frame): 

115 return 

116 

117 # Set tile 

118 self.__frame = frame 

119 self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.mode)] 

120 

121 def load(self) -> Image.core.PixelAccess | None: 

122 if self.tile: 

123 # We need to load the image data for this frame 

124 data, timescale, pts_in_timescales, duration_in_timescales = ( 

125 self._decoder.get_frame(self.__frame) 

126 ) 

127 self.info["timestamp"] = round(1000 * (pts_in_timescales / timescale)) 

128 self.info["duration"] = round(1000 * (duration_in_timescales / timescale)) 

129 

130 if self.fp and self._exclusive_fp: 

131 self.fp.close() 

132 self.fp = BytesIO(data) 

133 

134 return super().load() 

135 

136 def load_seek(self, pos: int) -> None: 

137 pass 

138 

139 def tell(self) -> int: 

140 return self.__frame 

141 

142 

143def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: 

144 _save(im, fp, filename, save_all=True) 

145 

146 

147def _save( 

148 im: Image.Image, fp: IO[bytes], filename: str | bytes, save_all: bool = False 

149) -> None: 

150 info = im.encoderinfo.copy() 

151 if save_all: 

152 append_images = list(info.get("append_images", [])) 

153 else: 

154 append_images = [] 

155 

156 total = 0 

157 for ims in [im] + append_images: 

158 total += getattr(ims, "n_frames", 1) 

159 

160 quality = info.get("quality", 75) 

161 if not isinstance(quality, int) or quality < 0 or quality > 100: 

162 msg = "Invalid quality setting" 

163 raise ValueError(msg) 

164 

165 duration = info.get("duration", 0) 

166 subsampling = info.get("subsampling", "4:2:0") 

167 speed = info.get("speed", 6) 

168 max_threads = info.get("max_threads", _get_default_max_threads()) 

169 codec = info.get("codec", "auto") 

170 if codec != "auto" and not _avif.encoder_codec_available(codec): 

171 msg = "Invalid saving codec" 

172 raise ValueError(msg) 

173 range_ = info.get("range", "full") 

174 tile_rows_log2 = info.get("tile_rows", 0) 

175 tile_cols_log2 = info.get("tile_cols", 0) 

176 alpha_premultiplied = bool(info.get("alpha_premultiplied", False)) 

177 autotiling = bool(info.get("autotiling", tile_rows_log2 == tile_cols_log2 == 0)) 

178 

179 icc_profile = info.get("icc_profile", im.info.get("icc_profile")) 

180 exif_orientation = 1 

181 if exif := info.get("exif"): 

182 if isinstance(exif, Image.Exif): 

183 exif_data = exif 

184 else: 

185 exif_data = Image.Exif() 

186 exif_data.load(exif) 

187 if ExifTags.Base.Orientation in exif_data: 

188 exif_orientation = exif_data.pop(ExifTags.Base.Orientation) 

189 exif = exif_data.tobytes() if exif_data else b"" 

190 elif isinstance(exif, Image.Exif): 

191 exif = exif_data.tobytes() 

192 

193 xmp = info.get("xmp") 

194 

195 if isinstance(xmp, str): 

196 xmp = xmp.encode("utf-8") 

197 

198 advanced = info.get("advanced") 

199 if advanced is not None: 

200 if isinstance(advanced, dict): 

201 advanced = advanced.items() 

202 try: 

203 advanced = tuple(advanced) 

204 except TypeError: 

205 invalid = True 

206 else: 

207 invalid = any(not isinstance(v, tuple) or len(v) != 2 for v in advanced) 

208 if invalid: 

209 msg = ( 

210 "advanced codec options must be a dict of key-value string " 

211 "pairs or a series of key-value two-tuples" 

212 ) 

213 raise ValueError(msg) 

214 

215 # Setup the AVIF encoder 

216 enc = _avif.AvifEncoder( 

217 im.size, 

218 subsampling, 

219 quality, 

220 speed, 

221 max_threads, 

222 codec, 

223 range_, 

224 tile_rows_log2, 

225 tile_cols_log2, 

226 alpha_premultiplied, 

227 autotiling, 

228 icc_profile or b"", 

229 exif or b"", 

230 exif_orientation, 

231 xmp or b"", 

232 advanced, 

233 ) 

234 

235 # Add each frame 

236 frame_idx = 0 

237 frame_duration = 0 

238 cur_idx = im.tell() 

239 is_single_frame = total == 1 

240 try: 

241 for ims in [im] + append_images: 

242 # Get number of frames in this image 

243 nfr = getattr(ims, "n_frames", 1) 

244 

245 for idx in range(nfr): 

246 ims.seek(idx) 

247 

248 # Make sure image mode is supported 

249 frame = ims 

250 rawmode = ims.mode 

251 if ims.mode not in {"RGB", "RGBA"}: 

252 rawmode = "RGBA" if ims.has_transparency_data else "RGB" 

253 frame = ims.convert(rawmode) 

254 

255 # Update frame duration 

256 if isinstance(duration, (list, tuple)): 

257 frame_duration = duration[frame_idx] 

258 else: 

259 frame_duration = duration 

260 

261 # Append the frame to the animation encoder 

262 enc.add( 

263 frame.tobytes("raw", rawmode), 

264 frame_duration, 

265 frame.size, 

266 rawmode, 

267 is_single_frame, 

268 ) 

269 

270 # Update frame index 

271 frame_idx += 1 

272 

273 if not save_all: 

274 break 

275 

276 finally: 

277 im.seek(cur_idx) 

278 

279 # Get the final output from the encoder 

280 data = enc.finish() 

281 if data is None: 

282 msg = "cannot write file as AVIF (encoder returned None)" 

283 raise OSError(msg) 

284 

285 fp.write(data) 

286 

287 

288Image.register_open(AvifImageFile.format, AvifImageFile, _accept) 

289if SUPPORTED: 

290 Image.register_save(AvifImageFile.format, _save) 

291 Image.register_save_all(AvifImageFile.format, _save_all) 

292 Image.register_extensions(AvifImageFile.format, [".avif", ".avifs"]) 

293 Image.register_mime(AvifImageFile.format, "image/avif")