Allow saving some PA images as PNG

Author: akx

Tests/test_file_png.py

@@ -893,3 +893,43 @@ def core() -> None:
                 im.load()
 
         self._test_leak(core)
+
+
+@skip_unless_feature("zlib")
+class TestFilePngPA:
+    @pytest.mark.parametrize("palette_rawmode", ("RGB", "RGBA"))
+    def test_save_pa(self, tmp_path: Path, palette_rawmode: str) -> None:
+        test_file = tmp_path / "temp.png"
+        image = Image.new("PA", (64, 64))
+        palette: list[int] = []
+        for color in ((255, 0, 0), (0, 255, 0), (0, 0, 255), (128, 128, 128)):
+            palette += color
+            if palette_rawmode == "RGBA":
+                # This value is ignored for PA mode,
+                # but we'll want to test palette rawmode RGBA too
+                palette.append(42)
+        image.putpalette(palette, palette_rawmode)
+        px = image.load()
+        assert px is not None
+        for y in range(64):
+            for x in range(64):
+                index = (x // 16 + y // 16) % 4
+                px[x, y] = (index, (255, 128, 64, 0)[index])
+        image.save(test_file)
+        with Image.open(test_file) as reloaded:
+            # The image is loaded as P with an RGBA palette
+            # (because there is no PA mode in PNG),
+            # so we need to convert it back to PA for comparison.
+            assert_image_equal(image, reloaded.convert("PA"))
+
+    def test_save_pa_incompatible_alpha(self, tmp_path: Path) -> None:
+        # PA images can only be saved as PNG
+        # if each palette index is used with a single alpha value
+        image = Image.new("PA", (2, 1))
+        image.putpalette([255, 0, 0])
+        px = image.load()
+        assert px is not None
+        px[0, 0] = (0, 255)
+        px[1, 0] = (0, 128)  # same palette index, different alpha - this won't fly
+        with pytest.raises(OSError, match="multiple alpha values"):
+            image.save(tmp_path / "temp.png")

src/PIL/PngImagePlugin.py

@@ -1317,6 +1317,57 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     _save(im, fp, filename, save_all=True)
 
 
+def _pa_to_p(im: Image.Image) -> Image.Image:
+    """
+    Try to fold a `PA` image's alpha band into an RGBA palette.
+
+    :param im: A `PA` mode image.
+    :return: A `P` mode image with an RGBA palette.
+    """
+
+    # PNG has no palette+alpha mode like Pillow's PA mode.
+    # However, palette entries may carry alpha via the tRNS chunk:
+    # > For color type 3 (indexed-color), the tRNS chunk contains
+    # > a series of one-byte alpha values,
+    # > corresponding to entries in the PLTE chunk.
+    # > Each entry indicates that pixels of the corresponding
+    # palette index shall be treated as having the specified alpha value.
+    # - https://www.w3.org/TR/png-3/#11tRNS
+    # This means that in some lucky cases of PA images,
+    # if every palette index is used with exactly a single alpha value,
+    # we can fold the alpha band into an RGBA palette, and save the image
+    # as P, so the RGBA palette gets split back to PLTE + tRNS.
+
+    assert im.mode == "PA"
+    indexes, alpha = im.split()
+    index_alpha: dict[int, int] = {}
+    for index, alpha_value in set(zip(indexes.tobytes(), alpha.tobytes())):
+        if index_alpha.setdefault(index, alpha_value) != alpha_value:
+            msg = (
+                "cannot write mode PA as PNG: "
+                "a palette entry is used with multiple alpha values"
+            )
+            raise OSError(msg)
+
+    # If the palette happened to be RGBA, we deliberately drop its alpha here,
+    # because the alpha we care about will be in the `alpha` band.
+    palette = im.getpalette("RGB") or []
+    colors = max(len(palette) // 3, max(index_alpha, default=-1) + 1, 1)
+    rgba_palette = bytearray()
+    for i in range(colors):
+        # The raster could (unfortunately) be using palette indexes
+        # that don't have entries, hence the fallback to black.
+        rgba_palette += bytes(palette[i * 3 : i * 3 + 3] or (0, 0, 0))
+        # Palette entries not referenced by any pixel stay opaque.
+        rgba_palette.append(index_alpha.get(i, 255))
+
+    indexes.putpalette(rgba_palette, "RGBA")
+    indexes.load()  # sync the palette to the core image for the tRNS check
+    indexes.encoderinfo = im.encoderinfo
+    indexes.info = {k: v for k, v in im.info.items() if k != "transparency"}
+    return indexes
+
+
 def _save(
     im: Image.Image,
     fp: IO[bytes],
@@ -1326,6 +1377,9 @@ def _save(
 ) -> None:
     # save an image to disk (called by the save method)
 
+    if im.mode == "PA" and not save_all:
+        im = _pa_to_p(im)
+
     if save_all:
         default_image = im.encoderinfo.get(
             "default_image", im.info.get("default_image")