Initial support for building dists using mypyc.

src/python/pants/backend/python/goals/BUILD

@@ -52,6 +52,8 @@ python_tests(
     name="setup_py_integration",
     sources=["setup_py_integration_test.py"],
     dependencies=["testprojects/src/python:native_directory"],
+    # We want to make sure the native builds work for both macOS and Linux.
+    tags=["platform_specific_behavior"],
     timeout=180,
 )
 

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.get(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/goals/setup_py_integration_test.py

@@ -12,10 +12,13 @@
 from tempfile import TemporaryDirectory
 from textwrap import dedent
 
+import pytest
+
 from pants.testutil.pants_integration_test import run_pants, setup_tmpdir
 from pants.util.dirutil import safe_rmtree
 
 
+@pytest.mark.platform_specific_behavior
 def test_native_code() -> None:
     dist_dir = "dist"
     pyver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"

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.*"
@@ -40,7 +40,7 @@ def assert_ics(build_file: str, expected: list[str]) -> None:
         assert lockfile_request.interpreter_constraints == InterpreterConstraints(expected)
 
     # If no dependencies for python_distribution, fall back to global [python] constraints.
-    assert_ics("python_distribution(provides=setup_py(name='dist'))", [global_constraint])
+    assert_ics("python_distribution(provides=python_artifact(name='dist'))", [global_constraint])
 
     assert_ics(
         dedent(
@@ -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/BUILD

@@ -12,3 +12,11 @@ python_tests(
     # We want to make sure the default lockfile works for both macOS and Linux.
     tags=["platform_specific_behavior"],
 )
+python_tests(
+    name="mypyc_integration_test",
+    sources=["mypyc_integration_test.py"],
+    dependencies=["testprojects/src/python:mypyc_fib_directory"],
+    timeout=120,
+    # We want to make sure we can build with mypyc on both macOS and Linux.
+    tags=["platform_specific_behavior"],
+)

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

@@ -0,0 +1,96 @@
+# Copyright 2022 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
+from pants.util.strutil import softwrap
+
+
+class UsesMyPycField(BoolField):
+    alias = "uses_mypyc"
+    default = False
+    help = softwrap(
+        """
+        If true, this distribution is built using mypyc.
+
+        In this case, Pants will build the distribution in an environment that includes
+        mypy, as configured in the `[mypy]` subsystem, including plugins, config files,
+        extra type stubs, and the distribution's own requirements (which normally would not
+        be needed at build time, but in this case may provide necessary type annotations).
+
+        You will typically set this field on distributions whose setup.py uses
+        mypyc.build.mypycify(). See https://mypyc.readthedocs.io/en/latest/index.html .
+        """
+    )
+
+
+@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(),
+    ]