Improve mypy runs by adding the "incremental" mypy cache (#16276)

Leverages mypy's cache to allow for "incremental" runs. The cache is an `append_only_cache` which is a bit disingenuous here, but AFAICT mypy is process-safe w.r.t. the cache (I'll do more digging to be very sure).

Fixes #10864

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/backend/codegen/protobuf/python/rules.py

@@ -107,13 +107,11 @@ async def generate_python_from_protobuf(
         )
 
         if request.protocol_target.get(ProtobufGrpcToggleField).value:
-            mypy_info = await Get(PexResolveInfo, VenvPex, mypy_pex)
+            mypy_pex_info = await Get(PexResolveInfo, VenvPex, mypy_pex)
 
             # In order to generate stubs for gRPC code, we need mypy-protobuf 2.0 or above.
-            if any(
-                dist_info.project_name == "mypy-protobuf" and dist_info.version.major >= 2
-                for dist_info in mypy_info
-            ):
+            mypy_protobuf_info = mypy_pex_info.find("mypy-protobuf")
+            if mypy_protobuf_info and mypy_protobuf_info.version.major >= 2:
                 # TODO: Use `pex_path` once VenvPex stores a Pex field.
                 mypy_pex = await Get(
                     VenvPex,

src/python/pants/backend/python/lint/isort/rules.py

@@ -74,11 +74,9 @@ async def isort_fmt(request: IsortRequest, isort: Isort) -> FmtResult:
     # Isort 5+ changes how config files are handled. Determine which semantics we should use.
     is_isort5 = False
     if isort.config:
-        isort_info = await Get(PexResolveInfo, VenvPex, isort_pex)
-        is_isort5 = any(
-            dist_info.project_name == "isort" and dist_info.version.major >= 5
-            for dist_info in isort_info
-        )
+        isort_pex_info = await Get(PexResolveInfo, VenvPex, isort_pex)
+        isort_info = isort_pex_info.find("isort")
+        is_isort5 = isort_info is not None and isort_info.version.major >= 5
 
     input_digest = await Get(
         Digest, MergeDigests((request.snapshot.digest, config_files.snapshot.digest))

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

@@ -5,8 +5,11 @@
 
 import itertools
 from dataclasses import dataclass
+from hashlib import sha256
 from typing import Iterable, Optional, Tuple
 
+import packaging
+
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import PythonSourceField
 from pants.backend.python.typecheck.mypy.skip_field import SkipMyPyField
@@ -17,19 +20,26 @@
 )
 from pants.backend.python.util_rules import partition, pex_from_targets
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
-from pants.backend.python.util_rules.pex import Pex, PexRequest, VenvPex, VenvPexProcess
+from pants.backend.python.util_rules.pex import (
+    Pex,
+    PexRequest,
+    PexResolveInfo,
+    VenvPex,
+    VenvPexProcess,
+)
 from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
 from pants.backend.python.util_rules.pex_requirements import PexRequirements
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFiles,
     PythonSourceFilesRequest,
 )
+from pants.base.build_root import BuildRoot
 from pants.core.goals.check import REPORT_DIR, CheckRequest, CheckResult, CheckResults
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.collection import Collection
 from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests, RemovePrefix
 from pants.engine.process import FallibleProcessResult
-from pants.engine.rules import Get, MultiGet, collect_rules, rule
+from pants.engine.rules import Get, MultiGet, collect_rules, rule, rule_helper
 from pants.engine.target import CoarsenedTargets, FieldSet, Target
 from pants.engine.unions import UnionRule
 from pants.util.logging import LogLevel
@@ -69,9 +79,12 @@ class MyPyRequest(CheckRequest):
     name = MyPy.options_scope
 
 
-def generate_argv(
+@rule_helper
+async def _generate_argv(
     mypy: MyPy,
     *,
+    pex: VenvPex,
+    cache_dir: str,
     venv_python: str,
     file_list_path: str,
     python_version: Optional[str],
@@ -81,6 +94,16 @@ def generate_argv(
         args.append(f"--config-file={mypy.config}")
     if python_version:
         args.append(f"--python-version={python_version}")
+
+    mypy_pex_info = await Get(PexResolveInfo, VenvPex, pex)
+    mypy_info = mypy_pex_info.find("mypy")
+    assert mypy_info is not None
+    if mypy_info.version > packaging.version.Version("0.700"):
+        args.append("--skip-cache-mtime-check")
+        args.extend(("--cache-dir", cache_dir))
+    else:
+        # Don't bother caching
+        args.append("--cache-dir=/dev/null")
     args.append(f"@{file_list_path}")
     return tuple(args)
 
@@ -109,6 +132,7 @@ async def mypy_typecheck_partition(
     partition: MyPyPartition,
     config_file: MyPyConfigFile,
     first_party_plugins: MyPyFirstPartyPlugins,
+    build_root: BuildRoot,
     mypy: MyPy,
     python_setup: PythonSetup,
 ) -> CheckResult:
@@ -224,13 +248,17 @@ async def mypy_typecheck_partition(
         "MYPYPATH": ":".join(all_used_source_roots),
     }
 
+    cache_dir = f".cache/mypy_cache/{sha256(build_root.path.encode()).hexdigest()}"
+
     result = await Get(
         FallibleProcessResult,
         VenvPexProcess(
             mypy_pex,
-            argv=generate_argv(
+            argv=await _generate_argv(
                 mypy,
+                pex=mypy_pex,
                 venv_python=requirements_venv_pex.python.argv0,
+                cache_dir=cache_dir,
                 file_list_path=file_list_path,
                 python_version=config_file.python_version_to_autoset(
                     partition.interpreter_constraints, python_setup.interpreter_versions_universe
@@ -241,6 +269,7 @@ async def mypy_typecheck_partition(
             output_directories=(REPORT_DIR,),
             description=f"Run MyPy on {pluralize(len(python_files), 'file')}.",
             level=LogLevel.DEBUG,
+            append_only_caches={"mypy_cache": cache_dir},
         ),
     )
     report = await Get(Digest, RemovePrefix(result.output_digest, REPORT_DIR))

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

@@ -11,7 +11,7 @@
 from dataclasses import dataclass
 from pathlib import PurePath
 from textwrap import dedent
-from typing import Iterable, Iterator, Mapping
+from typing import Iterable, Iterator, Mapping, TypeVar
 
 import packaging.specifiers
 import packaging.version
@@ -949,6 +949,7 @@ class VenvPexProcess:
     execution_slot_variable: str | None
     concurrency_available: int
     cache_scope: ProcessCacheScope
+    append_only_caches: FrozenDict[str, str]
 
     def __init__(
         self,
@@ -966,6 +967,7 @@ def __init__(
         execution_slot_variable: str | None = None,
         concurrency_available: int = 0,
         cache_scope: ProcessCacheScope = ProcessCacheScope.SUCCESSFUL,
+        append_only_caches: Mapping[str, str] | None = None,
     ) -> None:
         self.venv_pex = venv_pex
         self.argv = tuple(argv)
@@ -980,6 +982,7 @@ def __init__(
         self.execution_slot_variable = execution_slot_variable
         self.concurrency_available = concurrency_available
         self.cache_scope = cache_scope
+        self.append_only_caches = FrozenDict(append_only_caches or {})
 
 
 @rule
@@ -998,6 +1001,12 @@ async def setup_venv_pex_process(
         if request.input_digest
         else venv_pex.digest
     )
+    append_only_caches: FrozenDict[str, str] = FrozenDict(
+        **pex_environment.in_sandbox(
+            working_directory=request.working_directory
+        ).append_only_caches,
+        **request.append_only_caches,
+    )
     return Process(
         argv=argv,
         description=request.description,
@@ -1007,9 +1016,7 @@ async def setup_venv_pex_process(
         env=request.extra_env,
         output_files=request.output_files,
         output_directories=request.output_directories,
-        append_only_caches=pex_environment.in_sandbox(
-            working_directory=request.working_directory
-        ).append_only_caches,
+        append_only_caches=append_only_caches,
         timeout_seconds=request.timeout_seconds,
         execution_slot_variable=request.execution_slot_variable,
         concurrency_available=request.concurrency_available,
@@ -1030,10 +1037,22 @@ class PexDistributionInfo:
     requires_dists: tuple[Requirement, ...]
 
 
+DefaultT = TypeVar("DefaultT")
+
+
 class PexResolveInfo(Collection[PexDistributionInfo]):
     """Information about all distributions resolved in a PEX file, as reported by `PEX_TOOLS=1
     repository info -v`."""
 
+    def find(
+        self, name: str, default: DefaultT | None = None
+    ) -> PexDistributionInfo | DefaultT | None:
+        """Returns the PexDistributionInfo with the given name, first one wins."""
+        try:
+            return next(info for info in self if info.project_name == name)
+        except StopIteration:
+            return default
+
 
 def parse_repository_info(repository_info: str) -> PexResolveInfo:
     def iter_dist_info() -> Iterator[PexDistributionInfo]: