From bed079a29e8837ac9e0d417a373c3365c3559b6f Mon Sep 17 00:00:00 2001 From: David Hotham Date: Sat, 4 Apr 2026 12:15:13 +0100 Subject: [PATCH] Add trust_wheel_record to install() So that callers who already have validate the RECORD do not need to repeat expensive hash calculations. --- src/installer/_core.py | 39 +++++- src/installer/destinations.py | 52 ++++++++ tests/test_core.py | 222 ++++++++++++++++++++++++++++++++++ tests/test_destinations.py | 41 +++++++ 4 files changed, 348 insertions(+), 6 deletions(-) diff --git a/src/installer/_core.py b/src/installer/_core.py index e23f06e7..a38fc807 100644 --- a/src/installer/_core.py +++ b/src/installer/_core.py @@ -67,6 +67,8 @@ def install( source: WheelSource, destination: WheelDestination, additional_metadata: dict[str, bytes], + *, + trust_wheel_record: bool = False, ) -> None: """Install wheel described by ``source`` into ``destination``. @@ -74,6 +76,12 @@ def install( :param destination: where to write the wheel. :param additional_metadata: additional metadata files to generate, usually generated by the caller. + :param trust_wheel_record: if ``True``, reuse the hash and size from the + wheel's ``RECORD`` file instead of recomputing them for every + installed file. This avoids expensive hash calculations for + large wheels when the caller has already validated the wheel + (e.g. via :py:meth:`~installer.sources.WheelSource.validate_record`). + Generated scripts and additional metadata are always hashed. """ root_scheme = _process_WHEEL_file(source) @@ -121,12 +129,31 @@ def install( source=source, root_scheme=root_scheme, ) - record = destination.write_file( - scheme=scheme, - path=destination_path, - stream=stream, - is_executable=is_executable, - ) + + # When trusting the wheel's RECORD we can skip hash computation for + # files that already have a hash recorded. Scripts-scheme files are + # excluded because their content may be modified by shebang rewriting. + if ( + trust_wheel_record + and source_record.hash_ is not None + and scheme != "scripts" + ): + destination.write_file_no_record( + scheme=scheme, + path=destination_path, + stream=stream, + is_executable=is_executable, + ) + record = RecordEntry( + destination_path, source_record.hash_, source_record.size + ) + else: + record = destination.write_file( + scheme=scheme, + path=destination_path, + stream=stream, + is_executable=is_executable, + ) written_records.append((scheme, record)) # Write all the installation-specific metadata diff --git a/src/installer/destinations.py b/src/installer/destinations.py index 71a8f9f4..494d27ab 100644 --- a/src/installer/destinations.py +++ b/src/installer/destinations.py @@ -2,6 +2,7 @@ import io import os +import shutil from collections.abc import Collection, Iterable from dataclasses import dataclass from pathlib import Path @@ -77,6 +78,26 @@ def write_file( """ raise NotImplementedError + def write_file_no_record( + self, + scheme: Scheme, + path: Union[str, "os.PathLike[str]"], + stream: BinaryIO, + is_executable: bool, + ) -> None: + """Write a file to correct ``path`` within the ``scheme``, without recording. + + Like :py:meth:`write_file`, but does not need to return a + :py:class:`~installer.records.RecordEntry`. Implementations may skip + hash computation for better performance. + + :param scheme: scheme to write the file in (like "purelib", "platlib" etc). + :param path: path within that scheme + :param stream: contents of the file + :param is_executable: whether the file should be made executable + """ + raise NotImplementedError + def finalize_installation( self, scheme: Scheme, @@ -212,6 +233,37 @@ def write_file( return self.write_to_fs(scheme, path_, stream, is_executable) + def write_file_no_record( + self, + scheme: Scheme, + path: Union[str, "os.PathLike[str]"], + stream: BinaryIO, + is_executable: bool, + ) -> None: + """Write a file without computing a hash or returning a record entry. + + :param scheme: scheme to write the file in (like "purelib", "platlib" etc). + :param path: path within that scheme + :param stream: contents of the file + :param is_executable: whether the file should be made executable + + - Ensures that an existing file is not being overwritten. + """ + target_path = self._path_with_destdir(scheme, os.fspath(path)) + if not self.overwrite_existing and target_path.exists(): + message = f"File already exists: {target_path!s}" + raise FileExistsError(message) + + parent_folder = target_path.parent + if not parent_folder.exists(): + parent_folder.mkdir(parents=True) + + with target_path.open("wb") as f: + shutil.copyfileobj(stream, f) + + if is_executable: + make_file_executable(target_path) + def write_script( self, name: str, module: str, attr: str, section: "ScriptSection" ) -> RecordEntry: diff --git a/tests/test_core.py b/tests/test_core.py index 121b9d3e..d1f5922e 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -987,3 +987,225 @@ def test_skips_pycache_and_warns(self, mock_destination): assert sub_good_path in record_paths assert top_pycache_path not in record_paths assert sub_pycache_path not in record_paths + + def test_trust_wheel_record_uses_source_hashes(self, mock_destination): + """When trust_wheel_record=True, records from the wheel's RECORD + should be used instead of whatever write_file returns. + """ + source = FakeWheelSource( + distribution="fancy", + version="1.0.0", + regular_files={ + "fancy/__init__.py": b"""\ + def main(): + print("I'm a fancy package") + """, + "fancy/__main__.py": b"""\ + if __name__ == "__main__": + from . import main + main() + """, + }, + dist_info_files={ + "top_level.txt": b"""\ + fancy + """, + "WHEEL": b"""\ + Wheel-Version: 1.0 + Generator: magic (1.0.0) + Root-Is-Purelib: true + Tag: py3-none-any + """, + "METADATA": b"""\ + Metadata-Version: 2.1 + Name: fancy + Version: 1.0.0 + """, + }, + ) + + install( + source=source, + destination=mock_destination, + additional_metadata={}, + trust_wheel_record=True, + ) + + # All wheel files (non-scripts scheme) should go through + # write_file_no_record instead of write_file. + no_record_calls = mock_destination.write_file_no_record.call_args_list + write_file_calls = mock_destination.write_file.call_args_list + no_record_paths = {c.kwargs["path"] for c in no_record_calls} + + # Regular files and dist-info files (except RECORD) should all use + # write_file_no_record. + assert "fancy/__init__.py" in no_record_paths + assert "fancy/__main__.py" in no_record_paths + + # write_file should NOT have been called for files that went through + # write_file_no_record. + write_file_paths = {c.kwargs["path"] for c in write_file_calls} + assert "fancy/__init__.py" not in write_file_paths + assert "fancy/__main__.py" not in write_file_paths + + # finalize_installation should receive RecordEntry objects with the + # source hashes (not the mock return values from write_file). + records = mock_destination.finalize_installation.call_args[1]["records"] + for _scheme, rec in records: + if isinstance(rec, RecordEntry) and rec.hash_ is not None: + assert rec.hash_.name == "sha256" + assert rec.hash_.value != "" + + def test_trust_wheel_record_falls_back_for_missing_hash(self, mock_destination): + """When trust_wheel_record=True but a record has no hash, + write_file should be used as a fallback. + """ + source = FakeWheelSource( + distribution="fancy", + version="1.0.0", + regular_files={ + "fancy/__init__.py": b"""\ + def main(): + print("I'm a fancy package") + """, + }, + dist_info_files={ + "top_level.txt": b"""\ + fancy + """, + "WHEEL": b"""\ + Wheel-Version: 1.0 + Generator: magic (1.0.0) + Root-Is-Purelib: true + Tag: py3-none-any + """, + "METADATA": b"""\ + Metadata-Version: 2.1 + Name: fancy + Version: 1.0.0 + """, + }, + ) + + # Patch get_contents to strip hash info from the first file. + original_get_contents = source.get_contents + + def patched_get_contents(): + for record, stream, is_exec in original_get_contents(): + path = record[0] + if path == "fancy/__init__.py": + yield (path, "", ""), stream, is_exec + else: + yield record, stream, is_exec + + source.get_contents = patched_get_contents + + install( + source=source, + destination=mock_destination, + additional_metadata={}, + trust_wheel_record=True, + ) + + # fancy/__init__.py has no hash, so it should go through write_file. + write_file_paths = { + c.kwargs["path"] for c in mock_destination.write_file.call_args_list + } + assert "fancy/__init__.py" in write_file_paths + + def test_trust_wheel_record_still_hashes_scripts(self, mock_destination): + """Scripts-scheme files should always be hashed (shebang rewriting + may change their content). + """ + source = FakeWheelSource( + distribution="fancy", + version="1.0.0", + regular_files={ + "fancy/__init__.py": b"""\ + def main(): + print("I'm a fancy package") + """, + "fancy-1.0.0.data/scripts/run_fancy": b"""\ + #!/usr/bin/env python3 + import fancy; fancy.main() + """, + }, + dist_info_files={ + "top_level.txt": b"""\ + fancy + """, + "WHEEL": b"""\ + Wheel-Version: 1.0 + Generator: magic (1.0.0) + Root-Is-Purelib: true + Tag: py3-none-any + """, + "METADATA": b"""\ + Metadata-Version: 2.1 + Name: fancy + Version: 1.0.0 + """, + }, + ) + + install( + source=source, + destination=mock_destination, + additional_metadata={}, + trust_wheel_record=True, + ) + + # The scripts-scheme file should go through write_file, not + # write_file_no_record. + write_file_paths = { + c.kwargs["path"] for c in mock_destination.write_file.call_args_list + } + assert "run_fancy" in write_file_paths + + def test_trust_wheel_record_still_hashes_additional_metadata( + self, mock_destination + ): + """Additional metadata should always be hashed since it's generated + by the caller and not present in the wheel's RECORD. + """ + source = FakeWheelSource( + distribution="fancy", + version="1.0.0", + regular_files={ + "fancy/__init__.py": b"""\ + def main(): + print("I'm a fancy package") + """, + }, + dist_info_files={ + "top_level.txt": b"""\ + fancy + """, + "WHEEL": b"""\ + Wheel-Version: 1.0 + Generator: magic (1.0.0) + Root-Is-Purelib: true + Tag: py3-none-any + """, + "METADATA": b"""\ + Metadata-Version: 2.1 + Name: fancy + Version: 1.0.0 + """, + }, + ) + + install( + source=source, + destination=mock_destination, + additional_metadata={ + "fun_file.txt": b"this should be in dist-info!", + }, + trust_wheel_record=True, + ) + + # Additional metadata should go through write_file. + write_file_paths = { + c.kwargs["path"] for c in mock_destination.write_file.call_args_list + } + assert "fancy-1.0.0.dist-info/fun_file.txt" in write_file_paths diff --git a/tests/test_destinations.py b/tests/test_destinations.py index 697fb314..40e6c418 100644 --- a/tests/test_destinations.py +++ b/tests/test_destinations.py @@ -24,6 +24,11 @@ def test_raises_not_implemented_error(self): scheme=None, path=None, stream=None, is_executable=False ) + with pytest.raises(NotImplementedError): + destination.write_file_no_record( + scheme=None, path=None, stream=None, is_executable=False + ) + with pytest.raises(NotImplementedError): destination.finalize_installation( scheme=None, @@ -197,3 +202,39 @@ def test_finalize_write_record(self, destination): def test_blocking_path_traversal(self, destination): with pytest.raises(ValueError): destination._path_with_destdir("purelib", "subdir/../../outside.txt") + + @pytest.mark.parametrize( + ("scheme", "path", "data", "expected"), + [ + pytest.param( + "data", "my_data.bin", b"my data", b"my data", id="normal file" + ), + pytest.param( + "data", + "data_folder/my_data.bin", + b"my data", + b"my data", + id="normal file in subfolder", + ), + ], + ) + def test_write_file_no_record(self, destination, scheme, path, data, expected): + destination.write_file_no_record(scheme, path, io.BytesIO(data), False) + file_data = (Path(destination.scheme_dict[scheme]) / path).read_bytes() + assert file_data == expected + + def test_write_file_no_record_executable(self, destination): + destination.write_file_no_record( + "data", "my_script.sh", io.BytesIO(b"#!/bin/sh\n"), True + ) + file_path = Path(destination.scheme_dict["data"]) / "my_script.sh" + assert file_path.read_bytes() == b"#!/bin/sh\n" + + def test_write_file_no_record_duplicate_raises(self, destination): + destination.write_file_no_record( + "data", "my_data.bin", io.BytesIO(b"my data"), False + ) + with pytest.raises(FileExistsError): + destination.write_file_no_record( + "data", "my_data.bin", io.BytesIO(b"my data"), False + )