From 2c852401f793ebe65629dae0c37ae69aea0c2dd9 Mon Sep 17 00:00:00 2001
From: Eric Arellano <14852634+Eric-Arellano@users.noreply.github.com>
Date: Tue, 12 Apr 2022 17:01:44 -0700
Subject: [PATCH] Fix when the same `python_requirement` defines both type stub
 and implementation (#15121)

Closes https://github.com/pantsbuild/pants/issues/15111. This is a particular edge case when the same python module has multiple entries, and those entries come from the same target.

The better fix is https://github.com/pantsbuild/pants/issues/14719.

[ci skip-rust]
[ci skip-build-wheels]
---
 .../dependency_inference/module_mapper.py     |  7 +++++
 .../module_mapper_test.py                     | 26 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/python/pants/backend/python/dependency_inference/module_mapper.py b/src/python/pants/backend/python/dependency_inference/module_mapper.py
index 18d5e171df5..b5d50cc6cfc 100644
--- a/src/python/pants/backend/python/dependency_inference/module_mapper.py
+++ b/src/python/pants/backend/python/dependency_inference/module_mapper.py
@@ -8,6 +8,7 @@
 import logging
 from collections import defaultdict
 from dataclasses import dataclass
+from functools import total_ordering
 from pathlib import PurePath
 from typing import DefaultDict, Iterable, Mapping, Tuple
 
@@ -39,10 +40,16 @@
 ResolveName = str
 
 
+@total_ordering
 class ModuleProviderType(enum.Enum):
     TYPE_STUB = enum.auto()
     IMPL = enum.auto()
 
+    def __lt__(self, other) -> bool:
+        if not isinstance(other, ModuleProviderType):
+            return NotImplemented
+        return self.name < other.name
+
 
 @dataclass(frozen=True, order=True)
 class ModuleProvider:
diff --git a/src/python/pants/backend/python/dependency_inference/module_mapper_test.py b/src/python/pants/backend/python/dependency_inference/module_mapper_test.py
index 8e59c7ed4af..6dcf8c39ff8 100644
--- a/src/python/pants/backend/python/dependency_inference/module_mapper_test.py
+++ b/src/python/pants/backend/python/dependency_inference/module_mapper_test.py
@@ -650,3 +650,29 @@ def get_owners(resolve: str | None) -> PythonModuleOwners:
         Address("", target_name="dep1"),
         Address("", target_name="dep2"),
     )
+
+
+def test_issue_15111(rule_runner: RuleRunner) -> None:
+    """Ensure we can handle when a single address implement multiple modules.
+
+    This is currently only possible with third-party targets.
+    """
+    rule_runner.write_files(
+        {"BUILD": "python_requirement(name='req', requirements=['docopt', 'types-docopt'])"}
+    )
+    rule_runner.set_options(["--python-enable-resolves"])
+    result = rule_runner.request(ThirdPartyPythonModuleMapping, [])
+    assert result == ThirdPartyPythonModuleMapping(
+        {
+            "python-default": FrozenDict(
+                {
+                    "docopt": (
+                        ModuleProvider(Address("", target_name="req"), ModuleProviderType.IMPL),
+                        ModuleProvider(
+                            Address("", target_name="req"), ModuleProviderType.TYPE_STUB
+                        ),
+                    ),
+                }
+            )
+        }
+    )