Improve SgiImagePlugin test coverage (#8896)

Author: radarhere

Tests/test_file_sgi.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from io import BytesIO
 from pathlib import Path
 
 import pytest
@@ -71,6 +72,15 @@ def test_invalid_file() -> None:
         SgiImagePlugin.SgiImageFile(invalid_file)
 
 
+def test_unsupported_image_mode() -> None:
+    with open("Tests/images/hopper.rgb", "rb") as fp:
+        data = fp.read()
+    data = data[:3] + b"\x03" + data[4:]
+    with pytest.raises(ValueError, match="Unsupported SGI image mode"):
+        with Image.open(BytesIO(data)):
+            pass
+
+
 def roundtrip(img: Image.Image, tmp_path: Path) -> None:
     out = tmp_path / "temp.sgi"
     img.save(out, format="sgi")
@@ -109,3 +119,11 @@ def test_unsupported_mode(tmp_path: Path) -> None:
 
     with pytest.raises(ValueError):
         im.save(out, format="sgi")
+
+
+def test_unsupported_number_of_bytes_per_pixel(tmp_path: Path) -> None:
+    im = hopper()
+    out = tmp_path / "temp.sgi"
+
+    with pytest.raises(ValueError, match="Unsupported number of bytes per pixel"):
+        im.save(out, bpc=3)

src/PIL/SgiImagePlugin.py

@@ -82,17 +82,10 @@ def _open(self) -> None:
         # zsize : channels count
         zsize = i16(s, 10)
 
-        # layout
-        layout = bpc, dimension, zsize
-
         # determine mode from bits/zsize
-        rawmode = ""
         try:
-            rawmode = MODES[layout]
+            rawmode = MODES[(bpc, dimension, zsize)]
         except KeyError:
-            pass
-
-        if rawmode == "":
             msg = "Unsupported SGI image mode"
             raise ValueError(msg)
 
@@ -156,24 +149,15 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     # Run-Length Encoding Compression - Unsupported at this time
     rle = 0
 
-    # Number of dimensions (x,y,z)
-    dim = 3
     # X Dimension = width / Y Dimension = height
     x, y = im.size
-    if im.mode == "L" and y == 1:
-        dim = 1
-    elif im.mode == "L":
-        dim = 2
     # Z Dimension: Number of channels
     z = len(im.mode)
-
-    if dim in {1, 2}:
-        z = 1
-
-    # assert we've got the right number of bands.
-    if len(im.getbands()) != z:
-        msg = f"incorrect number of bands in SGI write: {z} vs {len(im.getbands())}"
-        raise ValueError(msg)
+    # Number of dimensions (x,y,z)
+    if im.mode == "L":
+        dimension = 1 if y == 1 else 2
+    else:
+        dimension = 3
 
     # Minimum Byte value
     pinmin = 0
@@ -188,7 +172,7 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     fp.write(struct.pack(">h", magic_number))
     fp.write(o8(rle))
     fp.write(o8(bpc))
-    fp.write(struct.pack(">H", dim))
+    fp.write(struct.pack(">H", dimension))
     fp.write(struct.pack(">H", x))
     fp.write(struct.pack(">H", y))
     fp.write(struct.pack(">H", z))