Compute source roots in a rule. (#9888)

Previously we computed the source root for a path
outside the engine, via a utillity method. However we
want to support marking source roots using a sentinel
file, and this requires filesystem access in the computation,
meaning that has to be done in a rule.

This change introduces a set of rules for finding the source root
for a path, and converts all v2 code to use those rules.

[ci skip-rust-tests]
[ci skip-jvm-tests]

src/python/pants/backend/pants_info/list_backends.py

@@ -6,13 +6,14 @@
 from dataclasses import dataclass
 from typing import Optional, Sequence
 
+from pants.core.util_rules.strip_source_roots import SourceRootStrippedSources, StripSnapshotRequest
 from pants.engine.console import Console
 from pants.engine.fs import Digest, FileContent, FilesContent, PathGlobs, Snapshot
 from pants.engine.goal import Goal, GoalSubsystem, LineOriented
 from pants.engine.rules import goal_rule
 from pants.engine.selectors import Get
 from pants.option.global_options import GlobalOptions
-from pants.source.source_root import SourceRootConfig, SourceRoots
+from pants.source.source_root import NoSourceRootError
 
 
 class BackendsOptions(LineOriented, GoalSubsystem):
@@ -73,10 +74,8 @@ class BackendInfo:
 
     @classmethod
     def create(
-        cls, file_content: FileContent, source_roots: SourceRoots, global_options: GlobalOptions
+        cls, file_content: FileContent, stripped_path: str, global_options: GlobalOptions
     ) -> "BackendInfo":
-        source_root = source_roots.strict_find_by_path(file_content.path)
-        stripped_path = file_content.path[len(source_root.path) + 1 :]
         module_name = os.path.dirname(stripped_path).replace(os.sep, ".")
 
         # NB: Both `build_file_aliases` and `target_types` are used by both V1 and V2.
@@ -162,17 +161,26 @@ def format_section(
 
 @goal_rule
 async def list_backends(
-    backend_options: BackendsOptions,
-    source_roots_config: SourceRootConfig,
-    global_options: GlobalOptions,
-    console: Console,
+    backend_options: BackendsOptions, global_options: GlobalOptions, console: Console,
 ) -> Backends:
-    source_roots = source_roots_config.get_source_roots()
     discovered_register_pys = await Get[Snapshot](PathGlobs(["**/*/register.py"]))
     register_pys_content = await Get[FilesContent](Digest, discovered_register_pys.digest)
 
+    source_root_stripped_sources = await Get[SourceRootStrippedSources](
+        StripSnapshotRequest(snapshot=discovered_register_pys)
+    )
+    file_to_stripped_file_mapping = source_root_stripped_sources.get_file_to_stripped_file_mapping()
+    stripped_paths = []
+    for fc in register_pys_content:
+        stripped_path = file_to_stripped_file_mapping.get(fc.path)
+        if stripped_path is None:
+            # This should never happen, but it's worth checking proactively.
+            raise NoSourceRootError(fc.path)
+        stripped_paths.append(stripped_path)
+
     backend_infos = tuple(
-        BackendInfo.create(fc, source_roots, global_options) for fc in register_pys_content
+        BackendInfo.create(fc, stripped_path, global_options)
+        for fc, stripped_path in zip(register_pys_content, stripped_paths)
     )
     v1_backends = []
     v2_backends = []

src/python/pants/backend/pants_info/list_backends_test.py

@@ -8,11 +8,12 @@
     hackily_get_module_docstring,
     list_backends,
 )
+from pants.core.util_rules.strip_source_roots import SourceRootStrippedSources, StripSnapshotRequest
 from pants.engine.fs import EMPTY_SNAPSHOT, Digest, FileContent, FilesContent, PathGlobs, Snapshot
 from pants.option.global_options import GlobalOptions
-from pants.source.source_root import SourceRootConfig
 from pants.testutil.engine.util import MockConsole, MockGet, create_goal_subsystem, run_rule
 from pants.testutil.subsystem.util import global_subsystem_instance
+from pants.util.frozendict import FrozenDict
 
 
 def test_hackily_get_module_docstring() -> None:
@@ -149,11 +150,22 @@ def rules():
         ]
     )
     console = MockConsole(use_colors=False)
+
+    def mock_strip_source_root(_) -> SourceRootStrippedSources:
+        return SourceRootStrippedSources(
+            EMPTY_SNAPSHOT,
+            FrozenDict(
+                {
+                    "src/python": ("pants/backend/fortran/register.py", "pants/core/register.py"),
+                    "contrib/elixir/src/python": ("pants/contrib/elixir/register.py",),
+                }
+            ),
+        )
+
     run_rule(
         list_backends,
         rule_args=[
             create_goal_subsystem(BackendsOptions, sep="\\n", output_file=None),
-            global_subsystem_instance(SourceRootConfig),
             global_subsystem_instance(GlobalOptions),
             console,
         ],
@@ -162,6 +174,11 @@ def rules():
             MockGet(
                 product_type=FilesContent, subject_type=Digest, mock=lambda _: all_register_pys
             ),
+            MockGet(
+                product_type=SourceRootStrippedSources,
+                subject_type=StripSnapshotRequest,
+                mock=mock_strip_source_root,
+            ),
         ],
     )
     assert console.stdout.getvalue() == dedent(

src/python/pants/backend/project_info/list_roots.py

@@ -7,8 +7,13 @@
 from pants.engine.fs import PathGlobs, Snapshot
 from pants.engine.goal import Goal, GoalSubsystem, LineOriented
 from pants.engine.rules import SubsystemRule, goal_rule, rule
-from pants.engine.selectors import Get
-from pants.source.source_root import AllSourceRoots, SourceRoot, SourceRootConfig
+from pants.engine.selectors import Get, MultiGet
+from pants.source.source_root import (
+    AllSourceRoots,
+    OptionalSourceRoot,
+    SourceRootConfig,
+    SourceRootRequest,
+)
 
 
 class RootsOptions(LineOriented, GoalSubsystem):
@@ -23,7 +28,6 @@ class Roots(Goal):
 
 @rule
 async def all_roots(source_root_config: SourceRootConfig) -> AllSourceRoots:
-
     source_roots = source_root_config.get_source_roots()
 
     all_paths: Set[str] = set()
@@ -43,17 +47,14 @@ async def all_roots(source_root_config: SourceRootConfig) -> AllSourceRoots:
 
     snapshot = await Get[Snapshot](PathGlobs(globs=sorted(all_paths)))
 
-    all_source_roots: Set[SourceRoot] = set()
-
     # The globs above can match on subdirectories of the source roots.
     # For instance, `src/*/` might match 'src/rust/' as well as
     # 'src/rust/engine/process_execution/bazel_protos/src/gen'.
     # So we use find_by_path to verify every candidate source root.
-    for directory in snapshot.dirs:
-        match: SourceRoot = source_roots.find_by_path(directory)
-        if match:
-            all_source_roots.add(match)
-
+    responses = await MultiGet(Get[OptionalSourceRoot](SourceRootRequest(d)) for d in snapshot.dirs)
+    all_source_roots = {
+        response.source_root for response in responses if response.source_root is not None
+    }
     return AllSourceRoots(all_source_roots)
 
 

src/python/pants/backend/project_info/list_roots_test.py

@@ -5,7 +5,12 @@
 
 from pants.backend.project_info import list_roots
 from pants.engine.fs import Digest, PathGlobs, Snapshot
-from pants.source.source_root import SourceRoot, SourceRootConfig
+from pants.source.source_root import (
+    OptionalSourceRoot,
+    SourceRoot,
+    SourceRootConfig,
+    SourceRootRequest,
+)
 from pants.testutil.engine.util import MockGet, run_rule
 from pants.testutil.goal_rule_test_base import GoalRuleTestBase
 from pants.testutil.test_base import TestBase
@@ -53,10 +58,20 @@ def provider_rule(path_globs: PathGlobs) -> Snapshot:
             )
             return Snapshot(Digest("abcdef", 10), (), dirs)
 
+        def source_root(req: SourceRootRequest) -> OptionalSourceRoot:
+            return OptionalSourceRoot(source_roots.find_by_path(req.path))
+
         output = run_rule(
             list_roots.all_roots,
             rule_args=[source_root_config],
-            mock_gets=[MockGet(product_type=Snapshot, subject_type=PathGlobs, mock=provider_rule)],
+            mock_gets=[
+                MockGet(product_type=Snapshot, subject_type=PathGlobs, mock=provider_rule),
+                MockGet(
+                    product_type=OptionalSourceRoot,
+                    subject_type=SourceRootRequest,
+                    mock=source_root,
+                ),
+            ],
         )
 
         self.assertEqual(
@@ -87,16 +102,27 @@ def test_all_roots_with_root_at_buildroot(self):
         )
 
         source_root_config = SourceRootConfig.global_instance()
+        source_roots = source_root_config.get_source_roots()
 
         # This function mocks out reading real directories off the file system
         def provider_rule(path_globs: PathGlobs) -> Snapshot:
             dirs = ("foo",)  # A python package at the buildroot.
             return Snapshot(Digest("abcdef", 10), (), dirs)
 
+        def source_root(req: SourceRootRequest) -> OptionalSourceRoot:
+            return OptionalSourceRoot(source_roots.find_by_path(req.path))
+
         output = run_rule(
             list_roots.all_roots,
             rule_args=[source_root_config],
-            mock_gets=[MockGet(product_type=Snapshot, subject_type=PathGlobs, mock=provider_rule)],
+            mock_gets=[
+                MockGet(product_type=Snapshot, subject_type=PathGlobs, mock=provider_rule),
+                MockGet(
+                    product_type=OptionalSourceRoot,
+                    subject_type=SourceRootRequest,
+                    mock=source_root,
+                ),
+            ],
         )
 
         self.assertEqual({SourceRoot("")}, set(output))
@@ -107,7 +133,7 @@ class RootsTest(GoalRuleTestBase):
 
     @classmethod
     def rules(cls):
-        return super().rules() + list_roots.rules()
+        return [*super().rules(), *list_roots.rules()]
 
     # Delete these *_deprecated tests in 1.30.0.dev0.
     def test_no_langs_deprecated(self):

src/python/pants/backend/python/rules/pytest_coverage.py

@@ -8,7 +8,7 @@
 from io import StringIO
 from pathlib import PurePath
 from textwrap import dedent
-from typing import Optional, Tuple
+from typing import Optional
 
 import pkg_resources
 
@@ -30,6 +30,7 @@
     FilesystemCoverageReport,
 )
 from pants.core.util_rules.determine_source_files import AllSourceFilesRequest, SourceFiles
+from pants.core.util_rules.strip_source_roots import SourceRootStrippedSources, StripSnapshotRequest
 from pants.engine.addresses import Address
 from pants.engine.fs import (
     AddPrefix,
@@ -45,7 +46,6 @@
 from pants.engine.target import Sources, Targets, TransitiveTargets
 from pants.engine.unions import UnionRule
 from pants.python.python_setup import PythonSetup
-from pants.source.source_root import SourceRootConfig
 
 # There are many moving parts in coverage, so here is a high level view of what's going on.
 
@@ -139,29 +139,24 @@ class CoverageConfig:
 
 
 @rule
-async def create_coverage_config(
-    coverage_config_request: CoverageConfigRequest, source_root_config: SourceRootConfig
-) -> CoverageConfig:
+async def create_coverage_config(coverage_config_request: CoverageConfigRequest) -> CoverageConfig:
     sources = await Get[SourceFiles](
         AllSourceFilesRequest((tgt.get(Sources) for tgt in coverage_config_request.targets))
     )
     init_injected = await Get[InitInjectedSnapshot](InjectInitRequest(sources.snapshot))
-    source_roots = source_root_config.get_source_roots()
+
+    source_root_stripped_sources = await Get[SourceRootStrippedSources](
+        StripSnapshotRequest(snapshot=init_injected.snapshot)
+    )
 
     # Generate a map from source root stripped source to its source root. eg:
     #  {'pants/testutil/subsystem/util.py': 'src/python'}. This is so that coverage reports
     #  referencing /chroot/path/pants/testutil/subsystem/util.py can be mapped back to the actual
     #  sources they reference when generating coverage reports.
-    def stripped_file_with_source_root(file_name: str) -> Tuple[str, str]:
-        source_root_object = source_roots.find_by_path(file_name)
-        source_root = source_root_object.path if source_root_object is not None else ""
-        stripped_path = file_name[len(source_root) + 1 :]
-        return stripped_path, source_root
-
-    stripped_files_to_source_roots = dict(
-        stripped_file_with_source_root(filename)
-        for filename in sorted(init_injected.snapshot.files)
-    )
+    stripped_files_to_source_roots = {}
+    for source_root, files in source_root_stripped_sources.root_to_relfiles.items():
+        for file in files:
+            stripped_files_to_source_roots[file] = source_root
 
     default_config = dedent(
         """

src/python/pants/backend/python/rules/run_setup_py.py

@@ -66,7 +66,7 @@
 from pants.engine.unions import UnionMembership
 from pants.option.custom_types import shell_str
 from pants.python.python_setup import PythonSetup
-from pants.source.source_root import SourceRootConfig
+from pants.source.source_root import SourceRoot, SourceRootRequest
 from pants.util.ordered_set import FrozenOrderedSet
 
 logger = logging.getLogger(__name__)
@@ -460,9 +460,7 @@ async def generate_chroot(request: SetupPyChrootRequest) -> SetupPyChroot:
 
 
 @rule
-async def get_sources(
-    request: SetupPySourcesRequest, source_root_config: SourceRootConfig
-) -> SetupPySources:
+async def get_sources(request: SetupPySourcesRequest) -> SetupPySources:
     targets = request.targets
     stripped_srcs_list = await MultiGet(
         Get[SourceRootStrippedSources](
@@ -494,7 +492,6 @@ async def get_sources(
     init_py_contents = await Get[FilesContent](Digest, init_pys_snapshot.digest)
 
     packages, namespace_packages, package_data = find_packages(
-        source_roots=source_root_config.get_source_roots(),
         tgts_and_stripped_srcs=list(zip(targets, stripped_srcs_list)),
         init_py_contents=init_py_contents,
         py2=request.py2,
@@ -508,14 +505,11 @@ async def get_sources(
 
 
 @rule
-async def get_ancestor_init_py(
-    targets: Targets, source_root_config: SourceRootConfig
-) -> AncestorInitPyFiles:
+async def get_ancestor_init_py(targets: Targets) -> AncestorInitPyFiles:
     """Find any ancestor __init__.py files for the given targets.
 
     Includes sibling __init__.py files. Returns the files stripped of their source roots.
     """
-    source_roots = source_root_config.get_source_roots()
     sources = await Get[SourceFiles](
         AllSourceFilesRequest(
             (tgt.get(Sources) for tgt in targets),
@@ -525,12 +519,15 @@ async def get_ancestor_init_py(
     )
     # Find the ancestors of all dirs containing .py files, including those dirs themselves.
     source_dir_ancestors: Set[Tuple[str, str]] = set()  # Items are (src_root, path incl. src_root).
-    for fp in sources.snapshot.files:
-        source_dir_ancestor = os.path.dirname(fp)
-        source_root = source_roots.strict_find_by_path(fp).path
+    source_roots = await MultiGet(
+        Get[SourceRoot](SourceRootRequest, SourceRootRequest.for_file(path))
+        for path in sources.files
+    )
+    for path, source_root in zip(sources.files, source_roots):
+        source_dir_ancestor = os.path.dirname(path)
         # Do not allow the repository root to leak (i.e., '.' should not be a package in setup.py).
-        while source_dir_ancestor != source_root:
-            source_dir_ancestors.add((source_root, source_dir_ancestor))
+        while source_dir_ancestor != source_root.path:
+            source_dir_ancestors.add((source_root.path, source_dir_ancestor))
             source_dir_ancestor = os.path.dirname(source_dir_ancestor)
 
     source_dir_ancestors_list = list(source_dir_ancestors)  # To force a consistent order.

src/python/pants/backend/python/rules/util.py

@@ -14,7 +14,7 @@
 from pants.engine.fs import FilesContent
 from pants.engine.target import Target
 from pants.python.python_setup import PythonSetup
-from pants.source.source_root import SourceRoots
+from pants.source.source_root import SourceRootError
 from pants.util.strutil import ensure_text
 
 # Convenient type alias for the pair (package name, data files in the package).
@@ -95,7 +95,6 @@ def _write_repr(o, indent=False, level=0):
 
 
 def find_packages(
-    source_roots: SourceRoots,
     tgts_and_stripped_srcs: Iterable[Tuple[Target, SourceRootStrippedSources]],
     init_py_contents: FilesContent,
     py2: bool,
@@ -125,7 +124,15 @@ def find_packages(
     # Now find all package_data.
     for tgt, stripped_srcs in tgts_and_stripped_srcs:
         if tgt.has_field(ResourcesSources):
-            source_root = source_roots.strict_find_by_path(tgt.address.spec_path).path
+            source_roots = list(stripped_srcs.root_to_relfiles.keys())
+            if len(source_roots) != 1:
+                # A single target is expected to be under a single source root, so this
+                # should never happen, but we check to be sure.
+                raise SourceRootError(
+                    f"Expected a single source root for target "
+                    f"{tgt.address.spec} but found: {', '.join(source_roots)}. "
+                )
+            source_root = source_roots[0]
             resource_dir_relpath = os.path.relpath(tgt.address.spec_path, source_root)
             # Find the closest enclosing package, if any.  Resources will be loaded relative to that.
             package: str = resource_dir_relpath.replace(os.path.sep, ".")