A helper method for indexing paths by source root (#23085)

Author: benjyw

src/python/pants/source/source_root.py

@@ -6,6 +6,7 @@
 import itertools
 import logging
 import os
+from collections import defaultdict
 from collections.abc import Iterable
 from dataclasses import dataclass
 from pathlib import PurePath
@@ -173,7 +174,7 @@ def __post_init__(self) -> None:
 
     @classmethod
     def for_files(cls, file_paths: Iterable[str]) -> SourceRootsRequest:
-        """Create a request for the source root for the given file."""
+        """Create a request for the source root for the given files."""
         return cls({PurePath(file_path) for file_path in file_paths}, ())
 
 
@@ -217,6 +218,12 @@ def debug_hint(self) -> str:
 class SourceRootsResult:
     path_to_root: FrozenDict[PurePath, SourceRoot]
 
+    def root_to_paths(self) -> FrozenDict[SourceRoot, tuple[PurePath, ...]]:
+        res = defaultdict(list)
+        for path, root in self.path_to_root.items():
+            res[root].append(path)
+        return FrozenDict((k, tuple(sorted(v))) for k, v in res.items())
+
 
 @dataclass(frozen=True)
 class OptionalSourceRootsResult:

src/python/pants/source/source_root_test.py

@@ -22,6 +22,7 @@
 from pants.source.source_root import rules as source_root_rules
 from pants.testutil.option_util import create_subsystem
 from pants.testutil.rule_runner import RuleRunner, run_rule_with_mocks
+from pants.util.frozendict import FrozenDict
 
 
 def _find_root(
@@ -287,3 +288,26 @@ def test_source_roots_request() -> None:
         PurePath("src/python/foo"): SourceRoot("src/python"),
         PurePath("src/python/baz/qux"): SourceRoot("src/python"),
     } == dict(res.path_to_root)
+
+
+def test_root_to_paths() -> None:
+    res = SourceRootsResult(
+        FrozenDict(
+            {
+                PurePath("src/python/foo/bar.py"): SourceRoot("src/python"),
+                PurePath("tests/python/foo/bar_test.py"): SourceRoot("tests/python"),
+                PurePath("src/python/foo"): SourceRoot("src/python"),
+                PurePath("src/python/baz/qux"): SourceRoot("src/python"),
+            }
+        )
+    )
+    assert res.root_to_paths() == FrozenDict(
+        {
+            SourceRoot("src/python"): (
+                PurePath("src/python/baz/qux"),
+                PurePath("src/python/foo"),
+                PurePath("src/python/foo/bar.py"),
+            ),
+            SourceRoot("tests/python"): (PurePath("tests/python/foo/bar_test.py"),),
+        }
+    )