diff --git a/pex/cache/dirs.py b/pex/cache/dirs.py index da04d1e28..aa7fd87d3 100644 --- a/pex/cache/dirs.py +++ b/pex/cache/dirs.py @@ -7,6 +7,7 @@ import os from pex.enum import Enum +from pex.exceptions import production_assert from pex.typing import TYPE_CHECKING, cast from pex.variables import ENV, Variables @@ -399,7 +400,14 @@ def create( if os.path.islink(wheel_dir): symlink_dir = os.path.dirname(wheel_dir) wheel_dir = os.path.realpath(wheel_dir) - wheel_hash = os.path.basename(os.path.dirname(wheel_dir)) + recorded_wheel_hash = os.path.basename(os.path.dirname(wheel_dir)) + if wheel_hash: + production_assert(wheel_hash == recorded_wheel_hash) + else: + wheel_hash = recorded_wheel_hash + elif wheel_hash is not None: + symlink_dir = os.path.dirname(wheel_dir) + wheel_dir = CacheDir.INSTALLED_WHEELS.path(wheel_hash, wheel_name, pex_root=pex_root) return cls( path=wheel_dir, diff --git a/pex/pip/installation.py b/pex/pip/installation.py index c4918ead4..c80c2110f 100644 --- a/pex/pip/installation.py +++ b/pex/pip/installation.py @@ -11,7 +11,7 @@ from pex import pep_427, pex_warnings, third_party from pex.atomic_directory import atomic_directory -from pex.cache.dirs import CacheDir, PipPexDir +from pex.cache.dirs import InstalledWheelDir, PipPexDir from pex.common import REPRODUCIBLE_BUILDS_ENV, CopyMode, pluralize, safe_mkdtemp from pex.dist_metadata import Requirement from pex.exceptions import production_assert @@ -28,7 +28,7 @@ from pex.result import Error, try_ from pex.targets import LocalInterpreter, RequiresPythonError, Targets from pex.tracer import TRACER -from pex.typing import TYPE_CHECKING +from pex.typing import TYPE_CHECKING, cast from pex.util import CacheHelper from pex.variables import ENV, Variables from pex.venv.virtualenv import InstallationChoice, Virtualenv @@ -208,22 +208,27 @@ def _install_wheel(wheel_path): # https://github.com/pex-tool/pex/issues/2556 wheel_hash = CacheHelper.hash(wheel_path, hasher=hashlib.sha256) wheel_name = os.path.basename(wheel_path) - destination = CacheDir.INSTALLED_WHEELS.path(wheel_hash, wheel_name) - with atomic_directory(destination) as atomic_dir: + installed_wheel_dir = InstalledWheelDir.create(wheel_name=wheel_name, install_hash=wheel_hash) + with atomic_directory(installed_wheel_dir) as atomic_dir: if not atomic_dir.is_finalized(): installed_wheel = pep_427.install_wheel_chroot( wheel_path=wheel_path, destination=atomic_dir.work_dir ) - runtime_key_dir = CacheDir.INSTALLED_WHEELS.path( - installed_wheel.fingerprint - or CacheHelper.dir_hash(atomic_dir.work_dir, hasher=hashlib.sha256) + runtime_key_dir = InstalledWheelDir.create( + wheel_name=wheel_name, + install_hash=( + installed_wheel.fingerprint + or CacheHelper.dir_hash(atomic_dir.work_dir, hasher=hashlib.sha256) + ), + wheel_hash=wheel_hash, ) - with atomic_directory(runtime_key_dir) as runtime_atomic_dir: + production_assert(runtime_key_dir.symlink_dir is not None) + with atomic_directory(cast(str, runtime_key_dir.symlink_dir)) as runtime_atomic_dir: if not runtime_atomic_dir.is_finalized(): source_path = os.path.join(runtime_atomic_dir.work_dir, wheel_name) - relative_target_path = os.path.relpath(destination, runtime_key_dir) + relative_target_path = os.path.relpath(installed_wheel_dir, runtime_key_dir) os.symlink(relative_target_path, source_path) - return destination + return installed_wheel_dir def _bootstrap_pip( diff --git a/pex/third_party/__init__.py b/pex/third_party/__init__.py index ff9fe3baa..18019205f 100644 --- a/pex/third_party/__init__.py +++ b/pex/third_party/__init__.py @@ -24,6 +24,7 @@ if TYPE_CHECKING: from typing import Container, Iterable, Iterator, List, Optional, Tuple + from pex.cache.dirs import InstalledWheelDir # noqa from pex.interpreter import PythonInterpreter @@ -628,10 +629,10 @@ def expose_installed_wheels( dists, # type: Iterable[str] interpreter=None, # type: Optional[PythonInterpreter] ): - # type: (...) -> Iterator[str] + # type: (...) -> Iterator[InstalledWheelDir] from pex.atomic_directory import atomic_directory - from pex.cache.dirs import CacheDir + from pex.cache.dirs import InstalledWheelDir from pex.pep_376 import InstalledWheel for path in expose(dists, interpreter=interpreter): @@ -642,12 +643,14 @@ def expose_installed_wheels( install_hash = installed_wheel.fingerprint or CacheHelper.dir_hash( path, hasher=hashlib.sha256 ) - wheel_path = CacheDir.INSTALLED_WHEELS.path(install_hash, wheel_file_name) - with atomic_directory(wheel_path) as atomic_dir: + installed_wheel_dir = InstalledWheelDir.create( + wheel_name=wheel_file_name, install_hash=install_hash + ) + with atomic_directory(installed_wheel_dir) as atomic_dir: if not atomic_dir.is_finalized(): for _src, _dst in iter_copytree(path, atomic_dir.work_dir, copy_mode=CopyMode.LINK): pass - yield wheel_path + yield installed_wheel_dir # Implicitly install an importer for vendored code on the first import of pex.third_party.