Retain and index downloads when locking.

pex/resolve/lockfile/__init__.py

@@ -3,23 +3,30 @@
 
 from __future__ import absolute_import
 
+import os
+import shutil
+
 from pex import resolver
-from pex.common import pluralize, safe_open
+from pex.common import pluralize, safe_mkdtemp, safe_open
 from pex.requirements import LocalProjectRequirement, VCSRequirement
 from pex.resolve import resolvers
-from pex.resolve.locked_resolve import LockConfiguration
+from pex.resolve.locked_resolve import Artifact, LockConfiguration
+from pex.resolve.lockfile.download_manager import DownloadedArtifact, DownloadManager
 from pex.resolve.lockfile.lockfile import Lockfile as Lockfile  # For re-export.
 from pex.resolve.requirement_configuration import RequirementConfiguration
 from pex.resolve.resolver_configuration import PipConfiguration
+from pex.resolver import Downloaded
 from pex.result import Error
 from pex.targets import Targets
 from pex.third_party.pkg_resources import Requirement
+from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING
+from pex.util import CacheHelper
 from pex.variables import ENV
 from pex.version import __version__
 
 if TYPE_CHECKING:
-    from typing import List, Text, Union
+    from typing import List, Mapping, Optional, Text, Union
 
 
 class ParseError(Exception):
@@ -77,6 +84,51 @@ def store(
         json.dump(json_codec.as_json_data(lockfile), fp, sort_keys=True)
 
 
+class CreateLockDownloadManager(DownloadManager):
+    @classmethod
+    def create(
+        cls,
+        download_dir,  # type: str
+        downloaded,  # type: Downloaded
+        pex_root=None,  # type: Optional[str]
+    ):
+        # type: (...) -> CreateLockDownloadManager
+
+        artifacts_by_filename = {
+            artifact.filename: artifact
+            for locked_resolve in downloaded.locked_resolves
+            for locked_requirement in locked_resolve.locked_requirements
+            for artifact in locked_requirement.iter_artifacts()
+        }
+        path_by_artifact = {
+            artifacts_by_filename[f]: os.path.join(root, f)
+            for root, _, files in os.walk(download_dir)
+            for f in files
+        }
+        return cls(path_by_artifact=path_by_artifact, pex_root=pex_root)
+
+    def __init__(
+        self,
+        path_by_artifact,  # type: Mapping[Artifact, str]
+        pex_root=None,  # type: Optional[str]
+    ):
+        super(CreateLockDownloadManager, self).__init__(pex_root=pex_root)
+        self._path_by_artifact = path_by_artifact
+
+    def store_all(self):
+        for artifact in self._path_by_artifact:
+            self.store(artifact)
+
+    def save(
+        self,
+        artifact,  # type: Artifact
+        path,  # type: str
+    ):
+        # type: (...) -> str
+        shutil.move(self._path_by_artifact[artifact], path)
+        return CacheHelper.hash(path)
+
+
 def create(
     lock_configuration,  # type: LockConfiguration
     requirement_configuration,  # type: RequirementConfiguration
@@ -120,6 +172,8 @@ def create(
         for constraint in requirement_configuration.parse_constraints(network_configuration)
     )
 
+    dest = safe_mkdtemp()
+
     try:
         downloaded = resolver.download(
             targets=targets,
@@ -140,13 +194,17 @@ def create(
             build_isolation=pip_configuration.build_isolation,
             max_parallel_jobs=pip_configuration.max_jobs,
             lock_configuration=lock_configuration,
-            # We're just out for the lock data and not the distribution files downloaded to produce
-            # that data.
-            dest=None,
+            dest=dest,
         )
     except resolvers.ResolveError as e:
         return Error(str(e))
 
+    with TRACER.timed("Indexing downloads"):
+        create_lock_download_manager = CreateLockDownloadManager.create(
+            download_dir=dest, downloaded=downloaded
+        )
+        create_lock_download_manager.store_all()
+
     return Lockfile.create(
         pex_version=__version__,
         style=lock_configuration.style,

pex/resolve/lockfile/download_manager.py

@@ -0,0 +1,56 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+import os
+
+from pex.common import atomic_directory
+from pex.resolve.locked_resolve import Artifact
+from pex.typing import TYPE_CHECKING
+from pex.variables import ENV
+
+if TYPE_CHECKING:
+    from typing import Optional
+
+    import attr  # vendor:skip
+else:
+    from pex.third_party import attr
+
+
+@attr.s(frozen=True)
+class DownloadedArtifact(object):
+    path = attr.ib()  # type: str
+
+    def fingerprint(self):
+        # type: () -> str
+        with open(os.path.join(os.path.dirname(self.path), "sha1"), "rb") as fp:
+            return str(fp.read().decode("ascii"))
+
+
+class DownloadManager(object):
+    def __init__(self, pex_root=None):
+        # type: (Optional[str]) -> None
+        self._download_dir = os.path.join(pex_root or ENV.PEX_ROOT, "downloads")
+
+    def store(self, artifact):
+        # type: (Artifact) -> DownloadedArtifact
+
+        download_dir = os.path.join(self._download_dir, artifact.fingerprint.hash)
+        with atomic_directory(download_dir, exclusive=True) as atomic_dir:
+            if not atomic_dir.is_finalized():
+                dest = os.path.join(atomic_dir.work_dir, artifact.filename)
+                internal_fingerprint = self.save(artifact, dest)
+                with open(os.path.join(atomic_dir.work_dir, "sha1"), "wb") as fp:
+                    fp.write(internal_fingerprint.encode("ascii"))
+
+        return DownloadedArtifact(path=os.path.join(download_dir, artifact.filename))
+
+    def save(
+        self,
+        artifact,  # type: Artifact
+        path,  # type: str
+    ):
+        # type: (...) -> str
+        """Save the given `artifact` at `path` and return its sha1 hex digest."""
+        raise NotImplementedError()

tests/integration/cli/commands/test_lock.py

@@ -1,6 +1,7 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import hashlib
 import os
 import re
 from textwrap import dedent
@@ -22,6 +23,7 @@
 from pex.testing import IS_MAC, IS_PYPY, PY310, PY_VER, IntegResults, ensure_python_interpreter
 from pex.third_party.pkg_resources import Requirement
 from pex.typing import TYPE_CHECKING
+from pex.util import CacheHelper
 from pex.version import __version__
 
 if TYPE_CHECKING:
@@ -73,21 +75,49 @@ def test_create(tmpdir):
 def test_create_style(tmpdir):
     # type: (Any) -> None
 
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+
     def create_lock(
         style,  # type: str
         interpreter_constraint=None,  # type: Optional[str]
     ):
         # type: (...) -> LockedRequirement
         lock_file = os.path.join(str(tmpdir), "{}.lock".format(style))
-        args = ["lock", "create", "psutil==5.9.0", "-o", lock_file, "--style", style]
+        args = [
+            "lock",
+            "create",
+            "psutil==5.9.0",
+            "-o",
+            lock_file,
+            "--style",
+            style,
+            "--pex-root",
+            pex_root,
+        ]
         if interpreter_constraint:
             args.extend(["--interpreter-constraint", interpreter_constraint])
         run_pex3(*args).assert_success()
         lock = lockfile.load(lock_file)
         assert 1 == len(lock.locked_resolves)
         locked_resolve = lock.locked_resolves[0]
         assert 1 == len(locked_resolve.locked_requirements)
-        return locked_resolve.locked_requirements[0]
+        locked_requirement = locked_resolve.locked_requirements[0]
+        download_dir = os.path.join(
+            pex_root, "downloads", locked_requirement.artifact.fingerprint.hash
+        )
+        downloaded_artifact = os.path.join(download_dir, locked_requirement.artifact.filename)
+        assert os.path.exists(downloaded_artifact), (
+            "Expected the primary artifact to be downloaded as a side-effect of executing the lock "
+            "resolve."
+        )
+        downloaded_artifact_internal_fingerprint = os.path.join(download_dir, "sha1")
+        assert os.path.exists(downloaded_artifact_internal_fingerprint), (
+            "Expected the primary artifact to have an internal fingerprint established to short "
+            "circuit builds and installs."
+        )
+        with open(downloaded_artifact_internal_fingerprint) as fp:
+            assert CacheHelper.hash(downloaded_artifact, digest=hashlib.sha1()) == fp.read()
+        return locked_requirement
 
     # See: https://pypi.org/project/psutil/5.9.0/#files