Add header_only option for VTF.read()

docs/source/changelog.rst

@@ -11,6 +11,7 @@ Version (dev)
 * Fix casing not being preserved for names of FGD keyvalues during parsing.
 * Fix :py:meth:`PackList.write_soundscript_manifest() <srctools.packlist.PackList.write_soundscript_manifest>`,
 :py:meth:`~srctools.packlist.PackList.write_particles_manifest` and :py:meth:`~srctools.packlist.PackList.write_manifest` trying to write to a closed file.
+* Add `header_only` option for :py:meth:`VTF.read() <srctools.vtf.VTF.read>`, allowing reading only metadata if the image is not required.
 
 -------------
 Version 2.3.4

src/srctools/vtf.py

@@ -3,7 +3,8 @@
 This is designed to be used with the `Python Imaging Library`_ to do the editing of pixels or
 saving/loading standard image files.
 
-To compress to DXT formats, this uses the `libsquish`_ library.
+To compress to DXT formats, this uses the `libsquish`_ library. Currently, 16-bit HDR formats are
+not supported, only metdata can be read.
 
 .. _`Python Imaging Library`: https://pillow.readthedocs.io/en/stable/
 .. _`libsquish`: https://sourceforge.net/projects/libsquish/
@@ -377,7 +378,7 @@ def load(self) -> None:
         if getattr(stream, 'closed', False):
             warnings.warn(
                 'VTF image frame read after stream was closed!\n'
-                'If passing in a stream, close the VTF before closing '
+                'If passing in a stream, load the VTF before closing '
                 'the file.',
                 ResourceWarning,
                 source=stream,
@@ -690,8 +691,13 @@ def __init__(
         self.mipmap_count = mip_count
 
     @classmethod
-    def read(cls: 'Type[VTF]', file: IO[bytes]) -> 'VTF':
-        """Read in a VTF file."""
+    def read(cls: 'Type[VTF]', file: IO[bytes], header_only: bool = False) -> 'VTF':
+        """Read in a VTF file.
+
+        :param file: The file to read from, must be seekable.
+        :param header_only: If set, only read metadata, skip the frames entirely.
+           If accessed the image data will be opaque black.
+        """
         signature = file.read(4)
         if signature != b'VTF\0':
             raise ValueError('Bad file signature!')
@@ -799,12 +805,12 @@ def read(cls: 'Type[VTF]', file: IO[bytes]) -> 'VTF':
         if high_res_offset < 0:
             raise ValueError('Missing main image resource!')
 
-        # We don't implement these high-res formats.
+        # We don't implement these high-res formats, just return metadata.
         if fmt is ImageFormats.RGBA16161616 or fmt is ImageFormats.RGBA16161616F:
-            return vtf
+            header_only = True
 
         vtf._low_res = Frame(low_width, low_height)
-        if low_fmt is not ImageFormats.NONE:
+        if low_fmt is not ImageFormats.NONE and not header_only:
             if low_res_offset < 0:
                 raise ValueError('Missing low-res thumbnail resource!')
             vtf._low_res._fileinfo = (file, low_res_offset, low_fmt)
@@ -819,9 +825,10 @@ def read(cls: 'Type[VTF]', file: IO[bytes]) -> 'VTF':
                         depth_or_cube,
                         data_mipmap,
                     ] = Frame(mip_width, mip_height)
-                    # noinspection PyProtectedMember
-                    frame._fileinfo = (file, high_res_offset, fmt)
-                    high_res_offset += fmt.frame_size(mip_width, mip_height)
+                    if not header_only:
+                        # noinspection PyProtectedMember
+                        frame._fileinfo = (file, high_res_offset, fmt)
+                        high_res_offset += fmt.frame_size(mip_width, mip_height)
         return vtf
 
     def save(
@@ -955,6 +962,7 @@ def __exit__(
         exc_type: Type[BaseException], exc_val: BaseException, exc_tb: types.TracebackType,
     ) -> None:
         """Close the streams if any frames still have them open."""
+        self._low_res._fileinfo = None
         for frame in self._frames.values():
             frame._fileinfo = None
 

tests/test_vtf.py

@@ -106,9 +106,6 @@ def test_load(
 
     These samples were created using VTFEdit Reloaded.
     """
-    if fmt.name in ("ATI1N", "ATI2N"):
-        pytest.xfail("VTFEdit doesn't support these.")
-
     with open(datadir / f"sample_{fmt.name.lower()}.vtf", "rb") as f:
         vtf = VTF.read(f)
         assert vtf.format is fmt
@@ -127,6 +124,22 @@ def test_load(
     )
 
 
+@pytest.mark.parametrize("fmt", FORMATS, ids=lambda fmt: fmt.name.lower())
+def test_load_header_only(
+    fmt: ImageFormats,
+    datadir: Path,
+) -> None:
+    """Test loading only the header."""
+    with open(datadir / f"sample_{fmt.name.lower()}.vtf", "rb") as f:
+        vtf = VTF.read(f, header_only=True)
+        assert vtf.format is fmt
+    # Check loading doesn't need the stream, it's not using them.
+    img = vtf.get().to_PIL()
+    # Check image is 64x64 and fully black.
+    assert img.size == (64, 64)
+    assert img.getextrema() == ((0, 0), (0, 0), (0, 0), (255, 255))
+
+
 @pytest.mark.parametrize("fmt", FORMATS, ids=lambda fmt: fmt.name.lower())
 def test_save_bad_size(
     cy_py_format_funcs: str,