Fix PYTHONPATH not including protobuf targets using python_source_root (Cherry-pick of #11673)

src/python/pants/backend/python/util_rules/python_sources.py

@@ -13,7 +13,7 @@
 from pants.core.util_rules.stripped_source_files import StrippedSourceFiles
 from pants.engine.fs import MergeDigests, Snapshot
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import Sources, Target
+from pants.engine.target import HydratedSources, HydrateSourcesRequest, Sources, Target
 from pants.engine.unions import UnionMembership
 from pants.source.source_root import SourceRoot, SourceRootRequest
 from pants.util.logging import LogLevel
@@ -98,15 +98,39 @@ async def prepare_python_sources(
         MergeDigests((sources.snapshot.digest, missing_init_files.snapshot.digest)),
     )
 
-    source_root_objs = await MultiGet(
-        Get(SourceRoot, SourceRootRequest, SourceRootRequest.for_target(tgt))
-        for tgt in request.targets
-        if (
-            tgt.has_field(PythonSources)
-            or tgt.has_field(ResourcesSources)
-            or tgt.get(Sources).can_generate(PythonSources, union_membership)
-            or tgt.get(Sources).can_generate(ResourcesSources, union_membership)
+    # Codegen is able to generate code in any arbitrary location, unlike sources normally being
+    # rooted under the target definition. To determine source roots for these generated files, we
+    # cannot use the normal `SourceRootRequest.for_target()` and we instead must determine
+    # a source root for every individual generated file. So, we re-resolve the codegen sources here.
+    python_and_resources_targets = []
+    codegen_targets = []
+    for tgt in request.targets:
+        if tgt.has_field(PythonSources) or tgt.has_field(ResourcesSources):
+            python_and_resources_targets.append(tgt)
+        elif tgt.get(Sources).can_generate(PythonSources, union_membership) or tgt.get(
+            Sources
+        ).can_generate(ResourcesSources, union_membership):
+            codegen_targets.append(tgt)
+    codegen_sources = await MultiGet(
+        Get(
+            HydratedSources,
+            HydrateSourcesRequest(
+                tgt.get(Sources), for_sources_types=request.valid_sources_types, enable_codegen=True
+            ),
         )
+        for tgt in codegen_targets
+    )
+    source_root_requests = [
+        *(SourceRootRequest.for_target(tgt) for tgt in python_and_resources_targets),
+        *(
+            SourceRootRequest.for_file(f)
+            for sources in codegen_sources
+            for f in sources.snapshot.files
+        ),
+    ]
+
+    source_root_objs = await MultiGet(
+        Get(SourceRoot, SourceRootRequest, req) for req in source_root_requests
     )
     source_root_paths = {source_root_obj.path for source_root_obj in source_root_objs}
     return PythonSourceFiles(

src/python/pants/backend/python/util_rules/python_sources_test.py

@@ -6,6 +6,7 @@
 
 import pytest
 
+from pants.backend.codegen.protobuf.python import additional_fields
 from pants.backend.codegen.protobuf.python.rules import rules as protobuf_rules
 from pants.backend.codegen.protobuf.target_types import ProtobufLibrary
 from pants.backend.python.target_types import PythonSources
@@ -36,6 +37,7 @@ def rule_runner() -> RuleRunner:
     return RuleRunner(
         rules=[
             *python_sources_rules(),
+            *additional_fields.rules(),
             *protobuf_rules(),
             QueryRule(PythonSourceFiles, [PythonSourceFilesRequest]),
             QueryRule(StrippedPythonSourceFiles, [PythonSourceFilesRequest]),
@@ -213,17 +215,39 @@ def test_python_protobuf(rule_runner: RuleRunner) -> None:
             """
         ),
     )
+    rule_runner.create_file(
+        "src/protobuf/other_dir/f.proto",
+        dedent(
+            """\
+            syntax = "proto2";
+
+            package other_dir;
+            """
+        ),
+    )
     rule_runner.add_to_build_file("src/protobuf/dir", "protobuf_library()")
-    targets = [ProtobufLibrary({}, address=Address("src/protobuf/dir"))]
+    rule_runner.add_to_build_file(
+        "src/protobuf/other_dir", "protobuf_library(python_source_root='src/python')"
+    )
+    targets = [
+        ProtobufLibrary({}, address=Address("src/protobuf/dir")),
+        ProtobufLibrary({}, address=Address("src/protobuf/other_dir")),
+    ]
     backend_args = ["--backend-packages=pants.backend.codegen.protobuf.python"]
 
     stripped_result = get_stripped_sources(
-        rule_runner, targets, source_roots=["src/protobuf"], extra_args=backend_args
+        rule_runner, targets, source_roots=["src/protobuf", "src/python"], extra_args=backend_args
+    )
+    assert stripped_result.stripped_source_files.snapshot.files == (
+        "dir/f_pb2.py",
+        "other_dir/f_pb2.py",
     )
-    assert stripped_result.stripped_source_files.snapshot.files == ("dir/f_pb2.py",)
 
     unstripped_result = get_unstripped_sources(
-        rule_runner, targets, source_roots=["src/protobuf"], extra_args=backend_args
+        rule_runner, targets, source_roots=["src/protobuf", "src/python"], extra_args=backend_args
+    )
+    assert unstripped_result.source_files.snapshot.files == (
+        "src/protobuf/dir/f_pb2.py",
+        "src/python/other_dir/f_pb2.py",
     )
-    assert unstripped_result.source_files.snapshot.files == ("src/protobuf/dir/f_pb2.py",)
-    assert unstripped_result.source_roots == ("src/protobuf",)
+    assert unstripped_result.source_roots == ("src/protobuf", "src/python")