Add support for gRPC mypy stub files (#11658)

### Problem

Pants currently supports generating mypy stubs for the generated protobuf code, which makes mypy very happy. However, when type checking a codebase that has generated gRPC code as well, mypy throws quite the hissy fit due to the lack of gRPC stubs.

### Solution

mypy-protobuf can ([since version 2.0](https://github.com/dropbox/mypy-protobuf/blob/master/CHANGELOG.md#20) at least) generate gRPC stubs as well, and it can be used in Pants with some minor changes. It will, of course, only work on mypy-protobuf 2.0+.

### Result

Stubs are now generated for the generated gRPC code as well when mypy-protobuf >= 2 is used! Awesome!

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

@@ -14,6 +14,7 @@
     PexInterpreterConstraints,
     PexRequest,
     PexRequirements,
+    PexResolveInfo,
     VenvPex,
     VenvPexRequest,
 )
@@ -96,27 +97,47 @@ async def generate_python_from_protobuf(
     )
 
     protoc_gen_mypy_script = "protoc-gen-mypy"
+    protoc_gen_mypy_grpc_script = "protoc-gen-mypy_grpc"
     mypy_pex = None
+    mypy_request = PexRequest(
+        output_filename="mypy_protobuf.pex",
+        internal_only=True,
+        requirements=PexRequirements([python_protobuf_subsystem.mypy_plugin_version]),
+        # TODO(John Sirois): Fix these interpreter constraints to track the actual
+        #  python requirement of the mypy_plugin_version or else plumb an option for
+        #  manually setting the constraint to track what mypy_plugin_version needs:
+        #  https://github.com/pantsbuild/pants/issues/11565
+        # Here we guess a constraint that will likely work with any mypy_plugin_version
+        # selected.
+        interpreter_constraints=PexInterpreterConstraints(["CPython>=3.5"]),
+    )
+
     if python_protobuf_subsystem.mypy_plugin:
         mypy_pex = await Get(
             VenvPex,
             VenvPexRequest(
                 bin_names=[protoc_gen_mypy_script],
-                pex_request=PexRequest(
-                    output_filename="mypy_protobuf.pex",
-                    internal_only=True,
-                    requirements=PexRequirements([python_protobuf_subsystem.mypy_plugin_version]),
-                    # TODO(John Sirois): Fix these interpreter constraints to track the actual
-                    #  python requirement of the mypy_plugin_version or else plumb an option for
-                    #  manually setting the constraint to track what mypy_plugin_version needs:
-                    #  https://github.com/pantsbuild/pants/issues/11565
-                    # Here we guess a constraint that will likely work with any mypy_plugin_version
-                    # selected.
-                    interpreter_constraints=PexInterpreterConstraints(["CPython>=3.5"]),
-                ),
+                pex_request=mypy_request,
             ),
         )
 
+        if request.protocol_target.get(ProtobufGrpcToggle).value:
+            mypy_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
+            ):
+                # TODO: Use `pex_path` once VenvPex stores a Pex field.
+                mypy_pex = await Get(
+                    VenvPex,
+                    VenvPexRequest(
+                        bin_names=[protoc_gen_mypy_script, protoc_gen_mypy_grpc_script],
+                        pex_request=mypy_request,
+                    ),
+                )
+
     downloaded_grpc_plugin = (
         await Get(
             DownloadedExternalTool,
@@ -151,8 +172,17 @@ async def generate_python_from_protobuf(
         argv.extend(
             [f"--plugin=protoc-gen-grpc={downloaded_grpc_plugin.exe}", "--grpc_out", output_dir]
         )
-    argv.extend(target_sources_stripped.snapshot.files)
 
+        if mypy_pex and protoc_gen_mypy_grpc_script in mypy_pex.bin:
+            argv.extend(
+                [
+                    f"--plugin=protoc-gen-mypy_grpc={mypy_pex.bin[protoc_gen_mypy_grpc_script].argv0}",
+                    "--mypy_grpc_out",
+                    output_dir,
+                ]
+            )
+
+    argv.extend(target_sources_stripped.snapshot.files)
     result = await Get(
         ProcessResult,
         Process(

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

@@ -2,7 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 from textwrap import dedent
-from typing import List
+from typing import List, Optional
 
 import pytest
 
@@ -17,6 +17,28 @@
 from pants.source.source_root import NoSourceRootError
 from pants.testutil.rule_runner import QueryRule, RuleRunner
 
+GRPC_PROTO_STANZA = """
+syntax = "proto3";
+
+package dir1;
+
+// The greeter service definition.
+service Greeter {
+  // Sends a greeting
+  rpc SayHello (HelloRequest) returns (HelloReply) {}
+}
+
+// The request message containing the user's name.
+message HelloRequest {
+  string name = 1;
+}
+
+// The response message containing the greetings
+message HelloReply {
+  string message = 1;
+}
+"""
+
 
 @pytest.fixture
 def rule_runner() -> RuleRunner:
@@ -39,13 +61,16 @@ def assert_files_generated(
     expected_files: List[str],
     source_roots: List[str],
     mypy: bool = False,
+    mypy_plugin_version: Optional[str] = None,
 ) -> None:
     options = [
         "--backend-packages=pants.backend.codegen.protobuf.python",
         f"--source-root-patterns={repr(source_roots)}",
     ]
     if mypy:
         options.append("--python-protobuf-mypy-plugin")
+    if mypy_plugin_version:
+        options.append(f"--python-protobuf-mypy-plugin-version={mypy_plugin_version}")
     rule_runner.set_options(
         options,
         env_inherit={"PATH", "PYENV_ROOT", "HOME"},
@@ -240,34 +265,53 @@ def test_mypy_plugin(rule_runner: RuleRunner) -> None:
 def test_grpc(rule_runner: RuleRunner) -> None:
     rule_runner.create_file(
         "src/protobuf/dir1/f.proto",
-        dedent(
-            """\
-            syntax = "proto3";
+        dedent(GRPC_PROTO_STANZA),
+    )
+    rule_runner.add_to_build_file("src/protobuf/dir1", "protobuf_library(grpc=True)")
+    assert_files_generated(
+        rule_runner,
+        "src/protobuf/dir1",
+        source_roots=["src/protobuf"],
+        expected_files=["src/protobuf/dir1/f_pb2.py", "src/protobuf/dir1/f_pb2_grpc.py"],
+    )
 
-            package dir1;
 
-            // The greeter service definition.
-            service Greeter {
-              // Sends a greeting
-              rpc SayHello (HelloRequest) returns (HelloReply) {}
-            }
+def test_grpc_mypy_plugin(rule_runner: RuleRunner) -> None:
+    rule_runner.create_file(
+        "src/protobuf/dir1/f.proto",
+        dedent(GRPC_PROTO_STANZA),
+    )
+    rule_runner.add_to_build_file("src/protobuf/dir1", "protobuf_library(grpc=True)")
+    assert_files_generated(
+        rule_runner,
+        "src/protobuf/dir1",
+        source_roots=["src/protobuf"],
+        mypy=True,
+        expected_files=[
+            "src/protobuf/dir1/f_pb2.py",
+            "src/protobuf/dir1/f_pb2.pyi",
+            "src/protobuf/dir1/f_pb2_grpc.py",
+            "src/protobuf/dir1/f_pb2_grpc.pyi",
+        ],
+    )
 
-            // The request message containing the user's name.
-            message HelloRequest {
-              string name = 1;
-            }
 
-            // The response message containing the greetings
-            message HelloReply {
-              string message = 1;
-            }
-            """
-        ),
+def test_grpc_pre_v2_mypy_plugin(rule_runner: RuleRunner) -> None:
+    rule_runner.create_file(
+        "src/protobuf/dir1/f.proto",
+        dedent(GRPC_PROTO_STANZA),
     )
     rule_runner.add_to_build_file("src/protobuf/dir1", "protobuf_library(grpc=True)")
+
     assert_files_generated(
         rule_runner,
         "src/protobuf/dir1",
         source_roots=["src/protobuf"],
-        expected_files=["src/protobuf/dir1/f_pb2.py", "src/protobuf/dir1/f_pb2_grpc.py"],
+        mypy=True,
+        mypy_plugin_version="mypy-protobuf==1.24",
+        expected_files=[
+            "src/protobuf/dir1/f_pb2.py",
+            "src/protobuf/dir1/f_pb2.pyi",
+            "src/protobuf/dir1/f_pb2_grpc.py",
+        ],
     )