From 58dd7adc1d00802c53b4960b5e28b7b8d62109f9 Mon Sep 17 00:00:00 2001 From: Spencer Brown Date: Thu, 12 Sep 2024 18:13:19 +1000 Subject: [PATCH] Add header_only option for VTF.read() --- docs/source/changelog.rst | 1 + src/srctools/vtf.py | 28 ++++++++++++++++++---------- tests/test_vtf.py | 19 ++++++++++++++++--- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 06ef1bad..d484a27b 100644 --- a/docs/source/changelog.rst +++ b/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() `, :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() `, allowing reading only metadata if the image is not required. ------------- Version 2.3.4 diff --git a/src/srctools/vtf.py b/src/srctools/vtf.py index c49c0467..a9e775d8 100644 --- a/src/srctools/vtf.py +++ b/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 diff --git a/tests/test_vtf.py b/tests/test_vtf.py index 4b34568c..2ba78291 100644 --- a/tests/test_vtf.py +++ b/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,