diff --git a/src/python/pants/backend/codegen/protobuf/python/rules.py b/src/python/pants/backend/codegen/protobuf/python/rules.py index 026ad65521f..30b988b7c6e 100644 --- a/src/python/pants/backend/codegen/protobuf/python/rules.py +++ b/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( diff --git a/src/python/pants/backend/codegen/protobuf/python/rules_integration_test.py b/src/python/pants/backend/codegen/protobuf/python/rules_integration_test.py index bf6403c86e3..d983dff97da 100644 --- a/src/python/pants/backend/codegen/protobuf/python/rules_integration_test.py +++ b/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,6 +61,7 @@ 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", @@ -46,6 +69,8 @@ def assert_files_generated( ] 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", + ], )