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

1""" 

2Blizzard Mipmap Format (.blp) 

3Jerome Leclanche <jerome@leclan.ch> 

4 

5The contents of this file are hereby released in the public domain (CC0) 

6Full text of the CC0 license: 

7 https://creativecommons.org/publicdomain/zero/1.0/ 

8 

9BLP1 files, used mostly in Warcraft III, are not fully supported. 

10All types of BLP2 files used in World of Warcraft are supported. 

11 

12The BLP file structure consists of a header, up to 16 mipmaps of the 

13texture 

14 

15Texture sizes must be powers of two, though the two dimensions do 

16not have to be equal; 512x256 is valid, but 512x200 is not. 

17The first mipmap (mipmap #0) is the full size image; each subsequent 

18mipmap halves both dimensions. The final mipmap should be 1x1. 

19 

20BLP files come in many different flavours: 

21* JPEG-compressed (type == 0) - only supported for BLP1. 

22* RAW images (type == 1, encoding == 1). Each mipmap is stored as an 

23 array of 8-bit values, one per pixel, left to right, top to bottom. 

24 Each value is an index to the palette. 

25* DXT-compressed (type == 1, encoding == 2): 

26- DXT1 compression is used if alpha_encoding == 0. 

27 - An additional alpha bit is used if alpha_depth == 1. 

28 - DXT3 compression is used if alpha_encoding == 1. 

29 - DXT5 compression is used if alpha_encoding == 7. 

30""" 

31 

32from __future__ import annotations 

33 

34__lazy_modules__ = {"io", "struct"} 

35 

36import abc 

37import os 

38import struct 

39from enum import IntEnum 

40from io import BytesIO 

41 

42from . import Image, ImageFile 

43 

44TYPE_CHECKING = False 

45if TYPE_CHECKING: 

46 from typing import IO 

47 

48 

49class Format(IntEnum): 

50 JPEG = 0 

51 

52 

53class Encoding(IntEnum): 

54 UNCOMPRESSED = 1 

55 DXT = 2 

56 UNCOMPRESSED_RAW_BGRA = 3 

57 

58 

59class AlphaEncoding(IntEnum): 

60 DXT1 = 0 

61 DXT3 = 1 

62 DXT5 = 7 

63 

64 

65def unpack_565(i: int) -> tuple[int, int, int]: 

66 return ((i >> 11) & 0x1F) << 3, ((i >> 5) & 0x3F) << 2, (i & 0x1F) << 3 

67 

68 

69def decode_dxt1( 

70 data: bytes, alpha: bool = False 

71) -> tuple[bytearray, bytearray, bytearray, bytearray]: 

72 """ 

73 input: one "row" of data (i.e. will produce 4*width pixels) 

74 """ 

75 

76 blocks = len(data) // 8 # number of blocks in row 

77 ret = (bytearray(), bytearray(), bytearray(), bytearray()) 

78 

79 for block_index in range(blocks): 

80 # Decode next 8-byte block. 

81 idx = block_index * 8 

82 color0, color1, bits = struct.unpack_from("<HHI", data, idx) 

83 

84 r0, g0, b0 = unpack_565(color0) 

85 r1, g1, b1 = unpack_565(color1) 

86 

87 # Decode this block into 4x4 pixels 

88 # Accumulate the results onto our 4 row accumulators 

89 for j in range(4): 

90 for i in range(4): 

91 # get next control op and generate a pixel 

92 

93 control = bits & 3 

94 bits = bits >> 2 

95 

96 a = 0xFF 

97 if control == 0: 

98 r, g, b = r0, g0, b0 

99 elif control == 1: 

100 r, g, b = r1, g1, b1 

101 elif control == 2: 

102 if color0 > color1: 

103 r = (2 * r0 + r1) // 3 

104 g = (2 * g0 + g1) // 3 

105 b = (2 * b0 + b1) // 3 

106 else: 

107 r = (r0 + r1) // 2 

108 g = (g0 + g1) // 2 

109 b = (b0 + b1) // 2 

110 elif control == 3: 

111 if color0 > color1: 

112 r = (2 * r1 + r0) // 3 

113 g = (2 * g1 + g0) // 3 

114 b = (2 * b1 + b0) // 3 

115 else: 

116 r, g, b, a = 0, 0, 0, 0 

117 

118 if alpha: 

119 ret[j].extend([r, g, b, a]) 

120 else: 

121 ret[j].extend([r, g, b]) 

122 

123 return ret 

124 

125 

126def decode_dxt3(data: bytes) -> tuple[bytearray, bytearray, bytearray, bytearray]: 

127 """ 

128 input: one "row" of data (i.e. will produce 4*width pixels) 

129 """ 

130 

131 blocks = len(data) // 16 # number of blocks in row 

132 ret = (bytearray(), bytearray(), bytearray(), bytearray()) 

133 

134 for block_index in range(blocks): 

135 idx = block_index * 16 

136 block = data[idx : idx + 16] 

137 # Decode next 16-byte block. 

138 bits = struct.unpack_from("<8B", block) 

139 color0, color1 = struct.unpack_from("<HH", block, 8) 

140 

141 (code,) = struct.unpack_from("<I", block, 12) 

142 

143 r0, g0, b0 = unpack_565(color0) 

144 r1, g1, b1 = unpack_565(color1) 

145 

146 for j in range(4): 

147 high = False # Do we want the higher bits? 

148 for i in range(4): 

149 alphacode_index = (4 * j + i) // 2 

150 a = bits[alphacode_index] 

151 if high: 

152 high = False 

153 a >>= 4 

154 else: 

155 high = True 

156 a &= 0xF 

157 a *= 17 # We get a value between 0 and 15 

158 

159 color_code = (code >> 2 * (4 * j + i)) & 0x03 

160 

161 if color_code == 0: 

162 r, g, b = r0, g0, b0 

163 elif color_code == 1: 

164 r, g, b = r1, g1, b1 

165 elif color_code == 2: 

166 r = (2 * r0 + r1) // 3 

167 g = (2 * g0 + g1) // 3 

168 b = (2 * b0 + b1) // 3 

169 elif color_code == 3: 

170 r = (2 * r1 + r0) // 3 

171 g = (2 * g1 + g0) // 3 

172 b = (2 * b1 + b0) // 3 

173 

174 ret[j].extend([r, g, b, a]) 

175 

176 return ret 

177 

178 

179def decode_dxt5(data: bytes) -> tuple[bytearray, bytearray, bytearray, bytearray]: 

180 """ 

181 input: one "row" of data (i.e. will produce 4 * width pixels) 

182 """ 

183 

184 blocks = len(data) // 16 # number of blocks in row 

185 ret = (bytearray(), bytearray(), bytearray(), bytearray()) 

186 

187 for block_index in range(blocks): 

188 idx = block_index * 16 

189 block = data[idx : idx + 16] 

190 # Decode next 16-byte block. 

191 a0, a1 = struct.unpack_from("<BB", block) 

192 

193 bits = struct.unpack_from("<6B", block, 2) 

194 alphacode1 = bits[2] | (bits[3] << 8) | (bits[4] << 16) | (bits[5] << 24) 

195 alphacode2 = bits[0] | (bits[1] << 8) 

196 

197 color0, color1 = struct.unpack_from("<HH", block, 8) 

198 

199 (code,) = struct.unpack_from("<I", block, 12) 

200 

201 r0, g0, b0 = unpack_565(color0) 

202 r1, g1, b1 = unpack_565(color1) 

203 

204 for j in range(4): 

205 for i in range(4): 

206 # get next control op and generate a pixel 

207 alphacode_index = 3 * (4 * j + i) 

208 

209 if alphacode_index <= 12: 

210 alphacode = (alphacode2 >> alphacode_index) & 0x07 

211 elif alphacode_index == 15: 

212 alphacode = (alphacode2 >> 15) | ((alphacode1 << 1) & 0x06) 

213 else: # alphacode_index >= 18 and alphacode_index <= 45 

214 alphacode = (alphacode1 >> (alphacode_index - 16)) & 0x07 

215 

216 if alphacode == 0: 

217 a = a0 

218 elif alphacode == 1: 

219 a = a1 

220 elif a0 > a1: 

221 a = ((8 - alphacode) * a0 + (alphacode - 1) * a1) // 7 

222 elif alphacode == 6: 

223 a = 0 

224 elif alphacode == 7: 

225 a = 255 

226 else: 

227 a = ((6 - alphacode) * a0 + (alphacode - 1) * a1) // 5 

228 

229 color_code = (code >> 2 * (4 * j + i)) & 0x03 

230 

231 if color_code == 0: 

232 r, g, b = r0, g0, b0 

233 elif color_code == 1: 

234 r, g, b = r1, g1, b1 

235 elif color_code == 2: 

236 r = (2 * r0 + r1) // 3 

237 g = (2 * g0 + g1) // 3 

238 b = (2 * b0 + b1) // 3 

239 elif color_code == 3: 

240 r = (2 * r1 + r0) // 3 

241 g = (2 * g1 + g0) // 3 

242 b = (2 * b1 + b0) // 3 

243 

244 ret[j].extend([r, g, b, a]) 

245 

246 return ret 

247 

248 

249class BLPFormatError(NotImplementedError): 

250 pass 

251 

252 

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

254 return prefix.startswith((b"BLP1", b"BLP2")) 

255 

256 

257class BlpImageFile(ImageFile.ImageFile): 

258 """ 

259 Blizzard Mipmap Format 

260 """ 

261 

262 format = "BLP" 

263 format_description = "Blizzard Mipmap Format" 

264 

265 def _open(self) -> None: 

266 assert self.fp is not None 

267 self.magic = self.fp.read(4) 

268 if not _accept(self.magic): 

269 msg = f"Bad BLP magic {repr(self.magic)}" 

270 raise BLPFormatError(msg) 

271 

272 compression = struct.unpack("<i", self.fp.read(4))[0] 

273 if self.magic == b"BLP1": 

274 alpha = struct.unpack("<I", self.fp.read(4))[0] != 0 

275 else: 

276 encoding = struct.unpack("<b", self.fp.read(1))[0] 

277 alpha = struct.unpack("<b", self.fp.read(1))[0] != 0 

278 alpha_encoding = struct.unpack("<b", self.fp.read(1))[0] 

279 self.fp.seek(1, os.SEEK_CUR) # mips 

280 

281 self._size = struct.unpack("<II", self.fp.read(8)) 

282 

283 args: tuple[int, int, bool] | tuple[int, int, bool, int] 

284 if self.magic == b"BLP1": 

285 encoding = struct.unpack("<i", self.fp.read(4))[0] 

286 self.fp.seek(4, os.SEEK_CUR) # subtype 

287 

288 args = (compression, encoding, alpha) 

289 offset = 28 

290 else: 

291 args = (compression, encoding, alpha, alpha_encoding) 

292 offset = 20 

293 

294 decoder = self.magic.decode() 

295 

296 self._mode = "RGBA" if alpha else "RGB" 

297 self.tile = [ImageFile._Tile(decoder, None, offset, args)] 

298 

299 

300class _BLPBaseDecoder(abc.ABC, ImageFile.PyDecoder): 

301 _pulls_fd = True 

302 

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

304 try: 

305 self._read_header() 

306 self._load() 

307 except struct.error as e: 

308 msg = "Truncated BLP file" 

309 raise OSError(msg) from e 

310 return -1, 0 

311 

312 @abc.abstractmethod 

313 def _load(self) -> None: 

314 pass 

315 

316 def _read_header(self) -> None: 

317 self._offsets = struct.unpack("<16I", self._safe_read(16 * 4)) 

318 self._lengths = struct.unpack("<16I", self._safe_read(16 * 4)) 

319 

320 def _safe_read(self, length: int) -> bytes: 

321 assert self.fd is not None 

322 return ImageFile._safe_read(self.fd, length) 

323 

324 def _read_palette(self) -> list[tuple[int, int, int, int]]: 

325 ret = [] 

326 for i in range(256): 

327 try: 

328 b, g, r, a = struct.unpack("<4B", self._safe_read(4)) 

329 except struct.error: 

330 break 

331 ret.append((b, g, r, a)) 

332 return ret 

333 

334 def _read_bgra( 

335 self, palette: list[tuple[int, int, int, int]], alpha: bool 

336 ) -> bytearray: 

337 data = bytearray() 

338 _data = BytesIO(self._safe_read(self._lengths[0])) 

339 while True: 

340 try: 

341 (offset,) = struct.unpack("<B", _data.read(1)) 

342 except struct.error: 

343 break 

344 b, g, r, a = palette[offset] 

345 d: tuple[int, ...] = (r, g, b) 

346 if alpha: 

347 d += (a,) 

348 data.extend(d) 

349 return data 

350 

351 

352class BLP1Decoder(_BLPBaseDecoder): 

353 def _load(self) -> None: 

354 _compression, _encoding, alpha = self.args 

355 

356 if _compression == Format.JPEG: 

357 self._decode_jpeg_stream() 

358 

359 elif _compression == 1: 

360 if _encoding in (4, 5): 

361 palette = self._read_palette() 

362 data = self._read_bgra(palette, alpha) 

363 self.set_as_raw(data) 

364 else: 

365 msg = f"Unsupported BLP encoding {repr(_encoding)}" 

366 raise BLPFormatError(msg) 

367 else: 

368 msg = f"Unsupported BLP compression {repr(_encoding)}" 

369 raise BLPFormatError(msg) 

370 

371 def _decode_jpeg_stream(self) -> None: 

372 from .JpegImagePlugin import JpegImageFile 

373 

374 (jpeg_header_size,) = struct.unpack("<I", self._safe_read(4)) 

375 jpeg_header = self._safe_read(jpeg_header_size) 

376 assert self.fd is not None 

377 self._safe_read(self._offsets[0] - self.fd.tell()) # What IS this? 

378 data = self._safe_read(self._lengths[0]) 

379 data = jpeg_header + data 

380 image = JpegImageFile(BytesIO(data)) 

381 Image._decompression_bomb_check(image.size) 

382 if image.mode == "CMYK": 

383 args = image.tile[0].args 

384 assert isinstance(args, tuple) 

385 image.tile = [image.tile[0]._replace(args=(args[0], "CMYK"))] 

386 self.set_as_raw(image.convert("RGB").tobytes(), "BGR") 

387 

388 

389class BLP2Decoder(_BLPBaseDecoder): 

390 def _load(self) -> None: 

391 _compression, _encoding, alpha, _alpha_encoding = self.args 

392 

393 palette = self._read_palette() 

394 

395 assert self.fd is not None 

396 self.fd.seek(self._offsets[0]) 

397 

398 if _compression == 1: 

399 # Uncompressed or DirectX compression 

400 

401 if _encoding == Encoding.UNCOMPRESSED: 

402 data = self._read_bgra(palette, alpha) 

403 

404 elif _encoding == Encoding.DXT: 

405 data = bytearray() 

406 if _alpha_encoding == AlphaEncoding.DXT1: 

407 linesize = (self.state.xsize + 3) // 4 * 8 

408 for yb in range((self.state.ysize + 3) // 4): 

409 for d in decode_dxt1(self._safe_read(linesize), alpha): 

410 data += d 

411 

412 elif _alpha_encoding == AlphaEncoding.DXT3: 

413 linesize = (self.state.xsize + 3) // 4 * 16 

414 for yb in range((self.state.ysize + 3) // 4): 

415 for d in decode_dxt3(self._safe_read(linesize)): 

416 data += d 

417 

418 elif _alpha_encoding == AlphaEncoding.DXT5: 

419 linesize = (self.state.xsize + 3) // 4 * 16 

420 for yb in range((self.state.ysize + 3) // 4): 

421 for d in decode_dxt5(self._safe_read(linesize)): 

422 data += d 

423 else: 

424 msg = f"Unsupported alpha encoding {repr(_alpha_encoding)}" 

425 raise BLPFormatError(msg) 

426 else: 

427 msg = f"Unknown BLP encoding {repr(_encoding)}" 

428 raise BLPFormatError(msg) 

429 

430 else: 

431 msg = f"Unknown BLP compression {repr(_compression)}" 

432 raise BLPFormatError(msg) 

433 

434 self.set_as_raw(data) 

435 

436 

437class BLPEncoder(ImageFile.PyEncoder): 

438 _pushes_fd = True 

439 

440 def _write_palette(self) -> bytes: 

441 data = b"" 

442 assert self.im is not None 

443 palette = self.im.getpalette("RGBA") 

444 for i in range(len(palette) // 4): 

445 r, g, b, a = palette[i * 4 : (i + 1) * 4] 

446 data += struct.pack("<4B", b, g, r, a) 

447 while len(data) < 256 * 4: 

448 data += b"\x00" * 4 

449 return data 

450 

451 def encode(self, bufsize: int) -> tuple[int, int, bytes]: 

452 palette_data = self._write_palette() 

453 

454 offset = 20 + 16 * 4 * 2 + len(palette_data) 

455 data = struct.pack("<16I", offset, *((0,) * 15)) 

456 

457 assert self.im is not None 

458 w, h = self.im.size 

459 data += struct.pack("<16I", w * h, *((0,) * 15)) 

460 

461 data += palette_data 

462 

463 for y in range(h): 

464 for x in range(w): 

465 data += struct.pack("<B", self.im.getpixel((x, y))) 

466 

467 return len(data), 0, data 

468 

469 

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

471 if im.mode != "P": 

472 msg = "Unsupported BLP image mode" 

473 raise ValueError(msg) 

474 

475 magic = b"BLP1" if im.encoderinfo.get("blp_version") == "BLP1" else b"BLP2" 

476 fp.write(magic) 

477 

478 assert im.palette is not None 

479 fp.write(struct.pack("<i", 1)) # Uncompressed or DirectX compression 

480 

481 alpha_depth = 1 if im.palette.mode == "RGBA" else 0 

482 if magic == b"BLP1": 

483 fp.write(struct.pack("<L", alpha_depth)) 

484 else: 

485 fp.write(struct.pack("<b", Encoding.UNCOMPRESSED)) 

486 fp.write(struct.pack("<b", alpha_depth)) 

487 fp.write(struct.pack("<b", 0)) # alpha encoding 

488 fp.write(struct.pack("<b", 0)) # mips 

489 fp.write(struct.pack("<II", *im.size)) 

490 if magic == b"BLP1": 

491 fp.write(struct.pack("<i", 5)) 

492 fp.write(struct.pack("<i", 0)) 

493 

494 ImageFile._save(im, fp, [ImageFile._Tile("BLP", args=im.mode)]) 

495 

496 

497Image.register_open(BlpImageFile.format, BlpImageFile, _accept) 

498Image.register_extension(BlpImageFile.format, ".blp") 

499Image.register_decoder("BLP1", BLP1Decoder) 

500Image.register_decoder("BLP2", BLP2Decoder) 

501 

502Image.register_save(BlpImageFile.format, _save) 

503Image.register_encoder("BLP", BLPEncoder)