Only use alpha flag for bcn if alpha exists

Tests/test_file_vtf.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import struct
+from io import BytesIO
 from pathlib import Path
 
 import pytest
 
 from PIL import Image
 from PIL.VtfImagePlugin import (
+    VTFException,
+    VtfImageFile,
     VtfPF,
     _closest_power,
     _get_mipmap_count,
@@ -45,7 +49,7 @@ def test_closest_power(size: int, expected_size: int) -> None:
     ],
 )
 def test_get_mipmap_count(size: tuple[int, int], expected_count: int) -> None:
-    assert _get_mipmap_count(*size) == expected_count
+    assert _get_mipmap_count(size) == expected_count
 
 
 @pytest.mark.parametrize(
@@ -67,7 +71,7 @@ def test_get_mipmap_count(size: tuple[int, int], expected_count: int) -> None:
 def test_get_texture_size(
     pixel_format: VtfPF, size: tuple[int, int], expected_size: int
 ) -> None:
-    assert _get_texture_size(pixel_format, *size) == expected_size
+    assert _get_texture_size(pixel_format, size) == expected_size
 
 
 @pytest.mark.parametrize(
@@ -97,6 +101,27 @@ def test_vtf_read(
             assert_image_equal(converted_e, f)
 
 
+def test_invalid_file() -> None:
+    invalid_file = "Tests/images/flower.jpg"
+
+    with pytest.raises(SyntaxError, match="not a VTF file"):
+        VtfImageFile(invalid_file)
+
+
+def test_vtf_read_unsupported_version() -> None:
+    b = BytesIO(b"VTF\x00" + struct.pack("<2I", 7, 5))
+    with pytest.raises(VTFException, match=r"Unsupported VTF version: \(7, 5\)"):
+        with Image.open(b):
+            pass
+
+
+def test_vtf_read_unsupported_pixel_format() -> None:
+    b = BytesIO(b"VTF\x00" + struct.pack("<2I40xI7x", 7, 2, 1))
+    with pytest.raises(VTFException, match="Unsupported VTF pixel format: 1"):
+        with Image.open(b):
+            pass
+
+
 @pytest.mark.parametrize(
     "pixel_format, file_path, expected_mode, epsilon",
     [
@@ -131,3 +156,17 @@ def test_vtf_save(
                 assert_image_similar(im, expected, epsilon)
             else:
                 assert_image_equal(im, expected)
+
+
+def test_vtf_save_unsupported_mode(tmp_path: Path) -> None:
+    out = tmp_path / "temp.vtf"
+    im = Image.new("HSV", (1, 1))
+    with pytest.raises(OSError, match="cannot write mode HSV as VTF"):
+        im.save(out)
+
+
+def test_vtf_save_unsupported_version(tmp_path: Path) -> None:
+    out = tmp_path / "temp.vtf"
+    im = Image.new("L", (1, 1))
+    with pytest.raises(VTFException, match=r"Unsupported VTF version: \(7, 5\)"):
+        im.save(out, version=(7, 5))

src/PIL/VtfImagePlugin.py

@@ -126,31 +126,25 @@ BLOCK_COMPRESSED = (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA, VtfPF.DXT3, VtfPF.DXT5)
 HEADER_V70 = "2HI2H4x3f4xfIbI2b"
 
 
-def _get_texture_size(pixel_format: int, width: int, height: int) -> int:
+def _get_texture_size(pixel_format: int, size: tuple[int, int]) -> int:
-    if pixel_format in (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA):
+    for factor, pixel_formats in (
-        return width * height // 2
+        (0.5, (VtfPF.DXT1, VtfPF.DXT1_ONEBITALPHA)),
-    elif pixel_format in (
+        (1, (VtfPF.DXT3, VtfPF.DXT5, VtfPF.A8, VtfPF.I8)),
-        VtfPF.DXT3,
+        (2, (VtfPF.UV88, VtfPF.IA88)),
-        VtfPF.DXT5,
+        (3, (VtfPF.RGB888, VtfPF.BGR888)),
-        VtfPF.A8,
+        (4, (VtfPF.RGBA8888,)),
-        VtfPF.I8,
     ):
-        return width * height
+        if pixel_format in pixel_formats:
-    elif pixel_format in (VtfPF.UV88, VtfPF.IA88):
+            return int(size[0] * size[1] * factor)
-        return width * height * 2
-    elif pixel_format in (VtfPF.RGB888, VtfPF.BGR888):
-        return width * height * 3
-    elif pixel_format == VtfPF.RGBA8888:
-        return width * height * 4
     msg = f"Unsupported VTF pixel format: {pixel_format}"
     raise VTFException(msg)
 
 
-def _get_mipmap_count(width: int, height: int) -> int:
+def _get_mipmap_count(size: tuple[int, int]) -> int:
     mip_count = 1
     while True:
-        mip_width = width >> mip_count
+        mip_width = size[0] >> mip_count
-        mip_height = height >> mip_count
+        mip_height = size[1] >> mip_count
         if mip_width == 0 and mip_height == 0:
             return mip_count
         mip_count += 1
@@ -177,7 +171,7 @@ def _write_image(fp: IO[bytes], im: Image.Image, pixel_format: VtfPF) -> None:
 
     codec_name = "bcn" if pixel_format in BLOCK_COMPRESSED else "raw"
     tile = [ImageFile._Tile(codec_name, (0, 0) + im.size, fp.tell(), encoder_args)]
-    ImageFile._save(im, fp, tile, _get_texture_size(pixel_format, *im.size))
+    ImageFile._save(im, fp, tile, _get_texture_size(pixel_format, im.size))
 
 
 def _closest_power(x: int) -> int:
@@ -205,7 +199,6 @@ class VtfImageFile(ImageFile.ImageFile):
                 "<I" + HEADER_V70, self.fp.read(struct.calcsize("<I" + HEADER_V70))
             )
         )
-        self.fp.seek(header.header_size)
 
         pixel_format = header.pixel_format
         if pixel_format in (
@@ -255,16 +248,16 @@ class VtfImageFile(ImageFile.ImageFile):
 
         self._size = (header.width, header.height)
 
-        data_start = self.fp.tell()
+        data_start = header.header_size
         data_start += _get_texture_size(
-            header.low_pixel_format, header.low_width, header.low_height
+            header.low_pixel_format, (header.low_width, header.low_height)
         )
         min_res = 4 if pixel_format in BLOCK_COMPRESSED else 1
         for mip_id in range(header.mipmap_count - 1, 0, -1):
             mip_width = max(header.width >> mip_id, min_res)
             mip_height = max(header.height >> mip_id, min_res)
 
-            data_start += _get_texture_size(pixel_format, mip_width, mip_height)
+            data_start += _get_texture_size(pixel_format, (mip_width, mip_height))
 
         self.tile = [ImageFile._Tile(codec_name, (0, 0) + self.size, data_start, args)]
 
@@ -293,20 +286,14 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
 
     flags = CompiledVtfFlags(0)
 
-    if pixel_format in (
+    if (
-        VtfPF.DXT1,
+        pixel_format in (VtfPF.DXT1, VtfPF.DXT3, VtfPF.DXT5)
-        VtfPF.DXT3,
+        and im.mode in ("RGBA", "LA")
-        VtfPF.DXT5,
+    ) or pixel_format in (VtfPF.RGBA8888, VtfPF.BGRA8888, VtfPF.A8, VtfPF.IA88):
-        VtfPF.RGBA8888,
-        VtfPF.BGRA8888,
-        VtfPF.A8,
-        VtfPF.IA88,
-    ):
         flags |= CompiledVtfFlags.EIGHTBITALPHA
     im = im.resize((_closest_power(im.width), _closest_power(im.height)))
-    width, height = im.size
 
-    mipmap_count = _get_mipmap_count(width, height) if generate_mips else 0
+    mipmap_count = _get_mipmap_count(im.size) if generate_mips else 0
 
     thumb = im.convert("RGBA")
     thumb.thumbnail((16, 16))
@@ -314,8 +301,8 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
 
     header = VTFHeader(
         0,
-        width,
+        im.width,
-        height,
+        im.height,
         flags,
         1,
         0,
@@ -343,15 +330,14 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     header_length = 16 + len(header_bytes)
     if version >= (7, 3):
         header_length += 16  # Resource entries
-    fp.write(b"VTF\x00" + struct.pack("<2II", *version, header_length))
+    fp.write(b"VTF\x00" + struct.pack("<2II", *version, header_length) + header_bytes)
-    fp.write(header_bytes)
 
     if version > (7, 2):
         # Resource entries
         for tag, offset in {
             b"\x01\x00\x00": header_length,  # Low-res
             b"\x30\x00\x00": header_length
-            + _get_texture_size(VtfPF.DXT1, thumb.width, thumb.height),  # High-res
+            + _get_texture_size(VtfPF.DXT1, thumb.size),  # High-res
         }.items():
             fp.write(tag + b"\x00")  # Tag, flags
             fp.write(struct.pack("<I", offset))