Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/PIL/WebPImagePlugin.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

206 statements  

1from __future__ import annotations 

2 

3__lazy_modules__ = {"io"} 

4 

5 

6from . import Image, ImageFile 

7from ._binary import i16le as i16 

8from ._binary import i32le as i32 

9 

10try: 

11 from . import _webp 

12 

13 SUPPORTED = True 

14except ImportError: 

15 SUPPORTED = False 

16 

17TYPE_CHECKING = False 

18if TYPE_CHECKING: 

19 from typing import IO, Any 

20 

21 

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

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

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

25 is_valid_vp8_mode = prefix[12:16] in {b"VP8 ", b"VP8X", b"VP8L"} 

26 

27 if is_riff_file_format and is_webp_file and is_valid_vp8_mode: 

28 if not SUPPORTED: 

29 return ( 

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

31 ) 

32 return True 

33 return False 

34 

35 

36def _get_size(s: bytes) -> tuple[int, int] | None: 

37 # Read the canvas size from the file header, without decoding 

38 chunk = s[12:16] 

39 if chunk == b"VP8 ": 

40 if len(s) >= 30 and s[23:26] == b"\x9d\x01\x2a": 

41 return i16(s, 26) & 0x3FFF, i16(s, 28) & 0x3FFF 

42 elif chunk == b"VP8X": 

43 if len(s) >= 30: 

44 return i32(s[24:27] + b"\0") + 1, i32(s[27:30] + b"\0") + 1 

45 elif chunk == b"VP8L": 

46 if len(s) >= 25 and s[20] == 0x2F: 

47 bits = i32(s, 21) 

48 return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1 

49 return None 

50 

51 

52class WebPImageFile(ImageFile.ImageFile): 

53 format = "WEBP" 

54 format_description = "WebP image" 

55 __loaded = 0 

56 __logical_frame = 0 

57 

58 def _open(self) -> None: 

59 assert self.fp is not None 

60 s = self.fp.read(30) 

61 if not _accept(s): 

62 msg = "not a WEBP file" 

63 raise SyntaxError(msg) 

64 

65 # The decoder allocates memory for the canvas when it is created, 

66 # so check the size before creating it. 

67 if size := _get_size(s): 

68 # Image.open() runs this check afterwards, 

69 # so only raise errors here, to avoid warning twice 

70 Image._decompression_bomb_check(size, error_only=True) 

71 s += self.fp.read() 

72 

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

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

75 self._decoder = _webp.WebPAnimDecoder(s) 

76 

77 # Get info from decoder 

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

79 self._decoder.get_info() 

80 ) 

81 self.info["background"] = ( 

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

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

84 bgcolor & 0xFF, # B 

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

86 ) 

87 self.is_animated = self.n_frames > 1 

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

89 

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

91 for key, chunk_name in { 

92 "icc_profile": "ICCP", 

93 "exif": "EXIF", 

94 "xmp": "XMP ", 

95 }.items(): 

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

97 self.info[key] = value 

98 

99 # Initialize seek state 

100 self._reset(reset=False) 

101 

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

103 if "exif" not in self.info: 

104 return None 

105 return self.getexif()._get_merged_dict() 

106 

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

108 if not self._seek_check(frame): 

109 return 

110 self.info.pop("timestamp", None) 

111 self.info.pop("duration", None) 

112 

113 # Set logical frame to requested position 

114 self.__logical_frame = frame 

115 

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

117 if reset: 

118 self._decoder.reset() 

119 self.__physical_frame = 0 

120 self.__loaded = -1 

121 self.__timestamp = 0 

122 

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

124 # Get next frame 

125 ret = self._decoder.get_next() 

126 self.__physical_frame += 1 

127 

128 # Check if an error occurred 

129 if ret is None: 

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

131 self.seek(0) 

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

133 raise EOFError(msg) 

134 

135 # Compute duration 

136 data, timestamp = ret 

137 duration = timestamp - self.__timestamp 

138 self.__timestamp = timestamp 

139 

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

141 timestamp -= duration 

142 return data, timestamp, duration 

143 

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

145 if self.__physical_frame == frame: 

146 return # Nothing to do 

147 if frame < self.__physical_frame: 

148 self._reset() # Rewind to beginning 

149 while self.__physical_frame < frame: 

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

151 

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

153 if self.__loaded != self.__logical_frame: 

154 self._seek(self.__logical_frame) 

155 

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

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

158 self.__loaded = self.__logical_frame 

159 

160 if self.fp and self._exclusive_fp: 

161 self.fp.close() 

162 self.fp = None 

163 self.frombytes(data, "raw", self.rawmode) 

164 

165 return super().load() 

166 

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

168 pass 

169 

170 def tell(self) -> int: 

171 return self.__logical_frame 

172 

173 

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

175 # Make sure image mode is supported 

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

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

178 return im 

179 

180 

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

182 encoderinfo = im.encoderinfo.copy() 

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

184 

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

186 # will preserve non-alpha modes 

187 total = 0 

188 for ims in [im] + append_images: 

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

190 if total == 1: 

191 _save(im, fp, filename) 

192 return 

193 

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

195 if "background" in encoderinfo: 

196 background = encoderinfo["background"] 

197 elif "background" in im.info: 

198 background = im.info["background"] 

199 if isinstance(background, int): 

200 # GifImagePlugin stores a global color table index in 

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

202 palette = im.getpalette() 

203 if palette: 

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

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

206 else: 

207 background = (background, background, background, 255) 

208 

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

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

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

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

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

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

215 verbose = False 

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

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

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

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

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

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

222 if isinstance(exif, Image.Exif): 

223 exif = exif.tobytes() 

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

225 if allow_mixed: 

226 lossless = False 

227 

228 # Sensible keyframe defaults are from gif2webp.c script 

229 if kmin is None: 

230 kmin = 9 if lossless else 3 

231 if kmax is None: 

232 kmax = 17 if lossless else 5 

233 

234 # Validate background color 

235 if ( 

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

237 or len(background) != 4 

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

239 ): 

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

241 raise OSError(msg) 

242 

243 # Convert to packed uint 

244 bg_r, bg_g, bg_b, bg_a = background 

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

246 

247 # Setup the WebP animation encoder 

248 enc = _webp.WebPAnimEncoder( 

249 im.size, 

250 background, 

251 loop, 

252 minimize_size, 

253 kmin, 

254 kmax, 

255 allow_mixed, 

256 verbose, 

257 ) 

258 

259 # Add each frame 

260 frame_idx = 0 

261 timestamp = 0 

262 cur_idx = im.tell() 

263 try: 

264 for ims in [im] + append_images: 

265 # Get number of frames in this image 

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

267 

268 for idx in range(nfr): 

269 ims.seek(idx) 

270 

271 frame = _convert_frame(ims) 

272 

273 # Append the frame to the animation encoder 

274 enc.add( 

275 frame.getim(), 

276 round(timestamp), 

277 lossless, 

278 quality, 

279 alpha_quality, 

280 method, 

281 ) 

282 

283 # Update timestamp and frame index 

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

285 timestamp += duration[frame_idx] 

286 else: 

287 timestamp += duration 

288 frame_idx += 1 

289 

290 finally: 

291 im.seek(cur_idx) 

292 

293 # Force encoder to flush frames 

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

295 

296 # Get the final output from the encoder 

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

298 if data is None: 

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

300 raise OSError(msg) 

301 

302 fp.write(data) 

303 

304 

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

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

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

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

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

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

311 if isinstance(exif, Image.Exif): 

312 exif = exif.tobytes() 

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

314 exif = exif[6:] 

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

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

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

318 

319 im = _convert_frame(im) 

320 

321 data = _webp.WebPEncode( 

322 im.getim(), 

323 lossless, 

324 float(quality), 

325 float(alpha_quality), 

326 icc_profile, 

327 method, 

328 exact, 

329 exif, 

330 xmp, 

331 ) 

332 if data is None: 

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

334 raise OSError(msg) 

335 

336 fp.write(data) 

337 

338 

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

340if SUPPORTED: 

341 Image.register_save(WebPImageFile.format, _save) 

342 Image.register_save_all(WebPImageFile.format, _save_all) 

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

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