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

1from __future__ import annotations 

2 

3from io import BytesIO 

4 

5from . import Image, ImageFile 

6 

7try: 

8 from . import _webp 

9 

10 SUPPORTED = True 

11except ImportError: 

12 SUPPORTED = False 

13 

14TYPE_CHECKING = False 

15if TYPE_CHECKING: 

16 from typing import IO, Any 

17 

18_VP8_MODES_BY_IDENTIFIER = { 

19 b"VP8 ": "RGB", 

20 b"VP8X": "RGBA", 

21 b"VP8L": "RGBA", # lossless 

22} 

23 

24 

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

26 is_riff_file_format = prefix.startswith(b"RIFF") 

27 is_webp_file = prefix[8:12] == b"WEBP" 

28 is_valid_vp8_mode = prefix[12:16] in _VP8_MODES_BY_IDENTIFIER 

29 

30 if is_riff_file_format and is_webp_file and is_valid_vp8_mode: 

31 if not SUPPORTED: 

32 return ( 

33 "image file could not be identified because WEBP support not installed" 

34 ) 

35 return True 

36 return False 

37 

38 

39class WebPImageFile(ImageFile.ImageFile): 

40 format = "WEBP" 

41 format_description = "WebP image" 

42 __loaded = 0 

43 __logical_frame = 0 

44 

45 def _open(self) -> None: 

46 # Use the newer AnimDecoder API to parse the (possibly) animated file, 

47 # and access muxed chunks like ICC/EXIF/XMP. 

48 assert self.fp is not None 

49 self._decoder = _webp.WebPAnimDecoder(self.fp.read()) 

50 

51 # Get info from decoder 

52 self._size, self.info["loop"], bgcolor, self.n_frames, self.rawmode = ( 

53 self._decoder.get_info() 

54 ) 

55 self.info["background"] = ( 

56 (bgcolor >> 16) & 0xFF, # R 

57 (bgcolor >> 8) & 0xFF, # G 

58 bgcolor & 0xFF, # B 

59 (bgcolor >> 24) & 0xFF, # A 

60 ) 

61 self.is_animated = self.n_frames > 1 

62 self._mode = "RGB" if self.rawmode == "RGBX" else self.rawmode 

63 

64 # Attempt to read ICC / EXIF / XMP chunks from file 

65 for key, chunk_name in { 

66 "icc_profile": "ICCP", 

67 "exif": "EXIF", 

68 "xmp": "XMP ", 

69 }.items(): 

70 if value := self._decoder.get_chunk(chunk_name): 

71 self.info[key] = value 

72 

73 # Initialize seek state 

74 self._reset(reset=False) 

75 

76 def _getexif(self) -> dict[int, Any] | None: 

77 if "exif" not in self.info: 

78 return None 

79 return self.getexif()._get_merged_dict() 

80 

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

82 if not self._seek_check(frame): 

83 return 

84 

85 # Set logical frame to requested position 

86 self.__logical_frame = frame 

87 

88 def _reset(self, reset: bool = True) -> None: 

89 if reset: 

90 self._decoder.reset() 

91 self.__physical_frame = 0 

92 self.__loaded = -1 

93 self.__timestamp = 0 

94 

95 def _get_next(self) -> tuple[bytes, int, int]: 

96 # Get next frame 

97 ret = self._decoder.get_next() 

98 self.__physical_frame += 1 

99 

100 # Check if an error occurred 

101 if ret is None: 

102 self._reset() # Reset just to be safe 

103 self.seek(0) 

104 msg = "failed to decode next frame in WebP file" 

105 raise EOFError(msg) 

106 

107 # Compute duration 

108 data, timestamp = ret 

109 duration = timestamp - self.__timestamp 

110 self.__timestamp = timestamp 

111 

112 # libwebp gives frame end, adjust to start of frame 

113 timestamp -= duration 

114 return data, timestamp, duration 

115 

116 def _seek(self, frame: int) -> None: 

117 if self.__physical_frame == frame: 

118 return # Nothing to do 

119 if frame < self.__physical_frame: 

120 self._reset() # Rewind to beginning 

121 while self.__physical_frame < frame: 

122 self._get_next() # Advance to the requested frame 

123 

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

125 if self.__loaded != self.__logical_frame: 

126 self._seek(self.__logical_frame) 

127 

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

129 data, self.info["timestamp"], self.info["duration"] = self._get_next() 

130 self.__loaded = self.__logical_frame 

131 

132 # Set tile 

133 if self.fp and self._exclusive_fp: 

134 self.fp.close() 

135 self.fp = BytesIO(data) 

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

137 

138 return super().load() 

139 

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

141 pass 

142 

143 def tell(self) -> int: 

144 return self.__logical_frame 

145 

146 

147def _convert_frame(im: Image.Image) -> Image.Image: 

148 # Make sure image mode is supported 

149 if im.mode not in ("RGBX", "RGBA", "RGB"): 

150 im = im.convert("RGBA" if im.has_transparency_data else "RGB") 

151 return im 

152 

153 

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

155 encoderinfo = im.encoderinfo.copy() 

156 append_images = list(encoderinfo.get("append_images", [])) 

157 

158 # If total frame count is 1, then save using the legacy API, which 

159 # will preserve non-alpha modes 

160 total = 0 

161 for ims in [im] + append_images: 

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

163 if total == 1: 

164 _save(im, fp, filename) 

165 return 

166 

167 background: int | tuple[int, ...] = (0, 0, 0, 0) 

168 if "background" in encoderinfo: 

169 background = encoderinfo["background"] 

170 elif "background" in im.info: 

171 background = im.info["background"] 

172 if isinstance(background, int): 

173 # GifImagePlugin stores a global color table index in 

174 # info["background"]. So it must be converted to an RGBA value 

175 palette = im.getpalette() 

176 if palette: 

177 r, g, b = palette[background * 3 : (background + 1) * 3] 

178 background = (r, g, b, 255) 

179 else: 

180 background = (background, background, background, 255) 

181 

182 duration = im.encoderinfo.get("duration", im.info.get("duration", 0)) 

183 loop = im.encoderinfo.get("loop", 0) 

184 minimize_size = im.encoderinfo.get("minimize_size", False) 

185 kmin = im.encoderinfo.get("kmin", None) 

186 kmax = im.encoderinfo.get("kmax", None) 

187 allow_mixed = im.encoderinfo.get("allow_mixed", False) 

188 verbose = False 

189 lossless = im.encoderinfo.get("lossless", False) 

190 quality = im.encoderinfo.get("quality", 80) 

191 alpha_quality = im.encoderinfo.get("alpha_quality", 100) 

192 method = im.encoderinfo.get("method", 0) 

193 icc_profile = im.encoderinfo.get("icc_profile") or "" 

194 exif = im.encoderinfo.get("exif", "") 

195 if isinstance(exif, Image.Exif): 

196 exif = exif.tobytes() 

197 xmp = im.encoderinfo.get("xmp", "") 

198 if allow_mixed: 

199 lossless = False 

200 

201 # Sensible keyframe defaults are from gif2webp.c script 

202 if kmin is None: 

203 kmin = 9 if lossless else 3 

204 if kmax is None: 

205 kmax = 17 if lossless else 5 

206 

207 # Validate background color 

208 if ( 

209 not isinstance(background, (list, tuple)) 

210 or len(background) != 4 

211 or not all(0 <= v < 256 for v in background) 

212 ): 

213 msg = f"Background color is not an RGBA tuple clamped to (0-255): {background}" 

214 raise OSError(msg) 

215 

216 # Convert to packed uint 

217 bg_r, bg_g, bg_b, bg_a = background 

218 background = (bg_a << 24) | (bg_r << 16) | (bg_g << 8) | (bg_b << 0) 

219 

220 # Setup the WebP animation encoder 

221 enc = _webp.WebPAnimEncoder( 

222 im.size, 

223 background, 

224 loop, 

225 minimize_size, 

226 kmin, 

227 kmax, 

228 allow_mixed, 

229 verbose, 

230 ) 

231 

232 # Add each frame 

233 frame_idx = 0 

234 timestamp = 0 

235 cur_idx = im.tell() 

236 try: 

237 for ims in [im] + append_images: 

238 # Get number of frames in this image 

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

240 

241 for idx in range(nfr): 

242 ims.seek(idx) 

243 

244 frame = _convert_frame(ims) 

245 

246 # Append the frame to the animation encoder 

247 enc.add( 

248 frame.getim(), 

249 round(timestamp), 

250 lossless, 

251 quality, 

252 alpha_quality, 

253 method, 

254 ) 

255 

256 # Update timestamp and frame index 

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

258 timestamp += duration[frame_idx] 

259 else: 

260 timestamp += duration 

261 frame_idx += 1 

262 

263 finally: 

264 im.seek(cur_idx) 

265 

266 # Force encoder to flush frames 

267 enc.add(None, round(timestamp), lossless, quality, alpha_quality, 0) 

268 

269 # Get the final output from the encoder 

270 data = enc.assemble(icc_profile, exif, xmp) 

271 if data is None: 

272 msg = "cannot write file as WebP (encoder returned None)" 

273 raise OSError(msg) 

274 

275 fp.write(data) 

276 

277 

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

279 lossless = im.encoderinfo.get("lossless", False) 

280 quality = im.encoderinfo.get("quality", 80) 

281 alpha_quality = im.encoderinfo.get("alpha_quality", 100) 

282 icc_profile = im.encoderinfo.get("icc_profile") or "" 

283 exif = im.encoderinfo.get("exif", b"") 

284 if isinstance(exif, Image.Exif): 

285 exif = exif.tobytes() 

286 if exif.startswith(b"Exif\x00\x00"): 

287 exif = exif[6:] 

288 xmp = im.encoderinfo.get("xmp", "") 

289 method = im.encoderinfo.get("method", 4) 

290 exact = 1 if im.encoderinfo.get("exact") else 0 

291 

292 im = _convert_frame(im) 

293 

294 data = _webp.WebPEncode( 

295 im.getim(), 

296 lossless, 

297 float(quality), 

298 float(alpha_quality), 

299 icc_profile, 

300 method, 

301 exact, 

302 exif, 

303 xmp, 

304 ) 

305 if data is None: 

306 msg = "cannot write file as WebP (encoder returned None)" 

307 raise OSError(msg) 

308 

309 fp.write(data) 

310 

311 

312Image.register_open(WebPImageFile.format, WebPImageFile, _accept) 

313if SUPPORTED: 

314 Image.register_save(WebPImageFile.format, _save) 

315 Image.register_save_all(WebPImageFile.format, _save_all) 

316 Image.register_extension(WebPImageFile.format, ".webp") 

317 Image.register_mime(WebPImageFile.format, "image/webp")