[internal] Simplify Python dependency inference handling of type stubs (

pantsbuild#14017)

Our Python dependency inference code is tricky because we allow you to have one type stub and implementation, whereas that would normally result in ambiguity.

The original implementation was much more complex than necessary because it had each specific module mapping handle ambiguity, e.g. the first-party vs third-party module mappings. Now, we simply have each mapping record _all_ owners and whether each owner is a type stub vs. implementation. Then, our final `ModuleOwners` rule can easily determine ambiguity.

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

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

@@ -12,13 +12,12 @@
 from pants.backend.python.dependency_inference.module_mapper import (
     FirstPartyPythonMappingImpl,
     FirstPartyPythonMappingImplMarker,
+    ModuleProvider,
+    ModuleProviderType,
 )
 from pants.core.util_rules.stripped_source_files import StrippedFileName, StrippedFileNameRequest
-from pants.engine.addresses import Address
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import Target
 from pants.engine.unions import UnionRule
-from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
 
 
@@ -41,35 +40,21 @@ async def map_protobuf_to_python_modules(
         for tgt in protobuf_targets
     )
 
-    # NB: There should be only one address per module, else it's ambiguous.
-    modules_to_addresses: dict[str, tuple[Address]] = {}
-    modules_with_multiple_owners: DefaultDict[str, set[Address]] = defaultdict(set)
-
-    def add_module(module: str, tgt: Target) -> None:
-        if module in modules_to_addresses:
-            modules_with_multiple_owners[module].update(
-                {*modules_to_addresses[module], tgt.address}
-            )
-        else:
-            modules_to_addresses[module] = (tgt.address,)
-
+    modules_to_providers: DefaultDict[str, list[ModuleProvider]] = defaultdict(list)
     for tgt, stripped_file in zip(protobuf_targets, stripped_file_per_target):
         # 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_file.value, suffix="_pb2"), tgt)
+        modules_to_providers[proto_path_to_py_module(stripped_file.value, suffix="_pb2")].append(
+            ModuleProvider(tgt.address, ModuleProviderType.IMPL)
+        )
         if tgt.get(ProtobufGrpcToggleField).value:
-            add_module(proto_path_to_py_module(stripped_file.value, suffix="_pb2_grpc"), tgt)
-
-    # Remove modules with ambiguous owners.
-    for ambiguous_module in modules_with_multiple_owners:
-        modules_to_addresses.pop(ambiguous_module)
+            modules_to_providers[
+                proto_path_to_py_module(stripped_file.value, suffix="_pb2_grpc")
+            ].append(ModuleProvider(tgt.address, ModuleProviderType.IMPL))
 
     return FirstPartyPythonMappingImpl(
-        mapping=FrozenDict(sorted(modules_to_addresses.items())),
-        ambiguous_modules=FrozenDict(
-            (k, tuple(sorted(v))) for k, v in sorted(modules_with_multiple_owners.items())
-        ),
+        (k, tuple(sorted(v))) for k, v in sorted(modules_to_providers.items())
     )
 
 

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

@@ -1,6 +1,8 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from __future__ import annotations
+
 import pytest
 
 from pants.backend.codegen.protobuf.python import additional_fields, python_protobuf_module_mapper
@@ -9,11 +11,14 @@
 )
 from pants.backend.codegen.protobuf.target_types import ProtobufSourcesGeneratorTarget
 from pants.backend.codegen.protobuf.target_types import rules as python_protobuf_target_types_rules
-from pants.backend.python.dependency_inference.module_mapper import FirstPartyPythonMappingImpl
+from pants.backend.python.dependency_inference.module_mapper import (
+    FirstPartyPythonMappingImpl,
+    ModuleProvider,
+    ModuleProviderType,
+)
 from pants.core.util_rules import stripped_source_files
 from pants.engine.addresses import Address
 from pants.testutil.rule_runner import QueryRule, RuleRunner
-from pants.util.frozendict import FrozenDict
 
 
 @pytest.fixture
@@ -37,7 +42,7 @@ def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
             "root1/protos/f1.proto": "",
             "root1/protos/f2.proto": "",
             "root1/protos/BUILD": "protobuf_sources()",
-            # These protos would result in the same module name, so neither should be used.
+            # These protos will result in the same module name.
             "root1/two_owners/f.proto": "",
             "root1/two_owners/BUILD": "protobuf_sources()",
             "root2/two_owners/f.proto": "",
@@ -49,21 +54,21 @@ def test_map_first_party_modules_to_addresses(rule_runner: RuleRunner) -> None:
         }
     )
     result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()])
+
+    def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]:
+        return tuple(ModuleProvider(addr, ModuleProviderType.IMPL) for addr in addresses)
+
     assert result == FirstPartyPythonMappingImpl(
-        mapping=FrozenDict(
-            {
-                "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"),),
-            }
-        ),
-        ambiguous_modules=FrozenDict(
-            {
-                "two_owners.f_pb2": (
+        {
+            "protos.f1_pb2": providers([Address("root1/protos", relative_file_path="f1.proto")]),
+            "protos.f2_pb2": providers([Address("root1/protos", relative_file_path="f2.proto")]),
+            "tests.f_pb2": providers([Address("root1/tests", relative_file_path="f.proto")]),
+            "tests.f_pb2_grpc": providers([Address("root1/tests", relative_file_path="f.proto")]),
+            "two_owners.f_pb2": providers(
+                [
                     Address("root1/two_owners", relative_file_path="f.proto"),
                     Address("root2/two_owners", relative_file_path="f.proto"),
-                )
-            }
-        ),
+                ]
+            ),
+        }
     )