Add dependency inference for Python imports of Protobuf/gRPC (pantsbu…

…ild#11195)

Closes pantsbuild#11184. Protobuf and gRPC are very predictable with the modules they generate. Further, we strip source roots, so we don't need to worry about where the files are located or the `python_source_root` field.

The only challenge is exposing an entry point for `[python-infer].imports`. We add a hook for different implementations to contribute to the final merged first-party module mapping. We use this hook with our original implementation so that it runs in parallel with all plugin implementations.

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

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

@@ -0,0 +1,69 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from typing import Dict, Set, Tuple
+
+from pants.backend.codegen.protobuf.target_types import ProtobufGrpcToggle, ProtobufSources
+from pants.backend.python.dependency_inference.module_mapper import (
+    FirstPartyPythonMappingImpl,
+    FirstPartyPythonMappingImplMarker,
+)
+from pants.base.specs import AddressSpecs, DescendantAddresses
+from pants.core.util_rules.stripped_source_files import StrippedSourceFileNames
+from pants.engine.addresses import Address
+from pants.engine.rules import Get, MultiGet, collect_rules, rule
+from pants.engine.target import SourcesPathsRequest, Target, Targets
+from pants.engine.unions import UnionRule
+from pants.util.logging import LogLevel
+
+
+def proto_path_to_py_module(stripped_path: str, *, suffix: str) -> str:
+    return stripped_path.replace(".proto", suffix).replace("/", ".")
+
+
+# This is only used to register our implementation with the plugin hook via unions.
+class PythonProtobufMappingMarker(FirstPartyPythonMappingImplMarker):
+    pass
+
+
+@rule(desc="Creating map of Protobuf targets to generated Python modules", level=LogLevel.DEBUG)
+async def map_protobuf_to_python_modules(
+    _: PythonProtobufMappingMarker,
+) -> FirstPartyPythonMappingImpl:
+    all_expanded_targets = await Get(Targets, AddressSpecs([DescendantAddresses("")]))
+    protobuf_targets = tuple(tgt for tgt in all_expanded_targets if tgt.has_field(ProtobufSources))
+    stripped_sources_per_target = await MultiGet(
+        Get(StrippedSourceFileNames, SourcesPathsRequest(tgt[ProtobufSources]))
+        for tgt in protobuf_targets
+    )
+
+    modules_to_addresses: Dict[str, Tuple[Address]] = {}
+    modules_with_multiple_owners: Set[str] = set()
+
+    def add_module(module: str, tgt: Target) -> None:
+        if module in modules_to_addresses:
+            modules_with_multiple_owners.add(module)
+        else:
+            modules_to_addresses[module] = (tgt.address,)
+
+    for tgt, stripped_sources in zip(protobuf_targets, stripped_sources_per_target):
+        for stripped_f in stripped_sources:
+            # NB: We don't consider the MyPy plugin, which generates `_pb2.pyi`. The stubs end up
+            # sharing the same module as the implementation `_pb2.py`. Because both generated files
+            # come from the same original Protobuf target, we're covered.
+            add_module(proto_path_to_py_module(stripped_f, suffix="_pb2"), tgt)
+            if tgt.get(ProtobufGrpcToggle).value:
+                add_module(proto_path_to_py_module(stripped_f, suffix="_pb2_grpc"), tgt)
+
+    # Remove modules with ambiguous owners.
+    for ambiguous_module in modules_with_multiple_owners:
+        modules_to_addresses.pop(ambiguous_module)
+
+    return FirstPartyPythonMappingImpl(modules_to_addresses)
+
+
+def rules():
+    return (
+        *collect_rules(),
+        UnionRule(FirstPartyPythonMappingImplMarker, PythonProtobufMappingMarker),
+    )

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

@@ -0,0 +1,58 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import pytest
+
+from pants.backend.codegen.protobuf.python import additional_fields, python_protobuf_module_mapper
+from pants.backend.codegen.protobuf.python.python_protobuf_module_mapper import (
+    PythonProtobufMappingMarker,
+)
+from pants.backend.codegen.protobuf.target_types import ProtobufLibrary
+from pants.backend.python.dependency_inference.module_mapper import FirstPartyPythonMappingImpl
+from pants.core.util_rules import stripped_source_files
+from pants.engine.addresses import Address
+from pants.testutil.rule_runner import QueryRule, RuleRunner
+
+
+@pytest.fixture
+def rule_runner() -> RuleRunner:
+    return RuleRunner(
+        rules=[
+            *additional_fields.rules(),
+            *stripped_source_files.rules(),
+            *python_protobuf_module_mapper.rules(),
+            QueryRule(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker]),
+        ],
+        target_types=[ProtobufLibrary],
+    )
+
+
+def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
+    rule_runner.set_options(["--source-root-patterns=['root1', 'root2', 'root3']"])
+
+    # Two proto files belonging to the same target. We should use two file addresses.
+    rule_runner.create_files("root1/protos", ["f1.proto", "f2.proto"])
+    rule_runner.add_to_build_file("root1/protos", "protobuf_library()")
+
+    # These protos would result in the same module name, so neither should be used.
+    rule_runner.create_file("root1/two_owners/f.proto")
+    rule_runner.add_to_build_file("root1/two_owners", "protobuf_library()")
+    rule_runner.create_file("root2/two_owners/f.proto")
+    rule_runner.add_to_build_file("root2/two_owners", "protobuf_library()")
+
+    # A file with grpc. This also uses the `python_source_root` mechanism, which should be
+    # irrelevant to the module mapping because we strip source roots.
+    rule_runner.create_file("root1/tests/f.proto")
+    rule_runner.add_to_build_file(
+        "root1/tests", "protobuf_library(grpc=True, python_source_root='root3')"
+    )
+
+    result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()])
+    assert result == FirstPartyPythonMappingImpl(
+        {
+            "protos.f1_pb2": (Address("root1/protos", relative_file_path="f1.proto"),),
+            "protos.f2_pb2": (Address("root1/protos", relative_file_path="f2.proto"),),
+            "tests.f_pb2": (Address("root1/tests", relative_file_path="f.proto"),),
+            "tests.f_pb2_grpc": (Address("root1/tests", relative_file_path="f.proto"),),
+        }
+    )

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

@@ -7,7 +7,11 @@
 """
 
 from pants.backend.codegen import export_codegen_goal
-from pants.backend.codegen.protobuf.python import additional_fields, python_protobuf_subsystem
+from pants.backend.codegen.protobuf.python import (
+    additional_fields,
+    python_protobuf_module_mapper,
+    python_protobuf_subsystem,
+)
 from pants.backend.codegen.protobuf.python.rules import rules as python_rules
 from pants.backend.codegen.protobuf.target_types import ProtobufLibrary
 
@@ -17,6 +21,7 @@ def rules():
         *additional_fields.rules(),
         *python_protobuf_subsystem.rules(),
         *python_rules(),
+        *python_protobuf_module_mapper.rules(),
         *export_codegen_goal.rules(),
     ]
 

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

@@ -7,7 +7,7 @@
 from pants.backend.codegen.protobuf.python.additional_fields import PythonSourceRootField
 from pants.backend.codegen.protobuf.python.grpc_python_plugin import GrpcPythonPlugin
 from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import PythonProtobufSubsystem
-from pants.backend.codegen.protobuf.target_types import ProtobufGrcpToggle, ProtobufSources
+from pants.backend.codegen.protobuf.target_types import ProtobufGrpcToggle, ProtobufSources
 from pants.backend.python.target_types import PythonSources
 from pants.backend.python.util_rules import extract_pex, pex
 from pants.backend.python.util_rules.extract_pex import ExtractedPexDistributions
@@ -119,7 +119,7 @@ async def generate_python_from_protobuf(
             ExternalToolRequest,
             grpc_python_plugin.get_request(Platform.current),
         )
-        if request.protocol_target.get(ProtobufGrcpToggle).value
+        if request.protocol_target.get(ProtobufGrpcToggle).value
         else None
     )
 

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

@@ -15,7 +15,7 @@ class ProtobufSources(Sources):
     expected_file_extensions = (".proto",)
 
 
-class ProtobufGrcpToggle(BoolField):
+class ProtobufGrpcToggle(BoolField):
     """Whether to generate gRPC code or not."""
 
     alias = "grpc"
@@ -29,4 +29,4 @@ class ProtobufLibrary(Target):
     """
 
     alias = "protobuf_library"
-    core_fields = (*COMMON_TARGET_FIELDS, ProtobufDependencies, ProtobufSources, ProtobufGrcpToggle)
+    core_fields = (*COMMON_TARGET_FIELDS, ProtobufDependencies, ProtobufSources, ProtobufGrpcToggle)