Merge pull request #5402 from radarhere/dds

hugovk · web-flow · commit d0394d44c2c5 · 2021-06-28T18:50:21.000+03:00
diff --git a/Tests/test_file_dds.py b/Tests/test_file_dds.py
@@ -5,7 +5,7 @@
 
 from PIL import DdsImagePlugin, Image
 
-from .helper import assert_image_equal, assert_image_equal_tofile
+from .helper import assert_image_equal, assert_image_equal_tofile, hopper
 
 TEST_FILE_DXT1 = "Tests/images/dxt1-rgb-4bbp-noalpha_MipMaps-1.dds"
 TEST_FILE_DXT3 = "Tests/images/dxt3-argb-8bbp-explicitalpha_MipMaps-1.dds"
@@ -242,3 +242,27 @@ def test_unimplemented_pixel_format():
     with pytest.raises(NotImplementedError):
         with Image.open("Tests/images/unimplemented_pixel_format.dds"):
             pass
+
+
+def test_save_unsupported_mode(tmp_path):
+    out = str(tmp_path / "temp.dds")
+    im = hopper("HSV")
+    with pytest.raises(OSError):
+        im.save(out)
+
+
+@pytest.mark.parametrize(
+    ("mode", "test_file"),
+    [
+        ("RGB", "Tests/images/hopper.png"),
+        ("RGBA", "Tests/images/pil123rgba.png"),
+    ],
+)
+def test_save(mode, test_file, tmp_path):
+    out = str(tmp_path / "temp.dds")
+    with Image.open(test_file) as im:
+        assert im.mode == mode
+        im.save(out)
+
+        with Image.open(out) as reloaded:
+            assert_image_equal(im, reloaded)
diff --git a/docs/handbook/image-file-formats.rst b/docs/handbook/image-file-formats.rst
@@ -39,6 +39,13 @@ The :py:meth:`~PIL.Image.open` method sets the following
 **compression**
     Set to ``bmp_rle`` if the file is run-length encoded.
 
+DDS
+^^^
+
+DDS is a popular container texture format used in video games and natively supported
+by DirectX. Uncompressed RGB and RGBA can be read, and (since 8.3.0) written. DXT1,
+DXT3 (since 3.4.0) and DXT5 pixel formats can be read, only in ``RGBA`` mode.
+
 DIB
 ^^^
 
@@ -1042,17 +1049,6 @@ is commonly used in fax applications. The DCX decoder can read files containing
 When the file is opened, only the first image is read. You can use
 :py:meth:`~PIL.Image.Image.seek` or :py:mod:`~PIL.ImageSequence` to read other images.
 
-
-DDS
-^^^
-
-DDS is a popular container texture format used in video games and natively
-supported by DirectX.
-Currently, uncompressed RGB data and DXT1, DXT3, and DXT5 pixel formats are
-supported, and only in ``RGBA`` mode.
-
-.. versionadded:: 3.4.0 DXT3
-
 FLI, FLC
 ^^^^^^^^
 
diff --git a/src/PIL/DdsImagePlugin.py b/src/PIL/DdsImagePlugin.py
@@ -14,6 +14,7 @@
 from io import BytesIO
 
 from . import Image, ImageFile
+from ._binary import o32le as o32
 
 # Magic ("DDS ")
 DDS_MAGIC = 0x20534444
@@ -130,8 +131,8 @@ def _open(self):
         fourcc = header.read(4)
         (bitcount,) = struct.unpack("<I", header.read(4))
         masks = struct.unpack("<4I", header.read(16))
-        if pfflags & 0x40:
-            # DDPF_RGB - Texture contains uncompressed RGB data
+        if pfflags & DDPF_RGB:
+            # Texture contains uncompressed RGB data
             masks = {mask: ["R", "G", "B", "A"][i] for i, mask in enumerate(masks)}
             rawmode = ""
             if bitcount == 32:
@@ -201,9 +202,46 @@ def load_seek(self, pos):
         pass
 
 
+def _save(im, fp, filename):
+    if im.mode not in ("RGB", "RGBA"):
+        raise OSError(f"cannot write mode {im.mode} as DDS")
+
+    fp.write(
+        o32(DDS_MAGIC)
+        + o32(124)  # header size
+        + o32(
+            DDSD_CAPS | DDSD_HEIGHT | DDSD_WIDTH | DDSD_PITCH | DDSD_PIXELFORMAT
+        )  # flags
+        + o32(im.height)
+        + o32(im.width)
+        + o32((im.width * (32 if im.mode == "RGBA" else 24) + 7) // 8)  # pitch
+        + o32(0)  # depth
+        + o32(0)  # mipmaps
+        + o32(0) * 11  # reserved
+        + o32(32)  # pfsize
+        + o32(DDS_RGBA if im.mode == "RGBA" else DDPF_RGB)  # pfflags
+        + o32(0)  # fourcc
+        + o32(32 if im.mode == "RGBA" else 24)  # bitcount
+        + o32(0xFF0000)  # rbitmask
+        + o32(0xFF00)  # gbitmask
+        + o32(0xFF)  # bbitmask
+        + o32(0xFF000000 if im.mode == "RGBA" else 0)  # abitmask
+        + o32(DDSCAPS_TEXTURE)  # dwCaps
+        + o32(0)  # dwCaps2
+        + o32(0)  # dwCaps3
+        + o32(0)  # dwCaps4
+        + o32(0)  # dwReserved2
+    )
+    if im.mode == "RGBA":
+        r, g, b, a = im.split()
+        im = Image.merge("RGBA", (a, r, g, b))
+    ImageFile._save(im, fp, [("raw", (0, 0) + im.size, 0, (im.mode[::-1], 0, 1))])
+
+
 def _accept(prefix):
     return prefix[:4] == b"DDS "
 
 
 Image.register_open(DdsImageFile.format, DdsImageFile, _accept)
+Image.register_save(DdsImageFile.format, _save)
 Image.register_extension(DdsImageFile.format, ".dds")
