From 169025df6c557874473037972dc4615bc38e9661 Mon Sep 17 00:00:00 2001
From: Andrew Murray <radarhere@users.noreply.github.com>
Date: Fri, 25 Feb 2022 16:53:53 +1100
Subject: [PATCH] Added BLP saving

---
 Tests/test_file_blp.py    | 15 ++++++++++-
 src/PIL/BlpImagePlugin.py | 54 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/Tests/test_file_blp.py b/Tests/test_file_blp.py
index c2f8d08cb..0891d4053 100644
--- a/Tests/test_file_blp.py
+++ b/Tests/test_file_blp.py
@@ -2,7 +2,7 @@ import pytest
 
 from PIL import BlpImagePlugin, Image
 
-from .helper import assert_image_equal_tofile
+from .helper import assert_image_equal, assert_image_equal_tofile, hopper
 
 
 def test_load_blp1():
@@ -25,6 +25,19 @@ def test_load_blp2_dxt1a():
         assert_image_equal_tofile(im, "Tests/images/blp/blp2_dxt1a.png")
 
 
+def test_save(tmp_path):
+    im = hopper("P")
+    f = str(tmp_path / "temp.blp")
+    im.save(f)
+
+    with Image.open(f) as reloaded:
+        assert_image_equal(im.convert("RGB"), reloaded)
+
+    im = hopper()
+    with pytest.raises(ValueError):
+        im.save(f)
+
+
 @pytest.mark.parametrize(
     "test_file",
     [
diff --git a/src/PIL/BlpImagePlugin.py b/src/PIL/BlpImagePlugin.py
index 35fd3a771..dfd651867 100644
--- a/src/PIL/BlpImagePlugin.py
+++ b/src/PIL/BlpImagePlugin.py
@@ -267,6 +267,10 @@ class BLPFormatError(NotImplementedError):
     pass
 
 
+def _accept(prefix):
+    return prefix[:4] in (b"BLP1", b"BLP2")
+
+
 class BlpImageFile(ImageFile.ImageFile):
     """
     Blizzard Mipmap Format
@@ -304,7 +308,7 @@ class _BLPBaseDecoder(ImageFile.PyDecoder):
             self._read_blp_header()
             self._load()
         except struct.error as e:
-            raise OSError("Truncated Blp file") from e
+            raise OSError("Truncated BLP file") from e
         return 0, 0
 
     def _safe_read(self, length):
@@ -439,12 +443,54 @@ class BLP2Decoder(_BLPBaseDecoder):
         self.set_as_raw(bytes(data))
 
 
-def _accept(prefix):
-    return prefix[:4] in (b"BLP1", b"BLP2")
+class BLP2Encoder(ImageFile.PyEncoder):
+    _pushes_fd = True
+
+    def _write_palette(self):
+        data = b""
+        palette = self.im.getpalette("RGBA", "RGBA")
+        for i in range(256):
+            r, g, b, a = palette[i * 4 : (i + 1) * 4]
+            data += struct.pack("<4B", b, g, r, a)
+        return data
+
+    def encode(self, bufsize):
+        palette_data = self._write_palette()
+
+        offset = 20 + 16 * 4 * 2 + len(palette_data)
+        data = struct.pack("<16I", offset, *((0,) * 15))
+
+        w, h = self.im.size
+        data += struct.pack("<16I", w * h, *((0,) * 15))
+
+        data += palette_data
+
+        for y in range(h):
+            for x in range(w):
+                data += struct.pack("<B", self.im.getpixel((x, y)))
+
+        return len(data), 0, data
+
+
+def _save(im, fp, filename, save_all=False):
+    if im.mode != "P":
+        raise ValueError("Unsupported BLP image mode")
+
+    fp.write(b"BLP2")
+    fp.write(struct.pack("<i", 1))  # Uncompressed or DirectX compression
+    fp.write(struct.pack("<b", Encoding.UNCOMPRESSED))
+    fp.write(struct.pack("<b", 1 if im.palette.mode == "RGBA" else 0))
+    fp.write(struct.pack("<b", 0))  # alpha encoding
+    fp.write(struct.pack("<b", 0))  # mips
+    fp.write(struct.pack("<II", *im.size))
+
+    ImageFile._save(im, fp, [("BLP2", (0, 0) + im.size, 0, im.mode)])
 
 
 Image.register_open(BlpImageFile.format, BlpImageFile, _accept)
 Image.register_extension(BlpImageFile.format, ".blp")
-
 Image.register_decoder("BLP1", BLP1Decoder)
 Image.register_decoder("BLP2", BLP2Decoder)
+
+Image.register_save(BlpImageFile.format, _save)
+Image.register_encoder("BLP2", BLP2Encoder)