Improve support of MyPy requirements when Python 2 used (#10853)

### Problem

The wheels used by third-party code might not use the same constraints as what the MyPy Pex uses. This was causing errors (#10819), which we fixed in #10820.

However, that fix was not perfect. Even if there were valid typed Python 2 wheels, we would ignore them. WIth Python 3, we had to be careful that `mypy.pex` used the exact same constraints as `requirements.pex`, so that, for example, we don't run MyPy with Py36 but resolve requirements with Py37.

### Solution

Stop using `--pex-path` to join the `requirements.pex`, which has the effect of the wheels being put on `PYTHONPATH`. (We were then teaching MyPy to read from this PYTHONPATH through a custom launcher script.)

Instead, we extract the wheel folders from `requirements.pex`, and teach our custom launcher script to look in these extracted `.deps/` folders through a new `EXTRACTED_WHEELS` env var. MyPy will use this to resolve the requirements, without us actually including the wheels in the final PEX or modifying PYTHONPATH.

### Result

Our Python 2 test now fully works, even though one of the wheels is not compatible with Python 3.

We no longer need to align `mypy.pex` with `requirements.pex`.

#### Caveat: MyPy plugins must be installed with `--mypy-extra-requirements`

An initial implementation did not use `EXTRACTED_WHEELS` and instead loaded the wheels via `PYTHONPATH`. This allowed for you to specify MyPy plugins via normal dependencies on `python_requirement_library` targets, because that requirement would get loaded via PYTHONPATH. Now, you cannot do this; you must specify the requirement via `--mypy-extra-requirements`.

This caveat seems acceptable. For Pytest, it's convenient that we allow you to load plugins via normal requirements, as it allows you to only sometimes use the plugin. MyPy is much more uniform, however; MyPy runs in a single batch over the entire transitive closure, unlike Pytest running one test file at-a-time. Likewise, without Pants, it's hard to imagine how you could only sometimes use a MyPy plugin, given that MyPy only works with a single config file.

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

@@ -31,6 +31,7 @@
 from pants.backend.python.target_types import rules as target_type_rules
 from pants.backend.python.util_rules import (
     ancestor_files,
+    extract_pex,
     pex,
     pex_cli,
     pex_environment,
@@ -60,6 +61,7 @@ def rules():
     return (
         *coverage_py.rules(),
         *ancestor_files.rules(),
+        *extract_pex.rules(),
         *python_sources.rules(),
         *dependency_inference_rules.rules(),
         *pex.rules(),

src/python/pants/backend/python/typecheck/mypy/BUILD

@@ -6,5 +6,5 @@ python_library()
 python_integration_tests(
   name='integration',
   uses_pants_run=False,
-  timeout=300,
+  timeout=360,
 )

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

@@ -13,7 +13,8 @@
     PythonSources,
 )
 from pants.backend.python.typecheck.mypy.subsystem import MyPy
-from pants.backend.python.util_rules import pex_from_targets
+from pants.backend.python.util_rules import extract_pex, pex_from_targets
+from pants.backend.python.util_rules.extract_pex import ExtractedPexDistributions
 from pants.backend.python.util_rules.pex import (
     Pex,
     PexInterpreterConstraints,
@@ -70,14 +71,22 @@ def generate_args(mypy: MyPy, *, file_list_path: str) -> Tuple[str, ...]:
 # a sibling `<package>-stubs` package as per PEP-0561. Going further than that PEP, MyPy restricts
 # its search to `site-packages`. Since PEX deliberately isolates itself from `site-packages` as
 # part of its raison d'être, we monkey-patch `site.getsitepackages` to look inside the scrubbed
-# PEX sys.path before handing off to `mypy`.
+# PEX sys.path before handing off to `mypy`. This will find dependencies installed in the
+# `mypy.pex`, such as MyPy itself and any third-party plugins installed via
+# `--mypy-extra-requirements`.
 #
-# As a complication, MyPy does its own validation to ensure packages aren't both available in
-# site-packages and on the PYTHONPATH. As such, we elide all PYTHONPATH entries from artificial
-# site-packages we set up since MyPy will manually scan PYTHONPATH outside this PEX to find
-# packages. We also elide the values of PEX_EXTRA_SYS_PATH, which will be relative paths unlike
-# every other entry of sys.path. (We can't directly look for PEX_EXTRA_SYS_PATH because Pex scrubs
-# it.)
+# We also include the values from our custom env var `EXTRACTED_WHEELS` in this monkey-patch. For
+# user's third-party requirements, we don't include them in the `mypy.pex`, as the interpreter
+# constraints for their own code may be different than what's used to run MyPy, and this would
+# cause issues with Pex. Instead, we extract out the `.deps` folder from `requirements.pex`, and
+# set the env var `EXTRACTED_WHEELS` to point to each entry. This allows MyPy to know about user's
+# third-party requirements without having to set them on PYTHONPATH.
+#
+# Finally, we elide the values of PEX_EXTRA_SYS_PATH, which will point to user's first-party code's
+# source roots. MyPy validates that the same paths are not available both in site-packages and
+# PYTHONPATH, so we must not add this first-party code to site-packages. We use a heuristic of
+# looking for relative paths, as all other entries will be absolute paths. (We can't directly look
+# for PEX_EXTRA_SYS_PATH because Pex scrubs it.)
 #
 # See:
 #   https://mypy.readthedocs.io/en/stable/installed_packages.html#installed-packages
@@ -92,13 +101,9 @@ def generate_args(mypy: MyPy, *, file_list_path: str) -> Tuple[str, ...]:
         import site
         import sys
 
-        PYTHONPATH = frozenset(
-            os.path.realpath(p)
-            for p in os.environ.get('PYTHONPATH', '').split(os.pathsep)
-        )
         site.getsitepackages = lambda: [
-            p for p in sys.path
-            if os.path.realpath(p) not in PYTHONPATH and os.path.isabs(p)
+            *(p for p in sys.path if os.path.isabs(p)),
+            *os.environ.get('EXTRACTED_WHEELS').split(os.pathsep),
         ]
         site.getusersitepackages = lambda: ''  # i.e, the CWD.
 
@@ -149,16 +154,14 @@ async def mypy_typecheck(
     #     `python_version`. We don't want to make the user set this up. (If they do, MyPy will use
     #     `python_version`, rather than defaulting to the executing interpreter).
     #
-    #  * We must resolve third-party dependencies. This should use whatever the actual code's
-    #     constraints are. However, PEX will error if they are not compatible
-    #     with the interpreter being used to run MyPy. So, we should generally align the tool PEX
-    #     constraints with the requirements constraints.
+    #  * When resolving third-party requirements, we should use the actual requirements. Normally,
+    #     we would merge the requirements.pex with mypy.pex via `--pex-path`. However, this will
+    #     cause a runtime error if the interpreter constraints are different between the PEXes and
+    #     they have incompatible wheels.
     #
-    #     However, it's possible for the requirements' constraints to include Python 2 and not be
-    #     compatible with MyPy's >=3.5 requirement. If any of the requirements only have
-    #     Python 2 wheels and they are not compatible with Python 3, then Pex will error about
-    #     missing wheels. So, in this case, we set `PEX_IGNORE_ERRORS`, which will avoid erroring,
-    #     but may result in MyPy complaining that it cannot find certain distributions.
+    #     Instead, we teach MyPy about the requirements by extracting the distributions from
+    #     requirements.pex and setting EXTRACTED_WHEELS, which our custom launcher script then
+    #     looks for.
     code_interpreter_constraints = PexInterpreterConstraints.create_from_compatibility_fields(
         (
             tgt[PythonInterpreterCompatibility]
@@ -167,15 +170,17 @@ async def mypy_typecheck(
         ),
         python_setup,
     )
-    use_subsystem_constraints = (
-        not mypy.options.is_default("interpreter_constraints")
-        or code_interpreter_constraints.includes_python2()
-    )
-    tool_interpreter_constraints = (
-        PexInterpreterConstraints(mypy.interpreter_constraints)
-        if use_subsystem_constraints
-        else code_interpreter_constraints
-    )
+
+    if not mypy.options.is_default("interpreter_constraints"):
+        tool_interpreter_constraints = mypy.interpreter_constraints
+    elif code_interpreter_constraints.requires_python38_or_newer():
+        tool_interpreter_constraints = ("CPython>=3.8",)
+    elif code_interpreter_constraints.requires_python37_or_newer():
+        tool_interpreter_constraints = ("CPython>=3.7",)
+    elif code_interpreter_constraints.requires_python36_or_newer():
+        tool_interpreter_constraints = ("CPython>=3.6",)
+    else:
+        tool_interpreter_constraints = mypy.interpreter_constraints
 
     plugin_sources_request = Get(
         PythonSourceFiles, PythonSourceFilesRequest(plugin_transitive_targets.closure)
@@ -202,9 +207,8 @@ async def mypy_typecheck(
             requirements=PexRequirements(
                 itertools.chain(mypy.all_requirements, plugin_requirements)
             ),
-            interpreter_constraints=tool_interpreter_constraints,
+            interpreter_constraints=PexInterpreterConstraints(tool_interpreter_constraints),
             entry_point=PurePath(LAUNCHER_FILE.path).stem,
-            additional_args=("--pex-path", requirements_pex_request.input.output_filename),
         ),
     )
 
@@ -236,11 +240,15 @@ async def mypy_typecheck(
     python_files = "\n".join(
         f for f in typechecked_sources.source_files.snapshot.files if f.endswith(".py")
     )
-    file_list_digest = await Get(
+    create_file_list_request = Get(
         Digest,
         CreateDigest([FileContent(file_list_path, python_files.encode())]),
     )
 
+    file_list_digest, extracted_pex_distributions = await MultiGet(
+        create_file_list_request, Get(ExtractedPexDistributions, Pex, requirements_pex)
+    )
+
     merged_input_files = await Get(
         Digest,
         MergeDigests(
@@ -249,7 +257,7 @@ async def mypy_typecheck(
                 plugin_sources.source_files.snapshot.digest,
                 typechecked_srcs_snapshot.digest,
                 mypy_pex.digest,
-                requirements_pex.digest,
+                extracted_pex_distributions.digest,
                 config_digest,
             ]
         ),
@@ -258,19 +266,18 @@ async def mypy_typecheck(
     all_used_source_roots = sorted(
         set(itertools.chain(plugin_sources.source_roots, typechecked_sources.source_roots))
     )
-    extra_env = {"PEX_EXTRA_SYS_PATH": ":".join(all_used_source_roots)}
-    # If the constraints are different for the tool than for the requirements, we must tell Pex to
-    # ignore errors. Otherwise, we risk runtime errors about missing dependencies.
-    if code_interpreter_constraints != tool_interpreter_constraints:
-        extra_env["PEX_IGNORE_ERRORS"] = "true"
+    env = {
+        "PEX_EXTRA_SYS_PATH": ":".join(all_used_source_roots),
+        "EXTRACTED_WHEELS": ":".join(extracted_pex_distributions.wheel_directory_paths),
+    }
 
     result = await Get(
         FallibleProcessResult,
         PexProcess(
             mypy_pex,
             argv=generate_args(mypy, file_list_path=file_list_path),
             input_digest=merged_input_files,
-            extra_env=extra_env,
+            extra_env=env,
             description=f"Run MyPy on {pluralize(len(typechecked_srcs_snapshot.files), 'file')}.",
             level=LogLevel.DEBUG,
         ),
@@ -284,6 +291,7 @@ def rules():
     return [
         *collect_rules(),
         UnionRule(TypecheckRequest, MyPyRequest),
+        *extract_pex.rules(),
         *pants_bin.rules(),
         *pex_from_targets.rules(),
     ]

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

@@ -76,7 +76,7 @@ def add(x: int, y: int) -> int:
 GLOBAL_ARGS = (
     "--backend-packages=pants.backend.python",
     "--backend-packages=pants.backend.python.typecheck.mypy",
-    "--source-root-patterns=['src/python', 'tests/python']",
+    "--source-root-patterns=['/', 'src/python', 'tests/python']",
 )
 
 
@@ -355,19 +355,14 @@ def test_works_with_python27(rule_runner: RuleRunner) -> None:
 
     There was a bug that this would cause the runner PEX to fail to execute because it did not have
     Python 3 distributions of the requirements.
-
-    TODO(#10819): This support is not as robust as we'd like. We'll only use third-party
-    distributions if its wheel is also compatible with the Python 3 interpreter being used to run
-    MyPy. Is it possible to fix this to always include the Python 2 wheels and have MyPy respect
-    them?
     """
     rule_runner.add_to_build_file(
         "",
         dedent(
             """\
             # Both requirements are a) typed and b) compatible with Py2 and Py3. However, `x690`
-            # has a distinct wheel for Py2 vs. Py3, whereas libumi has a universal wheel. We only
-            # expect libumi to be usable by MyPy.
+            # has a distinct wheel for Py2 vs. Py3, whereas libumi has a universal wheel. We expect
+            # both to be usable, even though libumi is not compatible with Py3.
 
             python_requirement_library(
                 name="libumi",
@@ -398,7 +393,13 @@ def test_works_with_python27(rule_runner: RuleRunner) -> None:
     assert len(result) == 1
     assert result[0].exit_code == 1
     assert "Failed to execute PEX file" not in result[0].stderr
-    assert "Cannot find implementation or library stub for module named 'x690'" in result[0].stdout
+    assert (
+        "Cannot find implementation or library stub for module named 'x690'" not in result[0].stdout
+    )
+    assert (
+        "Cannot find implementation or library stub for module named 'libumi'"
+        not in result[0].stdout
+    )
     assert f"{PACKAGE}/py2.py:5: error: Unsupported operand types" in result[0].stdout
 
 
@@ -424,6 +425,33 @@ def test_works_with_python38(rule_runner: RuleRunner) -> None:
     assert "Success: no issues found" in result[0].stdout.strip()
 
 
+def test_mypy_shadows_requirements(rule_runner: RuleRunner) -> None:
+    """Test the behavior of a MyPy requirement shadowing a user's requirement.
+
+    The way we load requirements is complex. We want to ensure that things still work properly in
+    this edge case.
+    """
+    rule_runner.create_file("app.py", "import typed_ast\n")
+    rule_runner.add_to_build_file(
+        "",
+        dedent(
+            """\
+            python_requirement_library(
+                name='typed-ast',
+                requirements=['typed-ast==1.4.1'],
+            )
+
+            python_library(name="lib")
+            """
+        ),
+    )
+    tgt = rule_runner.get_target(Address("", target_name="lib"))
+    result = run_mypy(rule_runner, [tgt], additional_args=["--mypy-version=mypy==0.782"])
+    assert len(result) == 1
+    assert result[0].exit_code == 0
+    assert "Success: no issues found" in result[0].stdout
+
+
 def test_source_plugin(rule_runner: RuleRunner) -> None:
     # NB: We make this source plugin fairly complex by having it use transitive dependencies.
     # This is to ensure that we can correctly support plugins with dependencies.

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

@@ -0,0 +1,42 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from dataclasses import dataclass
+from typing import Tuple
+
+from pants.backend.python.util_rules.pex import Pex
+from pants.core.util_rules.archive import UnzipBinary
+from pants.engine.fs import Digest, Snapshot
+from pants.engine.process import FallibleProcessResult, Process
+from pants.engine.rules import Get, collect_rules, rule
+from pants.util.logging import LogLevel
+
+
+@dataclass(frozen=True)
+class ExtractedPexDistributions:
+    digest: Digest
+    wheel_directory_paths: Tuple[str, ...]
+
+
+@rule
+async def extract_distributions(pex: Pex, unzip_binary: UnzipBinary) -> ExtractedPexDistributions:
+    # We only unzip the `.deps` folder to avoid unnecessary work. Note that this will cause the
+    # process to fail if there is no `.deps` folder, so we need to use `FallibleProcessResult`.
+    argv = (*unzip_binary.extract_argv(archive_path=pex.name, output_dir="."), ".deps/*")
+    unzipped_pex = await Get(
+        FallibleProcessResult,
+        Process(
+            argv=argv,
+            input_digest=pex.digest,
+            output_directories=(".deps",),
+            description=f"Unzip {pex.name} to determine its distributions",
+            level=LogLevel.DEBUG,
+        ),
+    )
+    snapshot = await Get(Snapshot, Digest, unzipped_pex.output_digest)
+    directory_paths = tuple(sorted(d for d in snapshot.dirs if d.endswith(".whl")))
+    return ExtractedPexDistributions(snapshot.digest, directory_paths)
+
+
+def rules():
+    return collect_rules()