Intorduce pex3 venv create.

This new subcommand can create either a venv or just populate a flat
`sys.path` directory entry (ala `pip install --target`) given a set of
requirements to resolve, potentially from a lock or an existing PEX
file, but otherwise from indexes and find links repos. Unlike the
sibling `venv` `pex-tool` subcommand, the target can be selected and,
in the flat `sys.path` directory entry case, it can be a foreign
platform.

Fixes pex-tool#1752
Fixes pex-tool#2110

pex/cli/commands/venv.py

@@ -1,19 +1,38 @@
 # Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import logging
 import os.path
 from argparse import ArgumentParser, _ActionsContainer
 
+from pex import pex_warnings
 from pex.cli.command import BuildTimeCommand
 from pex.commands.command import JsonMixin, OutputMixin
-from pex.common import is_script
+from pex.common import is_script, pluralize
+from pex.dist_metadata import Distribution
+from pex.executor import Executor
+from pex.pex import PEX
 from pex.pex_info import PexInfo
-from pex.result import Error, Ok, Result
+from pex.resolve import configured_resolve, requirement_options, resolver_options, target_options
+from pex.resolve.resolver_configuration import (
+    LockRepositoryConfiguration,
+    PexRepositoryConfiguration,
+)
+from pex.result import Error, Ok, Result, try_
+from pex.targets import LocalInterpreter, Target, Targets
+from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING
+from pex.venv import installer, installer_options
+from pex.venv.install_scope import InstallScope
+from pex.venv.installer import Provenance
+from pex.venv.installer_configuration import InstallerConfiguration
 from pex.venv.virtualenv import Virtualenv
 
 if TYPE_CHECKING:
-    from typing import Any, Dict
+    from typing import Any, Dict, Iterable, Optional
+
+
+logger = logging.getLogger(__name__)
 
 
 class Venv(OutputMixin, JsonMixin, BuildTimeCommand):
@@ -27,6 +46,31 @@ def _add_inspect_arguments(cls, parser):
         cls.add_output_option(parser, entity="venv information")
         cls.add_json_options(parser, entity="venv information")
 
+    @classmethod
+    def _add_create_arguments(cls, parser):
+        # type: (_ActionsContainer) -> None
+        parser.add_argument(
+            "-d",
+            "--dir",
+            dest="dest_dir",
+            metavar="VENV_DIR",
+            required=True,
+            help="The directory to create the venv in.",
+        )
+        parser.add_argument(
+            "--flat",
+            action="store_true",
+            default=False,
+            help=(
+                "Instead of creating a venv, populate a directory that can be used as a flat"
+                "sys.path entry for the selected Python."
+            ),
+        )
+        installer_options.register(parser)
+        target_options.register(parser, include_platforms=True)
+        resolver_options.register(parser, include_pex_repository=True, include_lock=True)
+        requirement_options.register(parser)
+
     @classmethod
     def add_extra_arguments(
         cls,
@@ -44,6 +88,13 @@ def add_extra_arguments(
             include_verbosity=False,
         ) as inspect_parser:
             cls._add_inspect_arguments(inspect_parser)
+        with subcommands.parser(
+            name="create",
+            help="Create a venv.",
+            func=cls._create,
+            include_verbosity=True,
+        ) as create_parser:
+            cls._add_create_arguments(create_parser)
 
     def _inspect(self):
         # type: () -> Result
@@ -101,3 +152,174 @@ def _inspect(self):
             out.write("\n")
 
         return Ok()
+
+    @staticmethod
+    def _install(
+        installer_configuration,  # type: InstallerConfiguration
+        provenance,  # type: Provenance
+        distributions,  # type: Iterable[Distribution]
+        dest_dir,  # type: str
+        hermetic_scripts,  # type: bool
+        pex=None,  # type: Optional[PEX]
+        venv=None,  # type: Optional[Virtualenv]
+    ):
+
+        if pex:
+            if installer_configuration.scope in (InstallScope.ALL, InstallScope.DEPS_ONLY):
+                if venv:
+                    installer.populate_venv_distributions(
+                        venv=venv,
+                        distributions=distributions,
+                        provenance=provenance,
+                        symlink=False,
+                        hermetic_scripts=hermetic_scripts,
+                    )
+                else:
+                    installer.populate_flat_distributions(
+                        dest_dir=dest_dir,
+                        distributions=distributions,
+                        provenance=provenance,
+                        symlink=False,
+                    )
+            if installer_configuration.scope in (InstallScope.ALL, InstallScope.SOURCE_ONLY):
+                if venv:
+                    installer.populate_venv_sources(
+                        venv=venv,
+                        pex=pex,
+                        provenance=provenance,
+                        bin_path=installer_configuration.bin_path,
+                        hermetic_scripts=hermetic_scripts,
+                    )
+                else:
+                    installer.populate_flat_sources(dst=dest_dir, pex=pex, provenance=provenance)
+        else:
+            if venv:
+                installer.populate_venv_distributions(
+                    venv=venv,
+                    distributions=distributions,
+                    provenance=provenance,
+                    symlink=False,
+                    hermetic_scripts=hermetic_scripts,
+                )
+            else:
+                installer.populate_flat_distributions(
+                    dest_dir=dest_dir,
+                    distributions=distributions,
+                    provenance=provenance,
+                    symlink=False,
+                )
+
+    def _create(self):
+        # type: () -> Result
+
+        installer_configuration = installer_options.configure(self.options)
+
+        dest_dir = self.options.dest_dir
+        update = os.path.exists(dest_dir) and not installer_configuration.force
+        flat = self.options.flat
+
+        subject = "flat sys.path directory entry" if flat else "venv"
+
+        venv = None  # type: Optional[Virtualenv]
+        if update and not flat:
+            venv = Virtualenv(venv_dir=dest_dir)
+            target = LocalInterpreter.create(venv.interpreter)  # type: Target
+            targets = Targets.from_target(target)
+        else:
+            targets = target_options.configure(self.options).resolve_targets()
+            target = try_(
+                targets.require_unique_target(
+                    purpose="creating a {subject}".format(subject=subject)
+                )
+            )
+            if not flat:
+                venv = Virtualenv.create(
+                    venv_dir=dest_dir,
+                    interpreter=target.get_interpreter(),
+                    force=installer_configuration.force,
+                    copies=installer_configuration.copies,
+                    prompt=installer_configuration.prompt,
+                )
+
+        if venv and installer_configuration.prompt != venv.custom_prompt:
+            logger.warning(
+                "Unable to apply custom --prompt {prompt!r} in {python} venv; continuing with the "
+                "default prompt.".format(
+                    prompt=installer_configuration.prompt, python=venv.interpreter.identity
+                )
+            )
+
+        requirement_configuration = requirement_options.configure(self.options)
+        resolver_configuration = resolver_options.configure(self.options)
+        with TRACER.timed("Resolving distributions"):
+            installed = configured_resolve.resolve(
+                targets=targets,
+                requirement_configuration=requirement_configuration,
+                resolver_configuration=resolver_configuration,
+            )
+
+        pex = None  # type: Optional[PEX]
+        lock = None  # type: Optional[str]
+        if isinstance(resolver_configuration, PexRepositoryConfiguration):
+            pex = PEX(resolver_configuration.pex_repository, interpreter=target.get_interpreter())
+        elif isinstance(resolver_configuration, LockRepositoryConfiguration):
+            lock = resolver_configuration.lock_file_path
+
+        with TRACER.timed(
+            "Installing {count} {wheels} in {subject} at {dest_dir}".format(
+                count=len(installed.installed_distributions),
+                wheels=pluralize(installed.installed_distributions, "wheel"),
+                subject=subject,
+                dest_dir=dest_dir,
+            )
+        ):
+            hermetic_scripts = not update and installer_configuration.hermetic_scripts
+            distributions = tuple(
+                installed_distribution.distribution
+                for installed_distribution in installed.installed_distributions
+            )
+            provenance = (
+                Provenance.create(venv=venv)
+                if venv
+                else Provenance(target_dir=dest_dir, target_python=target.get_interpreter().binary)
+            )
+            self._install(
+                installer_configuration=installer_configuration,
+                provenance=provenance,
+                distributions=distributions,
+                dest_dir=dest_dir,
+                hermetic_scripts=hermetic_scripts,
+                pex=pex,
+                venv=venv,
+            )
+            source = (
+                "PEX at {pex}".format(pex=pex.path())
+                if pex
+                else "lock at {lock}".format(lock=lock)
+                if lock
+                else "resolved requirements"
+            )
+            provenance.check_collisions(
+                collisions_ok=installer_configuration.collisions_ok, source=source
+            )
+
+        if venv and installer_configuration.pip:
+            with TRACER.timed("Installing Pip"):
+                try_(
+                    installer.ensure_pip_installed(
+                        venv,
+                        distributions=distributions,
+                        scope=installer_configuration.scope,
+                        collisions_ok=installer_configuration.collisions_ok,
+                        source=source,
+                    )
+                )
+
+        if installer_configuration.compile:
+            with TRACER.timed("Compiling venv sources"):
+                try:
+                    target.get_interpreter().execute(["-m", "compileall", dest_dir])
+                except Executor.NonZeroExit as non_zero_exit:
+                    pex_warnings.warn("ignoring compile error {}".format(repr(non_zero_exit)))
+
+        return Ok()

pex/pep_376.py

@@ -229,7 +229,67 @@ def stashed_path(self, *components):
         # type: (*str) -> str
         return os.path.join(self.prefix_dir, self.stash_dir, *components)
 
-    def reinstall(
+    @staticmethod
+    def _create_installed_file(
+        path,  # type: str
+        dest_dir,  # type: str
+    ):
+        # type: (...) -> InstalledFile
+        hasher = hashlib.sha256()
+        hashing.file_hash(path, digest=hasher)
+        return InstalledFile(
+            path=os.path.relpath(path, dest_dir),
+            hash=Hash.create(hasher),
+            size=os.stat(path).st_size,
+        )
+
+    def create_record(
+        self,
+        dst,  # type: str
+        installed_files,  # type: Iterable[InstalledFile]
+    ):
+        # type: (...) -> None
+
+        # The RECORD is a csv file with the path to each installed file in the 1st column.
+        # See: https://www.python.org/dev/peps/pep-0376/#record
+        with safe_open(os.path.join(dst, self.record_relpath), "w") as fp:
+            csv_writer = cast(
+                "CSVWriter",
+                csv.writer(fp, delimiter=",", quotechar='"', lineterminator="\n"),
+            )
+            for installed_file in sorted(installed_files, key=lambda installed: installed.path):
+                csv_writer.writerow(attr.astuple(installed_file, recurse=False))
+
+    def reinstall_flat(
+        self,
+        target_dir,  # type: str
+        symlink=False,  # type: bool
+    ):
+        # type: (...) -> Iterator[Tuple[str, str]]
+        """Re-installs the installed wheel in a flat target directory.
+
+        N.B.: A record of reinstalled files is returned in the form of an iterator that must be
+        consumed to drive the installation to completion.
+
+        If there is an error re-installing a file due to it already existing in the target
+        directory, the error is suppressed, and it's expected that the caller detects this by
+        comparing the record of installed files against those installed previously.
+
+        :return: An iterator over src -> dst pairs.
+        """
+        # >>> _create_installed_file(path=test.flat/pip-23.1.2.dist-info/AUTHORS.txt, dst=test.flat)
+        # >>> _create_installed_file(path=test.flat/pip-23.1.2.dist-info/AUTHORS.txt/requests-2.29.0.dist-info/METADATA, dst=test.flat/pip-23.1.2.dist-info/AUTHORS.txt)
+        installed_files = [InstalledFile(self.record_relpath)]
+        for src, dst in itertools.chain(
+            self._reinstall_stash(dest_dir=target_dir),
+            self._reinstall_site_packages(target_dir, symlink=symlink),
+        ):
+            installed_files.append(self._create_installed_file(path=dst, dest_dir=target_dir))
+            yield src, dst
+
+        self.create_record(target_dir, installed_files)
+
+    def reinstall_venv(
         self,
         venv,  # type: Virtualenv
         symlink=False,  # type: bool
@@ -255,50 +315,36 @@ def reinstall(
 
         installed_files = [InstalledFile(self.record_relpath)]
         for src, dst in itertools.chain(
-            self._reinstall_stash(venv),
+            self._reinstall_stash(dest_dir=venv.venv_dir, interpreter=venv.interpreter),
             self._reinstall_site_packages(site_packages_dir, symlink=symlink),
         ):
-            hasher = hashlib.sha256()
-            hashing.file_hash(dst, digest=hasher)
             installed_files.append(
-                InstalledFile(
-                    path=os.path.relpath(dst, site_packages_dir),
-                    hash=Hash.create(hasher),
-                    size=os.stat(dst).st_size,
-                )
+                self._create_installed_file(path=dst, dest_dir=site_packages_dir)
             )
-
             yield src, dst
 
-        # The RECORD is a csv file with the path to each installed file in the 1st column.
-        # See: https://www.python.org/dev/peps/pep-0376/#record
-        with safe_open(os.path.join(site_packages_dir, self.record_relpath), "w") as fp:
-            csv_writer = cast(
-                "CSVWriter",
-                csv.writer(fp, delimiter=",", quotechar='"', lineterminator="\n"),
-            )
-            for installed_file in sorted(installed_files, key=lambda installed: installed.path):
-                csv_writer.writerow(attr.astuple(installed_file, recurse=False))
+        self.create_record(site_packages_dir, installed_files)
 
-    def _reinstall_stash(self, venv):
-        # type: (Virtualenv) -> Iterator[Tuple[str, str]]
+    def _reinstall_stash(
+        self,
+        dest_dir,  # type: str
+        interpreter=None,  # type: Optional[PythonInterpreter]
+    ):
+        # type: (...) -> Iterator[Tuple[str, str]]
 
         link = True
         stash_abs_path = os.path.join(self.prefix_dir, self.stash_dir)
         for root, dirs, files in os.walk(stash_abs_path, topdown=True, followlinks=True):
-            for d in dirs:
-                src_relpath = os.path.relpath(os.path.join(root, d), stash_abs_path)
-                dst = InstalledFile.denormalized_path(
-                    path=os.path.join(venv.venv_dir, src_relpath), interpreter=venv.interpreter
-                )
-                safe_mkdir(dst)
-
+            dir_created = False
             for f in files:
                 src = os.path.join(root, f)
                 src_relpath = os.path.relpath(src, stash_abs_path)
                 dst = InstalledFile.denormalized_path(
-                    path=os.path.join(venv.venv_dir, src_relpath), interpreter=venv.interpreter
+                    path=os.path.join(dest_dir, src_relpath), interpreter=interpreter
                 )
+                if not dir_created:
+                    safe_mkdir(os.path.dirname(dst))
+                    dir_created = True
                 try:
                     # We only try to link regular files since linking a symlink on Linux can produce
                     # another symlink, which leaves open the possibility the src target could later