Make local dists available in test, run and repl.

[ci skip-rust]

[ci skip-build-wheels]

src/python/pants/backend/python/goals/pytest_runner.py

@@ -13,6 +13,7 @@
 )
 from pants.backend.python.subsystems.pytest import PyTest, PythonTestFieldSet
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.backend.python.util_rules.local_dists import LocalDistsPex, LocalDistsPexRequest
 from pants.backend.python.util_rules.pex import Pex, PexRequest, VenvPex, VenvPexProcess
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.backend.python.util_rules.python_sources import (
@@ -174,6 +175,12 @@ async def setup_pytest_for_target(
         PexFromTargetsRequest,
         PexFromTargetsRequest.for_requirements([request.field_set.address], internal_only=True),
     )
+    local_dists_get = Get(
+        LocalDistsPex,
+        LocalDistsPexRequest(
+            [request.field_set.address], interpreter_constraints=interpreter_constraints
+        ),
+    )
     pytest_pex_get = Get(
         Pex,
         PexRequest(
@@ -198,12 +205,14 @@ async def setup_pytest_for_target(
     (
         pytest_pex,
         requirements_pex,
+        local_dists,
         prepared_sources,
         field_set_source_files,
         extra_output_directory_digest,
     ) = await MultiGet(
         pytest_pex_get,
         requirements_pex_get,
+        local_dists_get,
         prepared_sources_get,
         field_set_source_files_get,
         extra_output_directory_digest_get,
@@ -216,7 +225,7 @@ async def setup_pytest_for_target(
             interpreter_constraints=interpreter_constraints,
             main=pytest.main,
             internal_only=True,
-            pex_path=[pytest_pex, requirements_pex],
+            pex_path=[pytest_pex, requirements_pex, local_dists.pex],
         ),
     )
     config_files_get = Get(

src/python/pants/backend/python/goals/pytest_runner_integration_test.py

@@ -11,17 +11,20 @@
 
 from pants.backend.python import target_types_rules
 from pants.backend.python.dependency_inference import rules as dependency_inference_rules
-from pants.backend.python.goals import package_pex_binary, pytest_runner
+from pants.backend.python.goals import package_pex_binary, pytest_runner, setup_py
 from pants.backend.python.goals.coverage_py import create_or_update_coverage_config
 from pants.backend.python.goals.pytest_runner import PytestPluginSetup, PytestPluginSetupRequest
+from pants.backend.python.macros.python_artifact import PythonArtifact
 from pants.backend.python.subsystems.pytest import PythonTestFieldSet
+from pants.backend.python.subsystems.setuptools import rules as setuptools_rules
 from pants.backend.python.target_types import (
     PexBinary,
+    PythonDistribution,
     PythonLibrary,
     PythonRequirementLibrary,
     PythonTests,
 )
-from pants.backend.python.util_rules import pex_from_targets
+from pants.backend.python.util_rules import local_dists, pex_from_targets
 from pants.core.goals.test import (
     TestDebugRequest,
     TestResult,
@@ -53,10 +56,20 @@ def rule_runner() -> RuleRunner:
             *package_pex_binary.rules(),
             get_filtered_environment,
             *target_types_rules.rules(),
+            *local_dists.rules(),
+            *setup_py.rules(),
+            *setuptools_rules(),
             QueryRule(TestResult, (PythonTestFieldSet,)),
             QueryRule(TestDebugRequest, (PythonTestFieldSet,)),
         ],
-        target_types=[PexBinary, PythonLibrary, PythonTests, PythonRequirementLibrary],
+        target_types=[
+            PexBinary,
+            PythonLibrary,
+            PythonTests,
+            PythonRequirementLibrary,
+            PythonDistribution,
+        ],
+        objects={"python_artifact": PythonArtifact},
     )
 
 
@@ -490,6 +503,48 @@ def test_additional_plugins():
     assert result.exit_code == 0
 
 
+def test_local_dists(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/foo/bar.py": "BAR = 'LOCAL DIST'",
+            f"{PACKAGE}/foo/setup.py": dedent(
+                """\
+                from setuptools import setup
+
+                setup(name="foo", version="9.8.7", packages=["foo"], package_dir={"foo": "."},)
+                """
+            ),
+            f"{PACKAGE}/foo/bar_test.py": dedent(
+                """
+                from foo.bar import BAR
+
+                def test_bar():
+                  assert BAR == "LOCAL DIST"
+                """
+            ),
+            f"{PACKAGE}/foo/BUILD": dedent(
+                """\
+                python_library(name="lib")
+
+                python_distribution(
+                    name="dist",
+                    dependencies=[":lib"],
+                    provides=python_artifact(name="foo", version="9.8.7", setup_script="setup.py"),
+                    setup_py_commands=["bdist_wheel",]
+                )
+
+                # Force-exclude any dep on bar.py, so the only way to consume it is via the dist.
+                python_tests(name="tests", sources=["bar_test.py"],
+                             dependencies=[":dist", "!!:lib"])
+                """
+            ),
+        }
+    )
+    tgt = rule_runner.get_target(Address(os.path.join(PACKAGE, "foo"), target_name="tests"))
+    result = run_pytest(rule_runner, tgt)
+    assert result.exit_code == 0
+
+
 def test_skip_tests(rule_runner: RuleRunner) -> None:
     rule_runner.write_files(
         {

src/python/pants/backend/python/goals/repl.py

@@ -1,7 +1,9 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+import os
 
 from pants.backend.python.subsystems.ipython import IPython
+from pants.backend.python.util_rules.local_dists import LocalDistsPex, LocalDistsPexRequest
 from pants.backend.python.util_rules.pex import Pex, PexRequest
 from pants.backend.python.util_rules.pex_environment import PexEnvironment
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
@@ -10,6 +12,7 @@
     PythonSourceFilesRequest,
 )
 from pants.core.goals.repl import ReplImplementation, ReplRequest
+from pants.engine.addresses import Addresses
 from pants.engine.fs import Digest, MergeDigests
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.unions import UnionRule
@@ -22,19 +25,38 @@ class PythonRepl(ReplImplementation):
 
 @rule(level=LogLevel.DEBUG)
 async def create_python_repl_request(repl: PythonRepl, pex_env: PexEnvironment) -> ReplRequest:
-    requirements_request = Get(
-        Pex,
+
+    # Note that we get an intermediate PexRequest here (instead of going straight to a Pex) so
+    # that we can get the interpreter constraints for use in local_dists_request.
+    requirements_pex_request = await Get(
+        PexRequest,
         PexFromTargetsRequest,
         PexFromTargetsRequest.for_requirements(
             (tgt.address for tgt in repl.targets), internal_only=True
         ),
     )
+    requirements_request = Get(Pex, PexRequest, requirements_pex_request)
+
+    local_dists_request = Get(
+        LocalDistsPex,
+        LocalDistsPexRequest(
+            Addresses(tgt.address for tgt in repl.targets),
+            interpreter_constraints=requirements_pex_request.interpreter_constraints,
+        ),
+    )
+
     sources_request = Get(
         PythonSourceFiles, PythonSourceFilesRequest(repl.targets, include_files=True)
     )
-    requirements_pex, sources = await MultiGet(requirements_request, sources_request)
+
+    requirements_pex, local_dists, sources = await MultiGet(
+        requirements_request, local_dists_request, sources_request
+    )
     merged_digest = await Get(
-        Digest, MergeDigests((requirements_pex.digest, sources.source_files.snapshot.digest))
+        Digest,
+        MergeDigests(
+            (requirements_pex.digest, local_dists.pex.digest, sources.source_files.snapshot.digest)
+        ),
     )
 
     complete_pex_env = pex_env.in_workspace()
@@ -46,6 +68,7 @@ async def create_python_repl_request(repl: PythonRepl, pex_env: PexEnvironment)
     extra_env = {
         **complete_pex_env.environment_dict(python_configured=requirements_pex.python is not None),
         "PEX_EXTRA_SYS_PATH": ":".join(chrooted_source_roots),
+        "PEX_PATH": repl.in_chroot(local_dists.pex.name),
     }
 
     return ReplRequest(digest=merged_digest, args=args, extra_env=extra_env)
@@ -59,8 +82,8 @@ class IPythonRepl(ReplImplementation):
 async def create_ipython_repl_request(
     repl: IPythonRepl, ipython: IPython, pex_env: PexEnvironment
 ) -> ReplRequest:
-    # Note that we get an intermediate PexRequest here (instead of going straight to a Pex)
-    # so that we can get the interpreter constraints for use in ipython_request.
+    # Note that we get an intermediate PexRequest here (instead of going straight to a Pex) so
+    # that we can get the interpreter constraints for use in ipython_request/local_dists_request.
     requirements_pex_request = await Get(
         PexRequest,
         PexFromTargetsRequest,
@@ -71,6 +94,14 @@ async def create_ipython_repl_request(
 
     requirements_request = Get(Pex, PexRequest, requirements_pex_request)
 
+    local_dists_request = Get(
+        LocalDistsPex,
+        LocalDistsPexRequest(
+            [tgt.address for tgt in repl.targets],
+            interpreter_constraints=requirements_pex_request.interpreter_constraints,
+        ),
+    )
+
     sources_request = Get(
         PythonSourceFiles, PythonSourceFilesRequest(repl.targets, include_files=True)
     )
@@ -86,13 +117,18 @@ async def create_ipython_repl_request(
         ),
     )
 
-    requirements_pex, sources, ipython_pex = await MultiGet(
-        requirements_request, sources_request, ipython_request
+    requirements_pex, local_dists, sources, ipython_pex = await MultiGet(
+        requirements_request, local_dists_request, sources_request, ipython_request
     )
     merged_digest = await Get(
         Digest,
         MergeDigests(
-            (requirements_pex.digest, sources.source_files.snapshot.digest, ipython_pex.digest)
+            (
+                requirements_pex.digest,
+                local_dists.pex.digest,
+                sources.source_files.snapshot.digest,
+                ipython_pex.digest,
+            )
         ),
     )
 
@@ -106,7 +142,12 @@ async def create_ipython_repl_request(
     chrooted_source_roots = [repl.in_chroot(sr) for sr in sources.source_roots]
     extra_env = {
         **complete_pex_env.environment_dict(python_configured=ipython_pex.python is not None),
-        "PEX_PATH": repl.in_chroot(requirements_pex_request.output_filename),
+        "PEX_PATH": os.pathsep.join(
+            [
+                repl.in_chroot(requirements_pex_request.output_filename),
+                repl.in_chroot(local_dists.pex.name),
+            ]
+        ),
         "PEX_EXTRA_SYS_PATH": ":".join(chrooted_source_roots),
     }
 

src/python/pants/backend/python/goals/run_pex_binary.py

@@ -9,6 +9,7 @@
     ResolvedPexEntryPoint,
     ResolvePexEntryPointRequest,
 )
+from pants.backend.python.util_rules.local_dists import LocalDistsPex, LocalDistsPexRequest
 from pants.backend.python.util_rules.pex import Pex, PexRequest
 from pants.backend.python.util_rules.pex_environment import PexEnvironment
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
@@ -48,6 +49,14 @@ async def create_pex_binary_run_request(
 
     requirements_request = Get(Pex, PexRequest, requirements_pex_request)
 
+    local_dists_request = Get(
+        LocalDistsPex,
+        LocalDistsPexRequest(
+            [field_set.address],
+            interpreter_constraints=requirements_pex_request.interpreter_constraints,
+        ),
+    )
+
     sources_request = Get(
         PythonSourceFiles, PythonSourceFilesRequest(transitive_targets.closure, include_files=True)
     )
@@ -74,14 +83,19 @@ async def create_pex_binary_run_request(
         ),
     )
 
-    requirements, sources, runner_pex = await MultiGet(
-        requirements_request, sources_request, runner_pex_request
+    requirements, local_dists, sources, runner_pex = await MultiGet(
+        requirements_request, local_dists_request, sources_request, runner_pex_request
     )
 
     merged_digest = await Get(
         Digest,
         MergeDigests(
-            [requirements.digest, sources.source_files.snapshot.digest, runner_pex.digest]
+            [
+                requirements.digest,
+                local_dists.pex.digest,
+                sources.source_files.snapshot.digest,
+                runner_pex.digest,
+            ]
         ),
     )
 
@@ -94,7 +108,12 @@ def in_chroot(relpath: str) -> str:
     chrooted_source_roots = [in_chroot(sr) for sr in sources.source_roots]
     extra_env = {
         **complete_pex_env.environment_dict(python_configured=runner_pex.python is not None),
-        "PEX_PATH": in_chroot(requirements_pex_request.output_filename),
+        "PEX_PATH": os.pathsep.join(
+            [
+                in_chroot(requirements_pex_request.output_filename),
+                in_chroot(local_dists.pex.name),
+            ],
+        ),
         "PEX_EXTRA_SYS_PATH": ":".join(chrooted_source_roots),
     }
 

src/python/pants/backend/python/goals/run_pex_binary_integration_test.py

@@ -158,3 +158,47 @@ def test_no_leak_pex_root_issues_12055() -> None:
         result = run_pants(args)
         result.assert_success()
         assert os.path.join(named_caches_dir, "pex_root") == result.stdout.strip()
+
+
+def test_local_dist() -> None:
+    sources = {
+        "foo/bar.py": "BAR = 'LOCAL DIST'",
+        "foo/setup.py": dedent(
+            """\
+            from setuptools import setup
+
+            # Double-brace the package_dir to avoid setup_tmpdir treating it as a format.
+            setup(name="foo", version="9.8.7", packages=["foo"], package_dir={{"foo": "."}},)
+            """
+        ),
+        "foo/main.py": "from foo.bar import BAR; print(BAR)",
+        "foo/BUILD": dedent(
+            """\
+            python_library(name="lib", sources=["bar.py", "setup.py"])
+
+            python_library(name="main_lib", sources=["main.py"])
+
+            python_distribution(
+                name="dist",
+                dependencies=[":lib"],
+                provides=python_artifact(name="foo", version="9.8.7", setup_script="setup.py"),
+                setup_py_commands=["bdist_wheel",]
+            )
+
+            pex_binary(
+                name="bin",
+                entry_point="main.py",
+                # Force-exclude any dep on bar.py, so the only way to consume it is via the dist.
+                dependencies=[":main_lib", ":dist", "!!:lib"])
+            """
+        ),
+    }
+    with setup_tmpdir(sources) as tmpdir:
+        args = [
+            "--backend-packages=pants.backend.python",
+            f"--source-root-patterns=['/{tmpdir}']",
+            "run",
+            f"{tmpdir}/foo/main.py",
+        ]
+        result = run_pants(args)
+        assert result.stdout == "LOCAL DIST\n"

src/python/pants/backend/python/register.py

@@ -33,6 +33,7 @@
 )
 from pants.backend.python.util_rules import (
     ancestor_files,
+    local_dists,
     pex,
     pex_cli,
     pex_environment,
@@ -59,6 +60,7 @@ def rules():
         *coverage_py.rules(),
         *tailor.rules(),
         *ancestor_files.rules(),
+        *local_dists.rules(),
         *python_sources.rules(),
         *dependency_inference_rules.rules(),
         *pex.rules(),