Prefix the entire setup.py chroot. (#12372)

More specifically:

Add a prefix to the chroot. This prefix only exists
at build time, and does not affect the built dist,
but allows us to capture setuptools output from input
dirs without capturing the entire sandbox.

No longer prefix just the sources within the chroot.
This was never particularly necessary, and I originally
implemented that somewhat arbitrarily. Note that this means
that from now on dists will have a different internal structure,
but one that is more usual and simpler, so probably better.

Run setup.py with its directory as the cwd. Custom setup.pys
will expect that.

This will be useful when we run the setuptools develop command
on setup.py scripts that build native code and write the .so or
other output files into the input dirs.

This is a replay of #12359 with a fix for the issue that caused it
to be reverted.

[ci skip-rust]

Author: benjyw

src/python/pants/BUILD

@@ -28,7 +28,7 @@ python_distribution(
     # used when creating the wheel, which is a good thing. We should be setting this ABI to ensure
     # consumers of pantsbuild.pants are using a compatible interpreter.
     # TODO(7344): the tuple syntax for ext_modules is deprecated. Use Extension once we support it.
-    ext_modules=[('native_engine', {'sources': ['src/pants/dummy.c']})],
+    ext_modules=[('native_engine', {'sources': ['pants/dummy.c']})],
   ).with_binaries(
     pants='src/python/pants/bin:pants',
   )

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

@@ -405,10 +405,6 @@ async def package_python_dist(
         return BuiltPackage(rel_chroot, (BuiltPackageArtifact(dirname),))
 
 
-# We write .py sources into the chroot under this dir.
-CHROOT_SOURCE_ROOT = "src"
-
-
 SETUP_BOILERPLATE = """
 # DO NOT EDIT THIS FILE -- AUTOGENERATED BY PANTS
 # Target: {target_address_spec}
@@ -433,22 +429,40 @@ async def run_setup_py(req: RunSetupPyRequest, setuptools: Setuptools) -> RunSet
             interpreter_constraints=req.interpreter_constraints,
         ),
     )
+
+    # We prefix the entire chroot, and run with this prefix as the cwd, so that we can capture any
+    # changes setup made within it (e.g., when running 'develop') without also capturing other
+    # artifacts of the pex process invocation.
+    chroot_prefix = "chroot"
+
     # The setuptools dist dir, created by it under the chroot (not to be confused with
     # pants's own dist dir, at the buildroot).
-    dist_dir = "dist/"
+    dist_dir = "dist"
+
+    prefixed_chroot = await Get(Digest, AddPrefix(req.chroot.digest, chroot_prefix))
+
+    # setup.py basically always expects to be run with the cwd as its own directory
+    # (e.g., paths in it are relative to that directory). This is true of the setup.py
+    # we generate and is overwhelmingly likely to be true for existing setup.py files,
+    # as there is no robust way to run them otherwise.
+    setup_script_reldir, setup_script_name = os.path.split(req.chroot.setup_script)
+    working_directory = os.path.join(chroot_prefix, setup_script_reldir)
+
     result = await Get(
         ProcessResult,
         VenvPexProcess(
             setuptools_pex,
-            argv=(req.chroot.setup_script, *req.args),
-            input_digest=req.chroot.digest,
+            argv=(setup_script_name, *req.args),
+            input_digest=prefixed_chroot,
+            working_directory=working_directory,
             # setuptools commands that create dists write them to the distdir.
             # TODO: Could there be other useful files to capture?
-            output_directories=(dist_dir,),
+            output_directories=(dist_dir,),  # Relative to the working_directory.
             description=f"Run setuptools for {req.exported_target.target.address}",
             level=LogLevel.DEBUG,
         ),
     )
+    # Note that output_digest paths are relative to the working_directory.
     output_digest = await Get(Digest, RemovePrefix(result.output_digest, dist_dir))
     return RunSetupPyResult(output_digest)
 
@@ -535,7 +549,6 @@ async def generate_chroot(request: SetupPyChrootRequest) -> SetupPyChroot:
     # specified `SetupKwargs(_allow_banned_keys=True)`.
     setup_kwargs.update(
         {
-            "package_dir": {"": CHROOT_SOURCE_ROOT, **setup_kwargs.get("package_dir", {})},
             "packages": (*sources.packages, *(setup_kwargs.get("packages", []))),
             "namespace_packages": (
                 *sources.namespace_packages,
@@ -595,13 +608,8 @@ async def generate_chroot(request: SetupPyChrootRequest) -> SetupPyChroot:
         FileContent("setup.py", setup_py_content),
         FileContent("MANIFEST.in", "include *.py".encode()),
     ]
-    extra_files_digest, src_digest = await MultiGet(
-        Get(Digest, CreateDigest(files_to_create)),
-        # Nest the sources under the src/ prefix.
-        Get(Digest, AddPrefix(sources.digest, CHROOT_SOURCE_ROOT)),
-    )
-
-    chroot_digest = await Get(Digest, MergeDigests((src_digest, extra_files_digest)))
+    extra_files_digest = await Get(Digest, CreateDigest(files_to_create))
+    chroot_digest = await Get(Digest, MergeDigests((sources.digest, extra_files_digest)))
     return SetupPyChroot(
         chroot_digest, "setup.py", FinalizedSetupKwargs(setup_kwargs, address=target.address)
     )

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

@@ -285,13 +285,13 @@ def test_generate_chroot(chroot_rule_runner: RuleRunner) -> None:
     assert_chroot(
         chroot_rule_runner,
         [
-            "src/files/README.txt",
-            "src/foo/qux/__init__.py",
-            "src/foo/qux/qux.py",
-            "src/foo/qux/qux.pyi",
-            "src/foo/resources/js/code.js",
-            "src/foo/__init__.py",
-            "src/foo/foo.py",
+            "files/README.txt",
+            "foo/qux/__init__.py",
+            "foo/qux/qux.py",
+            "foo/qux/qux.pyi",
+            "foo/resources/js/code.js",
+            "foo/__init__.py",
+            "foo/foo.py",
             "setup.py",
             "MANIFEST.in",
         ],
@@ -300,7 +300,6 @@ def test_generate_chroot(chroot_rule_runner: RuleRunner) -> None:
             "name": "foo",
             "version": "1.2.3",
             "plugin_demo": "hello world",
-            "package_dir": {"": "src"},
             "packages": ("foo", "foo.qux"),
             "namespace_packages": ("foo",),
             "package_data": {"foo": ("resources/js/code.js",)},
@@ -378,13 +377,12 @@ def test_binary_shorthand(chroot_rule_runner: RuleRunner) -> None:
     )
     assert_chroot(
         chroot_rule_runner,
-        ["src/project/app.py", "setup.py", "MANIFEST.in"],
+        ["project/app.py", "setup.py", "MANIFEST.in"],
         "setup.py",
         {
             "name": "bin",
             "version": "1.1.1",
             "plugin_demo": "hello world",
-            "package_dir": {"": "src"},
             "packages": ("project",),
             "namespace_packages": (),
             "install_requires": (),

src/python/pants/engine/process_test.py

@@ -5,7 +5,14 @@
 
 import pytest
 
-from pants.engine.fs import EMPTY_DIGEST, CreateDigest, Digest, DigestContents, FileContent
+from pants.engine.fs import (
+    EMPTY_DIGEST,
+    CreateDigest,
+    Digest,
+    DigestContents,
+    Directory,
+    FileContent,
+)
 from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.process import (
     BinaryPathRequest,
@@ -74,11 +81,24 @@ def test_env(rule_runner: RuleRunner) -> None:
     assert b"VAR2=VAL" in result.stdout
 
 
-def test_output_digest(rule_runner: RuleRunner) -> None:
+@pytest.mark.parametrize("working_directory", ["", "subdir"])
+def test_output_digest(rule_runner: RuleRunner, working_directory) -> None:
+    # Test that the output files are relative to the working directory, both in how
+    # they're specified, and their paths in the output_digest.
+    input_digest = (
+        rule_runner.request(
+            Digest,
+            [CreateDigest([Directory(working_directory)])],
+        )
+        if working_directory
+        else EMPTY_DIGEST
+    )
     process = Process(
+        input_digest=input_digest,
         argv=("/bin/bash", "-c", "echo -n 'European Burmese' > roland"),
         description="echo roland",
         output_files=("roland",),
+        working_directory=working_directory,
     )
     result = rule_runner.request(ProcessResult, [process])
     assert result.output_digest == Digest(