Simplify source root stripping.

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

@@ -6,7 +6,7 @@
 from pants.backend.python.target_types import PythonSources
 from pants.core.util_rules.determine_source_files import AllSourceFilesRequest, SourceFiles
 from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
-from pants.core.util_rules.strip_source_roots import representative_path_from_address
+from pants.core.util_rules.strip_source_roots import StrippedSourceFiles
 from pants.engine.addresses import Addresses
 from pants.engine.fs import AddPrefix, Digest, MergeDigests, RemovePrefix, Snapshot
 from pants.engine.platform import Platform
@@ -52,33 +52,36 @@ async def generate_python_from_protobuf(
         AllSourceFilesRequest(
             (tgt.get(Sources) for tgt in transitive_targets.closure),
             for_sources_types=(ProtobufSources,),
-            # NB: By stripping the source roots, we avoid having to set the value `--proto_path`
-            # for Protobuf imports to be discoverable.
-            strip_source_roots=True,
         ),
     )
-    stripped_target_sources_request = Get(
-        SourceFiles,
-        AllSourceFilesRequest([request.protocol_target[ProtobufSources]], strip_source_roots=True),
+    unstripped_target_sources_request = Get(
+        SourceFiles, AllSourceFilesRequest([request.protocol_target[ProtobufSources]]),
     )
 
     (
         downloaded_protoc_binary,
         create_output_dir_result,
-        all_sources,
-        stripped_target_sources,
+        all_sources_unstripped,
+        target_sources_unstripped,
     ) = await MultiGet(
         download_protoc_request,
         create_output_dir_request,
         all_sources_request,
-        stripped_target_sources_request,
+        unstripped_target_sources_request,
+    )
+
+    # NB: By stripping the source roots, we avoid having to set the value `--proto_path`
+    # for Protobuf imports to be discoverable.
+    all_sources_stripped, target_sources_stripped = await MultiGet(
+        Get(StrippedSourceFiles, SourceFiles, all_sources_unstripped),
+        Get(StrippedSourceFiles, SourceFiles, target_sources_unstripped),
     )
 
     input_digest = await Get(
         Digest,
         MergeDigests(
             (
-                all_sources.snapshot.digest,
+                all_sources_stripped.snapshot.digest,
                 downloaded_protoc_binary.digest,
                 create_output_dir_result.output_digest,
             )
@@ -92,7 +95,7 @@ async def generate_python_from_protobuf(
                 downloaded_protoc_binary.exe,
                 "--python_out",
                 output_dir,
-                *stripped_target_sources.snapshot.files,
+                *target_sources_stripped.snapshot.files,
             ),
             input_digest=input_digest,
             description=f"Generating Python sources from {request.protocol_target.address}.",
@@ -105,13 +108,7 @@ async def generate_python_from_protobuf(
     # including adding back the original source root.
     normalized_digest, source_root = await MultiGet(
         Get(Digest, RemovePrefix(result.output_digest, output_dir)),
-        Get(
-            SourceRoot,
-            SourceRootRequest,
-            SourceRootRequest.for_file(
-                representative_path_from_address(request.protocol_target.address)
-            ),
-        ),
+        Get(SourceRoot, SourceRootRequest, SourceRootRequest.for_target(request.protocol_target)),
     )
     source_root_restored = (
         await Get(Snapshot, AddPrefix(normalized_digest, source_root.path))

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

@@ -7,7 +7,7 @@
 from pants.backend.codegen.protobuf.python.rules import GeneratePythonFromProtobufRequest
 from pants.backend.codegen.protobuf.python.rules import rules as protobuf_rules
 from pants.backend.codegen.protobuf.target_types import ProtobufLibrary, ProtobufSources
-from pants.core.util_rules import determine_source_files
+from pants.core.util_rules import determine_source_files, strip_source_roots
 from pants.engine.addresses import Address
 from pants.engine.rules import RootRule
 from pants.engine.target import (
@@ -32,6 +32,7 @@ def rules(cls):
             *super().rules(),
             *protobuf_rules(),
             *determine_source_files.rules(),
+            *strip_source_roots.rules(),
             RootRule(GeneratePythonFromProtobufRequest),
         )
 

src/python/pants/backend/python/dependency_inference/module_mapper.py

@@ -7,10 +7,8 @@
 
 from pants.backend.python.target_types import PythonRequirementsField, PythonSources
 from pants.base.specs import AddressSpecs, DescendantAddresses
-from pants.core.util_rules.strip_source_roots import (
-    SourceRootStrippedSources,
-    StripSourcesFieldRequest,
-)
+from pants.core.util_rules.determine_source_files import AllSourceFilesRequest
+from pants.core.util_rules.strip_source_roots import StrippedSourceFiles
 from pants.engine.addresses import Address
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import (
@@ -75,9 +73,10 @@ async def map_first_party_modules_to_addresses() -> FirstPartyModuleToAddressMap
         for tgt in candidate_explicit_targets
     )
     stripped_sources_per_explicit_target = await MultiGet(
-        Get(SourceRootStrippedSources, StripSourcesFieldRequest(tgt[PythonSources]))
+        Get(StrippedSourceFiles, AllSourceFilesRequest([tgt[PythonSources]]))
         for tgt in candidate_explicit_targets
     )
+
     modules_to_addresses: Dict[str, Address] = {}
     modules_with_multiple_owners: Set[str] = set()
     for explicit_tgt, unstripped_sources, stripped_sources in zip(

src/python/pants/backend/python/dependency_inference/module_mapper_test.py

@@ -18,7 +18,7 @@
 )
 from pants.backend.python.target_types import PythonLibrary, PythonRequirementLibrary
 from pants.build_graph.build_file_aliases import BuildFileAliases
-from pants.core.util_rules import strip_source_roots
+from pants.core.util_rules import determine_source_files, strip_source_roots
 from pants.engine.addresses import Address
 from pants.engine.rules import RootRule
 from pants.python.python_requirement import PythonRequirement
@@ -84,6 +84,7 @@ def rules(cls):
         return (
             *super().rules(),
             *strip_source_roots.rules(),
+            *determine_source_files.rules(),
             map_first_party_modules_to_addresses,
             map_module_to_address,
             map_third_party_modules_to_addresses,

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

@@ -12,10 +12,8 @@
 from pants.backend.python.rules import ancestor_files
 from pants.backend.python.rules.ancestor_files import AncestorFiles, AncestorFilesRequest
 from pants.backend.python.target_types import PythonSources, PythonTestsSources
-from pants.core.util_rules.strip_source_roots import (
-    SourceRootStrippedSources,
-    StripSourcesFieldRequest,
-)
+from pants.core.util_rules.determine_source_files import AllSourceFilesRequest
+from pants.core.util_rules.strip_source_roots import StrippedSourceFiles
 from pants.engine.fs import Digest, DigestContents
 from pants.engine.internals.graph import Owners, OwnersRequest
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
@@ -92,7 +90,7 @@ async def infer_python_dependencies(
         return InferredDependencies()
 
     stripped_sources = await Get(
-        SourceRootStrippedSources, StripSourcesFieldRequest(request.sources_field)
+        StrippedSourceFiles, AllSourceFilesRequest([request.sources_field])
     )
     modules = tuple(
         PythonModule.create_from_stripped_path(PurePath(fp))
@@ -133,8 +131,7 @@ async def infer_python_init_dependencies(
     # Locate __init__.py files not already in the Snapshot.
     hydrated_sources = await Get(HydratedSources, HydrateSourcesRequest(request.sources_field))
     extra_init_files = await Get(
-        AncestorFiles,
-        AncestorFilesRequest("__init__.py", hydrated_sources.snapshot, sources_stripped=False),
+        AncestorFiles, AncestorFilesRequest("__init__.py", hydrated_sources.snapshot),
     )
 
     # And add dependencies on their owners.
@@ -161,8 +158,7 @@ async def infer_python_conftest_dependencies(
     # Locate conftest.py files not already in the Snapshot.
     hydrated_sources = await Get(HydratedSources, HydrateSourcesRequest(request.sources_field))
     extra_conftest_files = await Get(
-        AncestorFiles,
-        AncestorFilesRequest("conftest.py", hydrated_sources.snapshot, sources_stripped=False),
+        AncestorFiles, AncestorFilesRequest("conftest.py", hydrated_sources.snapshot),
     )
 
     # And add dependencies on their owners.

src/python/pants/backend/python/dependency_inference/rules_test.py

@@ -16,7 +16,7 @@
     PythonTests,
 )
 from pants.build_graph.build_file_aliases import BuildFileAliases
-from pants.core.util_rules import strip_source_roots
+from pants.core.util_rules import determine_source_files, strip_source_roots
 from pants.engine.addresses import Address
 from pants.engine.rules import RootRule
 from pants.engine.target import InferredDependencies, WrappedTarget
@@ -37,6 +37,7 @@ def rules(cls):
         return (
             *super().rules(),
             *strip_source_roots.rules(),
+            *determine_source_files.rules(),
             *dependency_inference_rules(),
             all_roots,
             RootRule(InferPythonDependencies),

src/python/pants/backend/python/lint/black/rules.py

@@ -111,7 +111,7 @@ async def setup(setup_request: SetupRequest, black: Black) -> Setup:
     all_source_files, requirements_pex, config_digest, specified_source_files = (
         await MultiGet(all_source_files_request, *requests)
         if setup_request.request.prior_formatter_result is None
-        else (SourceFiles(EMPTY_SNAPSHOT), *await MultiGet(*requests))
+        else (SourceFiles(EMPTY_SNAPSHOT, ()), *await MultiGet(*requests))
     )
     all_source_files_snapshot = (
         all_source_files.snapshot

src/python/pants/backend/python/lint/docformatter/rules.py

@@ -91,7 +91,7 @@ async def setup(setup_request: SetupRequest, docformatter: Docformatter) -> Setu
     all_source_files, requirements_pex, specified_source_files = (
         await MultiGet(all_source_files_request, *requests)
         if setup_request.request.prior_formatter_result is None
-        else (SourceFiles(EMPTY_SNAPSHOT), *await MultiGet(*requests))
+        else (SourceFiles(EMPTY_SNAPSHOT, ()), *await MultiGet(*requests))
     )
     all_source_files_snapshot = (
         all_source_files.snapshot

src/python/pants/backend/python/lint/isort/rules.py

@@ -115,7 +115,7 @@ async def setup(setup_request: SetupRequest, isort: Isort) -> Setup:
     all_source_files, requirements_pex, config_digest, specified_source_files = (
         await MultiGet(all_source_files_request, *requests)
         if setup_request.request.prior_formatter_result is None
-        else (SourceFiles(EMPTY_SNAPSHOT), *await MultiGet(*requests))
+        else (SourceFiles(EMPTY_SNAPSHOT, ()), *await MultiGet(*requests))
     )
     all_source_files_snapshot = (
         all_source_files.snapshot

src/python/pants/backend/python/lint/pylint/rules.py

@@ -16,10 +16,9 @@
     PexRequirements,
 )
 from pants.backend.python.rules.python_sources import (
-    StrippedPythonSources,
-    StrippedPythonSourcesRequest,
-    UnstrippedPythonSources,
-    UnstrippedPythonSourcesRequest,
+    PythonSourceFiles,
+    PythonSourceFilesRequest,
+    StrippedPythonSourceFiles,
 )
 from pants.backend.python.target_types import (
     PythonInterpreterCompatibility,
@@ -160,11 +159,10 @@ async def pylint_lint_partition(partition: PylintPartition, pylint: Pylint) -> L
     )
 
     prepare_plugin_sources_request = Get(
-        StrippedPythonSources, StrippedPythonSourcesRequest(partition.plugin_targets),
+        StrippedPythonSourceFiles, PythonSourceFilesRequest(partition.plugin_targets),
     )
     prepare_python_sources_request = Get(
-        UnstrippedPythonSources,
-        UnstrippedPythonSourcesRequest(partition.targets_with_dependencies),
+        PythonSourceFiles, PythonSourceFilesRequest(partition.targets_with_dependencies),
     )
     specified_source_files_request = Get(
         SourceFiles,
@@ -192,7 +190,10 @@ async def pylint_lint_partition(partition: PylintPartition, pylint: Pylint) -> L
     )
 
     prefixed_plugin_sources = (
-        await Get(Digest, AddPrefix(prepared_plugin_sources.snapshot.digest, "__plugins"))
+        await Get(
+            Digest,
+            AddPrefix(prepared_plugin_sources.stripped_source_files.snapshot.digest, "__plugins"),
+        )
         if pylint.source_plugins
         else EMPTY_DIGEST
     )
@@ -215,7 +216,7 @@ async def pylint_lint_partition(partition: PylintPartition, pylint: Pylint) -> L
                 pylint_runner_pex.digest,
                 config_digest,
                 prefixed_plugin_sources,
-                prepared_python_sources.snapshot.digest,
+                prepared_python_sources.source_files.snapshot.digest,
             )
         ),
     )
@@ -287,7 +288,7 @@ async def pylint_lint(
 
     partitions = (
         PylintPartition(
-            tuple(sorted(target_setups, key=lambda target_setup: target_setup.field_set.address)),
+            tuple(sorted(target_setups, key=lambda tgt_setup: tgt_setup.field_set.address)),
             interpreter_constraints,
             Targets(plugin_targets.closure),
         )

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

@@ -1,16 +1,12 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-import itertools
 import os
 from dataclasses import dataclass
-from pathlib import PurePath
 from typing import Sequence, Set
 
-from pants.core.util_rules.strip_source_roots import SourceRootStrippedSources, StripSnapshotRequest
 from pants.engine.fs import EMPTY_SNAPSHOT, PathGlobs, Snapshot
 from pants.engine.rules import Get, collect_rules, rule
-from pants.source.source_root import AllSourceRoots
 from pants.util.ordered_set import FrozenOrderedSet
 
 
@@ -31,25 +27,22 @@ class AncestorFilesRequest:
 
     name: str
     snapshot: Snapshot
-    sources_stripped: bool  # True iff snapshot has already had source roots stripped.
 
 
 @dataclass(frozen=True)
 class AncestorFiles:
     """Any ancestor files found."""
 
-    # Files in the snapshot are stripped iff the request had sources_stripped=True.
     snapshot: Snapshot
 
 
 def identify_missing_ancestor_files(name: str, sources: Sequence[str]) -> FrozenOrderedSet[str]:
     """Return the paths of potentially missing ancestor files.
 
-    NB: If the sources have not had their source roots (e.g., 'src/python') stripped, this
-    function will consider superfluous files at and above the source roots, (e.g.,
-    src/python/<name>, src/<name>). It is the caller's responsibility to filter these
-    out if necessary. If the sources have had their source roots stripped, then this function
-    will only identify consider files in actual packages.
+    NB: The sources are expected to not have had their source roots stripped.
+    Therefore this function will consider superfluous files at and above the source roots,
+    (e.g., src/python/<name>, src/<name>). It is the caller's responsibility to filter these
+    out if necessary.
     """
     packages: Set[str] = set()
     for source in sources:
@@ -69,33 +62,14 @@ def identify_missing_ancestor_files(name: str, sources: Sequence[str]) -> Frozen
 
 
 @rule
-async def find_missing_ancestor_files(
-    request: AncestorFilesRequest, all_source_roots: AllSourceRoots
-) -> AncestorFiles:
+async def find_missing_ancestor_files(request: AncestorFilesRequest) -> AncestorFiles:
     """Find any named ancestor files that exist on the filesystem but are not in the snapshot."""
     missing_ancestor_files = identify_missing_ancestor_files(request.name, request.snapshot.files)
     if not missing_ancestor_files:
         return AncestorFiles(EMPTY_SNAPSHOT)
 
-    if request.sources_stripped:
-        # If files are stripped, we don't know what source root they might live in, so we look
-        # up every source root.
-        roots = tuple(root.path for root in all_source_roots)
-        missing_ancestor_files = FrozenOrderedSet(
-            PurePath(root, f).as_posix()
-            for root, f in itertools.product(roots, missing_ancestor_files)
-        )
-
     # NB: This will intentionally _not_ error on any unmatched globs.
     discovered_ancestors_snapshot = await Get(Snapshot, PathGlobs(missing_ancestor_files))
-
-    if request.sources_stripped:
-        # We must now strip all discovered paths.
-        stripped_snapshot = await Get(
-            SourceRootStrippedSources, StripSnapshotRequest(discovered_ancestors_snapshot)
-        )
-        discovered_ancestors_snapshot = stripped_snapshot.snapshot
-
     return AncestorFiles(discovered_ancestors_snapshot)