Implement support for --exclude <req>. (#2281)

When excluding a requirement from a PEX, any resolved distribution
matching that requirement, as well as any of its transitive dependencies
not also needed by non-excluded requirements, are elided from the PEX.

At runtime these missing dependencies will not trigger boot resolve
errors, but they will cause errors if the modules they would have
provided are attempted to be imported. If the intention is to load the
modules from the runtime environment, then `--pex-inherit-path` /
`PEX_INHERIT_PATH` or `PEX_EXTRA_SYS_PATH` knobs must be used to allow
the PEX to see distributions installed in the runtime environment.
Clearly, you must know what you're doing to use this option and not
encounter runtime errors due to import errors. Be ware!

A forthcoming `--provided` option, with similar effects on the PEX
contents, will both automatically inherit any needed missing
distributions from the runtime environment and require all missing
distributions are found; failing fast if they are not.

Work towards #2097.

pex/bin/pex.py

@@ -24,6 +24,7 @@
     register_global_arguments,
 )
 from pex.common import die, is_pyc_dir, is_pyc_file, safe_mkdtemp
+from pex.dependency_manager import DependencyManager
 from pex.enum import Enum
 from pex.inherit_path import InheritPath
 from pex.interpreter_constraints import InterpreterConstraints
@@ -290,6 +291,22 @@ def configure_clp_pex_options(parser):
         ),
     )
 
+    group.add_argument(
+        "--exclude",
+        dest="excluded",
+        default=[],
+        type=str,
+        action="append",
+        help=(
+            "Specifies a requirement to exclude from the built PEX. Any distribution included in "
+            "the PEX's resolve that matches the requirement is excluded from the built PEX along "
+            "with all of its transitive dependencies that are not also required by other "
+            "non-excluded distributions. At runtime, the PEX will boot without checking the "
+            "excluded dependencies are available (say, via `--inherit-path`). This option can be "
+            "used multiple times."
+        ),
+    )
+
     group.add_argument(
         "--compile",
         "--no-compile",
@@ -808,11 +825,18 @@ def build_pex(
     pex_info.strip_pex_env = options.strip_pex_env
     pex_info.interpreter_constraints = interpreter_constraints
 
-    for requirements_pex in options.requirements_pexes:
-        pex_builder.add_from_requirements_pex(requirements_pex)
+    dependency_manager = DependencyManager()
+    excluded = list(options.excluded)  # type: List[str]
 
     with TRACER.timed(
-        "Resolving distributions ({})".format(
+        "Adding distributions from pexes: {}".format(" ".join(options.requirements_pexes))
+    ):
+        for requirements_pex in options.requirements_pexes:
+            requirements_pex_info = dependency_manager.add_from_pex(requirements_pex)
+            excluded.extend(requirements_pex_info.excluded)
+
+    with TRACER.timed(
+        "Resolving distributions for requirements: {}".format(
             " ".join(
                 itertools.chain.from_iterable(
                     (
@@ -824,22 +848,21 @@ def build_pex(
         )
     ):
         try:
-            result = resolve(
-                targets=targets,
-                requirement_configuration=requirement_configuration,
-                resolver_configuration=resolver_configuration,
-                compile_pyc=options.compile,
-                ignore_errors=options.ignore_errors,
-            )
-            for installed_dist in result.installed_distributions:
-                pex_builder.add_distribution(
-                    installed_dist.distribution, fingerprint=installed_dist.fingerprint
+            dependency_manager.add_from_installed(
+                resolve(
+                    targets=targets,
+                    requirement_configuration=requirement_configuration,
+                    resolver_configuration=resolver_configuration,
+                    compile_pyc=options.compile,
+                    ignore_errors=options.ignore_errors,
                 )
-                for direct_req in installed_dist.direct_requirements:
-                    pex_builder.add_requirement(direct_req)
+            )
         except Unsatisfiable as e:
             die(str(e))
 
+    with TRACER.timed("Configuring PEX dependencies"):
+        dependency_manager.configure(pex_builder, excluded=excluded)
+
     if options.entry_point:
         pex_builder.set_entry_point(options.entry_point)
     elif options.script:

pex/dependency_manager.py

@@ -0,0 +1,118 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+from collections import defaultdict
+
+from pex import pex_warnings
+from pex.dist_metadata import Requirement
+from pex.environment import PEXEnvironment
+from pex.exclude_configuration import ExcludeConfiguration
+from pex.fingerprinted_distribution import FingerprintedDistribution
+from pex.orderedset import OrderedSet
+from pex.pep_503 import ProjectName
+from pex.pex_builder import PEXBuilder
+from pex.pex_info import PexInfo
+from pex.resolve.resolvers import Installed
+from pex.tracer import TRACER
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import DefaultDict, Iterable, Iterator
+
+    import attr  # vendor:skip
+else:
+    from pex.third_party import attr
+
+
+@attr.s
+class DependencyManager(object):
+    _requirements = attr.ib(factory=OrderedSet)  # type: OrderedSet[Requirement]
+    _distributions = attr.ib(factory=OrderedSet)  # type: OrderedSet[FingerprintedDistribution]
+
+    def add_from_pex(self, pex):
+        # type: (str) -> PexInfo
+
+        pex_info = PexInfo.from_pex(pex)
+        self._requirements.update(Requirement.parse(req) for req in pex_info.requirements)
+
+        pex_environment = PEXEnvironment.mount(pex, pex_info=pex_info)
+        self._distributions.update(pex_environment.iter_distributions())
+
+        return pex_info
+
+    def add_from_installed(self, installed):
+        # type: (Installed) -> None
+
+        for installed_dist in installed.installed_distributions:
+            self._requirements.update(installed_dist.direct_requirements)
+            self._distributions.add(installed_dist.fingerprinted_distribution)
+
+    def configure(
+        self,
+        pex_builder,  # type: PEXBuilder
+        excluded=(),  # type: Iterable[str]
+    ):
+        # type: (...) -> None
+
+        exclude_configuration = ExcludeConfiguration.create(excluded)
+        exclude_configuration.configure(pex_builder.info)
+
+        dists_by_project_name = defaultdict(
+            OrderedSet
+        )  # type: DefaultDict[ProjectName, OrderedSet[FingerprintedDistribution]]
+        for dist in self._distributions:
+            dists_by_project_name[dist.distribution.metadata.project_name].add(dist)
+
+        root_requirements_by_project_name = defaultdict(
+            OrderedSet
+        )  # type: DefaultDict[ProjectName, OrderedSet[Requirement]]
+        for root_req in self._requirements:
+            root_requirements_by_project_name[root_req.project_name].add(root_req)
+
+        def iter_non_excluded_distributions(requirements):
+            # type: (Iterable[Requirement]) -> Iterator[FingerprintedDistribution]
+            for req in requirements:
+                candidate_dists = dists_by_project_name[req.project_name]
+                for candidate_dist in tuple(candidate_dists):
+                    if candidate_dist.distribution not in req:
+                        continue
+                    candidate_dists.discard(candidate_dist)
+
+                    excluded_by = exclude_configuration.excluded_by(candidate_dist.distribution)
+                    if excluded_by:
+                        excludes = " and ".join(map(str, excluded_by))
+                        TRACER.log(
+                            "Skipping adding {candidate}: excluded by {excludes}".format(
+                                candidate=candidate_dist.distribution, excludes=excludes
+                            )
+                        )
+                        for root_req in root_requirements_by_project_name[
+                            candidate_dist.distribution.metadata.project_name
+                        ]:
+                            if candidate_dist.distribution in root_req:
+                                pex_warnings.warn(
+                                    "The distribution {dist} was required by the input requirement "
+                                    "{root_req} but excluded by configured excludes: "
+                                    "{excludes}".format(
+                                        dist=candidate_dist.distribution,
+                                        root_req=root_req,
+                                        excludes=excludes,
+                                    )
+                                )
+                        continue
+
+                    yield candidate_dist
+                    for dep in iter_non_excluded_distributions(
+                        candidate_dist.distribution.requires()
+                    ):
+                        yield dep
+
+        for fingerprinted_dist in iter_non_excluded_distributions(self._requirements):
+            pex_builder.add_distribution(
+                dist=fingerprinted_dist.distribution, fingerprint=fingerprinted_dist.fingerprint
+            )
+
+        for requirement in self._requirements:
+            pex_builder.add_requirement(requirement)

pex/environment.py

@@ -12,6 +12,7 @@
 from pex import dist_metadata, pex_warnings, targets
 from pex.common import pluralize
 from pex.dist_metadata import Distribution, Requirement
+from pex.exclude_configuration import ExcludeConfiguration
 from pex.fingerprinted_distribution import FingerprintedDistribution
 from pex.inherit_path import InheritPath
 from pex.interpreter import PythonInterpreter
@@ -345,6 +346,7 @@ def _evaluate_marker(
     def _resolve_requirement(
         self,
         requirement,  # type: Requirement
+        exclude_configuration,  # type: ExcludeConfiguration
         resolved_dists_by_key,  # type: MutableMapping[_RequirementKey, FingerprintedDistribution]
         required,  # type: bool
         required_by=None,  # type: Optional[Distribution]
@@ -354,6 +356,16 @@ def _resolve_requirement(
         if requirement_key in resolved_dists_by_key:
             return
 
+        excluded_by = exclude_configuration.excluded_by(requirement)
+        if excluded_by:
+            TRACER.log(
+                "Skipping resolving {requirement}: excluded by {excludes}".format(
+                    requirement=requirement,
+                    excludes=" and ".join(map(str, excluded_by)),
+                )
+            )
+            return
+
         available_distributions = [
             ranked_dist
             for ranked_dist in self._available_ranked_dists_by_project_name[
@@ -406,6 +418,7 @@ def _resolve_requirement(
 
             for not_found in self._resolve_requirement(
                 dep_requirement,
+                exclude_configuration,
                 resolved_dists_by_key,
                 required,
                 required_by=resolved_distribution.distribution,
@@ -502,14 +515,21 @@ def resolve(self):
         # type: () -> Iterable[Distribution]
         if self._resolved_dists is None:
             all_reqs = [Requirement.parse(req) for req in self._pex_info.requirements]
+            exclude_configuration = ExcludeConfiguration.create(excluded=self._pex_info.excluded)
             self._resolved_dists = tuple(
                 fingerprinted_distribution.distribution
-                for fingerprinted_distribution in self.resolve_dists(all_reqs)
+                for fingerprinted_distribution in self.resolve_dists(
+                    all_reqs, exclude_configuration=exclude_configuration
+                )
             )
         return self._resolved_dists
 
-    def resolve_dists(self, reqs):
-        # type: (Iterable[Requirement]) -> Iterable[FingerprintedDistribution]
+    def resolve_dists(
+        self,
+        reqs,  # type: Iterable[Requirement]
+        exclude_configuration=ExcludeConfiguration(),  # type: ExcludeConfiguration
+    ):
+        # type: (...) -> Iterable[FingerprintedDistribution]
 
         self._update_candidate_distributions(self.iter_distributions())
 
@@ -536,6 +556,7 @@ def record_unresolved(dist_not_found):
             with TRACER.timed("Resolving {}".format(qualified_req_or_not_found.requirement), V=2):
                 for not_found in self._resolve_requirement(
                     requirement=qualified_req_or_not_found.requirement,
+                    exclude_configuration=exclude_configuration,
                     required=qualified_req_or_not_found.required,
                     resolved_dists_by_key=resolved_dists_by_key,
                 ):

pex/exclude_configuration.py

@@ -0,0 +1,36 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+from pex.dist_metadata import Distribution, Requirement
+from pex.pex_info import PexInfo
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Iterable, Tuple, Union
+
+    import attr  # vendor:skip
+else:
+    from pex.third_party import attr
+
+
+@attr.s(frozen=True)
+class ExcludeConfiguration(object):
+    @classmethod
+    def create(cls, excluded):
+        # type: (Iterable[str]) -> ExcludeConfiguration
+        return cls(excluded=tuple(Requirement.parse(req) for req in excluded))
+
+    _excluded = attr.ib(factory=tuple)  # type: Tuple[Requirement, ...]
+
+    def configure(self, pex_info):
+        # type: (PexInfo) -> None
+        for excluded in self._excluded:
+            pex_info.add_excluded(excluded)
+
+    def excluded_by(self, item):
+        # type: (Union[Distribution, Requirement]) -> Iterable[Requirement]
+        if isinstance(item, Distribution):
+            return tuple(req for req in self._excluded if item in req)
+        return tuple(req for req in self._excluded if item.project_name == req.project_name)

pex/pex_builder.py

@@ -29,7 +29,6 @@
 from pex.compiler import Compiler
 from pex.dist_metadata import Distribution, MetadataError
 from pex.enum import Enum
-from pex.environment import PEXEnvironment
 from pex.finders import get_entry_point_from_console_script, get_script_from_distributions
 from pex.interpreter import PythonInterpreter
 from pex.layout import Layout
@@ -374,21 +373,6 @@ def add_requirement(self, req):
         self._ensure_unfrozen("Adding a requirement")
         self._pex_info.add_requirement(req)
 
-    def add_from_requirements_pex(self, pex):
-        """Add requirements from an existing pex.
-
-        :param pex: The path to an existing .pex file or unzipped pex directory.
-        """
-        self._ensure_unfrozen("Adding from pex")
-        pex_info = PexInfo.from_pex(pex)
-        pex_environment = PEXEnvironment.mount(pex, pex_info=pex_info)
-        for fingerprinted_dist in pex_environment.iter_distributions():
-            self.add_distribution(
-                dist=fingerprinted_dist.distribution, fingerprint=fingerprinted_dist.fingerprint
-            )
-        for requirement in pex_info.requirements:
-            self.add_requirement(requirement)
-
     def set_executable(self, filename, env_filename=None):
         """Set the executable for this environment.