Initial support for building dists using mypyc.

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

@@ -12,6 +12,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from functools import partial
+from pathlib import PurePath
 from typing import Any, DefaultDict, Dict, List, Mapping, Tuple, cast
 
 from pants.backend.python.macros.python_artifact import PythonArtifact
@@ -41,6 +42,7 @@
 )
 from pants.backend.python.util_rules.dists import rules as dists_rules
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.backend.python.util_rules.pex import Pex
 from pants.backend.python.util_rules.pex_requirements import PexRequirements
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFiles,
@@ -80,6 +82,7 @@
 from pants.engine.unions import UnionMembership, UnionRule, union
 from pants.option.option_types import BoolOption, EnumOption
 from pants.option.subsystem import Subsystem
+from pants.source.source_root import SourceRootsRequest, SourceRootsResult
 from pants.util.docutil import doc_url
 from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
@@ -370,10 +373,31 @@ class NoDistTypeSelected(ValueError):
     pass
 
 
+@union
+@dataclass(frozen=True)
+class DistBuildEnvironmentRequest:
+    target_addresses: tuple[Address, ...]
+    interpreter_constraints: InterpreterConstraints
+
+    @classmethod
+    def is_applicable(cls, tgt: Target) -> bool:
+        # Union members should override.
+        return False
+
+
+@dataclass(frozen=True)
+class DistBuildEnvironment:
+    """Various extra information that might be needed to build a dist."""
+
+    extra_build_time_requirements: tuple[Pex, ...]
+    extra_build_time_inputs: Digest
+
+
 @rule
 async def package_python_dist(
     field_set: PythonDistributionFieldSet,
     python_setup: PythonSetup,
+    union_membership: UnionMembership,
 ) -> BuiltPackage:
     transitive_targets = await Get(TransitiveTargets, TransitiveTargetsRequest([field_set.address]))
     exported_target = ExportedTarget(transitive_targets.roots[0])
@@ -401,25 +425,66 @@ async def package_python_dist(
         ),
     )
 
+    # Find the source roots for the build-time 1stparty deps (e.g., deps of setup.py).
+    source_roots_result = await Get(
+        SourceRootsResult,
+        SourceRootsRequest(
+            files=[], dirs={PurePath(tgt.address.spec_path) for tgt in transitive_targets.closure}
+        ),
+    )
+    source_roots = tuple(sorted({sr.path for sr in source_roots_result.path_to_root.values()}))
+
+    # Get any extra build-time environment (e.g., native extension requirements).
+    build_env_requests = []
+    build_env_request_types = union_membership[DistBuildEnvironmentRequest]
+    for build_env_request_type in build_env_request_types:
+        if build_env_request_type.is_applicable(dist_tgt):
+            build_env_requests.append(
+                build_env_request_type(
+                    tuple(tt.address for tt in transitive_targets.closure), interpreter_constraints
+                )
+            )
+
+    build_envs = await MultiGet(
+        [
+            Get(DistBuildEnvironment, DistBuildEnvironmentRequest, build_env_request)
+            for build_env_request in build_env_requests
+        ]
+    )
+    extra_build_time_requirements = tuple(
+        itertools.chain.from_iterable(
+            build_env.extra_build_time_requirements for build_env in build_envs
+        )
+    )
+    input_digest = await Get(
+        Digest,
+        MergeDigests(
+            [chroot.digest, *(build_env.extra_build_time_inputs for build_env in build_envs)]
+        ),
+    )
+
     # We prefix the entire chroot, and run with this prefix as the cwd, so that we can capture
     # any changes setup made within it without also capturing other artifacts of the pex
     # process invocation.
     chroot_prefix = "chroot"
     working_directory = os.path.join(chroot_prefix, chroot.working_directory)
-    prefixed_chroot = await Get(Digest, AddPrefix(chroot.digest, chroot_prefix))
-    build_system = await Get(BuildSystem, BuildSystemRequest(prefixed_chroot, working_directory))
+    prefixed_input = await Get(Digest, AddPrefix(input_digest, chroot_prefix))
+    build_system = await Get(BuildSystem, BuildSystemRequest(prefixed_input, working_directory))
+
     setup_py_result = await Get(
         DistBuildResult,
         DistBuildRequest(
             build_system=build_system,
             interpreter_constraints=interpreter_constraints,
             build_wheel=wheel,
             build_sdist=sdist,
-            input=prefixed_chroot,
+            input=prefixed_input,
             working_directory=working_directory,
+            build_time_source_roots=source_roots,
             target_address_spec=exported_target.target.address.spec,
             wheel_config_settings=wheel_config_settings,
             sdist_config_settings=sdist_config_settings,
+            extra_build_time_requirements=extra_build_time_requirements,
         ),
     )
     dist_snapshot = await Get(Snapshot, Digest, setup_py_result.output)

src/python/pants/backend/python/subsystems/setuptools_test.py

@@ -23,7 +23,7 @@ def test_setup_lockfile_interpreter_constraints() -> None:
             QueryRule(GeneratePythonLockfile, [SetuptoolsLockfileSentinel]),
         ],
         target_types=[PythonSourcesGeneratorTarget, PythonDistribution],
-        objects={"setup_py": PythonArtifact},
+        objects={"python_artifact": PythonArtifact},
     )
 
     global_constraint = "==3.9.*"
@@ -49,7 +49,7 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
             python_distribution(
                 name="dist",
                 dependencies=[":lib"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
             """
         ),
@@ -62,7 +62,7 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
             python_distribution(
                 name="dist",
                 dependencies=[":lib"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
             """
         ),
@@ -75,7 +75,7 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
             python_distribution(
                 name="dist",
                 dependencies=[":lib"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
             """
         ),
@@ -94,14 +94,14 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
             python_distribution(
                 name="dist1",
                 dependencies=[":lib1"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
 
             python_sources(name="lib2", interpreter_constraints=["==3.5.*"])
             python_distribution(
                 name="dist2",
                 dependencies=[":lib2"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
             """
         ),
@@ -114,14 +114,14 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
             python_distribution(
                 name="dist1",
                 dependencies=[":lib1"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
 
             python_sources(name="lib2", interpreter_constraints=[">=3.5"])
             python_distribution(
                 name="dist2",
                 dependencies=[":lib2"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
             """
         ),
@@ -134,21 +134,21 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
             python_distribution(
                 name="dist1",
                 dependencies=[":lib1"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
 
             python_sources(name="lib2", interpreter_constraints=["==2.7.*"])
             python_distribution(
                 name="dist2",
                 dependencies=[":lib2"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
 
             python_sources(name="lib3", interpreter_constraints=[">=3.6"])
             python_distribution(
                 name="dist3",
                 dependencies=[":lib3"],
-                provides=setup_py(name="dist"),
+                provides=python_artifact(name="dist"),
             )
             """
         ),

src/python/pants/backend/python/typecheck/mypy/mypyc.py

@@ -0,0 +1,83 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pants.backend.python.goals.setup_py import DistBuildEnvironment, DistBuildEnvironmentRequest
+from pants.backend.python.target_types import PythonDistribution
+from pants.backend.python.typecheck.mypy.subsystem import (
+    MyPy,
+    MyPyConfigFile,
+    MyPyFirstPartyPlugins,
+)
+from pants.backend.python.util_rules import pex_from_targets
+from pants.backend.python.util_rules.pex import Pex, PexRequest
+from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
+from pants.backend.python.util_rules.pex_requirements import PexRequirements
+from pants.engine.rules import Get, MultiGet, collect_rules, rule
+from pants.engine.target import BoolField, Target
+from pants.engine.unions import UnionRule
+
+
+class UsesMyPycField(BoolField):
+    alias = "uses_mypyc"
+    default = False
+    help = "If true, this distribution is built using mypyc."
+
+
+@dataclass(frozen=True)
+class MyPycDistBuildEnvironmentRequest(DistBuildEnvironmentRequest):
+    @classmethod
+    def is_applicable(cls, tgt: Target) -> bool:
+        return tgt.get(UsesMyPycField).value
+
+
+@rule(desc="Get mypyc build environment")
+async def get_mypyc_build_environment(
+    request: MyPycDistBuildEnvironmentRequest,
+    first_party_plugins: MyPyFirstPartyPlugins,
+    mypy_config_file: MyPyConfigFile,
+    mypy: MyPy,
+) -> DistBuildEnvironment:
+    mypy_pex_get = Get(
+        Pex,
+        PexRequest,
+        mypy.to_pex_request(
+            interpreter_constraints=request.interpreter_constraints,
+            extra_requirements=first_party_plugins.requirement_strings,
+        ),
+    )
+    requirements_pex_get = Get(
+        Pex,
+        RequirementsPexRequest(
+            addresses=request.target_addresses,
+            hardcoded_interpreter_constraints=request.interpreter_constraints,
+        ),
+    )
+    extra_type_stubs_pex_get = Get(
+        Pex,
+        PexRequest(
+            output_filename="extra_type_stubs.pex",
+            internal_only=True,
+            requirements=PexRequirements(mypy.extra_type_stubs),
+            interpreter_constraints=request.interpreter_constraints,
+        ),
+    )
+    (mypy_pex, requirements_pex, extra_type_stubs_pex) = await MultiGet(
+        mypy_pex_get, requirements_pex_get, extra_type_stubs_pex_get
+    )
+    return DistBuildEnvironment(
+        extra_build_time_requirements=(mypy_pex, requirements_pex, extra_type_stubs_pex),
+        extra_build_time_inputs=mypy_config_file.digest,
+    )
+
+
+def rules():
+    return [
+        *collect_rules(),
+        UnionRule(DistBuildEnvironmentRequest, MyPycDistBuildEnvironmentRequest),
+        PythonDistribution.register_plugin_field(UsesMyPycField),
+        *pex_from_targets.rules(),
+    ]

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

@@ -7,9 +7,10 @@
 https://mypy.readthedocs.io/en/stable/.
 """
 
+from pants.backend.python.typecheck.mypy import mypyc
 from pants.backend.python.typecheck.mypy import rules as mypy_rules
 from pants.backend.python.typecheck.mypy import skip_field, subsystem
 
 
 def rules():
-    return (*mypy_rules.rules(), *skip_field.rules(), *subsystem.rules())
+    return (*mypy_rules.rules(), *mypyc.rules(), *skip_field.rules(), *subsystem.rules())

src/python/pants/backend/python/util_rules/dists.py

@@ -14,7 +14,7 @@
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.subsystems.setuptools import Setuptools
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
-from pants.backend.python.util_rules.pex import PexRequest, VenvPex, VenvPexProcess
+from pants.backend.python.util_rules.pex import Pex, PexRequest, VenvPex, VenvPexProcess
 from pants.backend.python.util_rules.pex import rules as pex_rules
 from pants.backend.python.util_rules.pex_requirements import EntireLockfile, PexRequirements
 from pants.base.glob_match_error_behavior import GlobMatchErrorBehavior
@@ -116,11 +116,14 @@ class DistBuildRequest:
     build_sdist: bool
     input: Digest
     working_directory: str  # Relpath within the input digest.
+    build_time_source_roots: tuple[str, ...]  # Source roots for 1st party build-time deps.
 
     target_address_spec: str | None = None  # Only needed for logging etc.
     wheel_config_settings: FrozenDict[str, tuple[str, ...]] | None = None
     sdist_config_settings: FrozenDict[str, tuple[str, ...]] | None = None
 
+    extra_build_time_requirements: tuple[Pex, ...] = tuple()
+
 
 @dataclass(frozen=True)
 class DistBuildResult:
@@ -191,6 +194,7 @@ async def run_pep517_build(request: DistBuildRequest, python_setup: PythonSetup)
             output_filename="build_backend.pex",
             internal_only=True,
             requirements=request.build_system.requires,
+            pex_path=request.extra_build_time_requirements,
             interpreter_constraints=request.interpreter_constraints,
         ),
     )
@@ -209,10 +213,11 @@ async def run_pep517_build(request: DistBuildRequest, python_setup: PythonSetup)
 
     merged_digest = await Get(Digest, MergeDigests((request.input, backend_shim_digest)))
 
+    extra_env = {
+        "PEX_EXTRA_SYS_PATH": os.pathsep.join(request.build_time_source_roots),
+    }
     if python_setup.macos_big_sur_compatibility and is_macos_big_sur():
-        extra_env = {"MACOSX_DEPLOYMENT_TARGET": "10.16"}
-    else:
-        extra_env = {}
+        extra_env["MACOSX_DEPLOYMENT_TARGET"] = "10.16"
 
     result = await Get(
         ProcessResult,