diff --git a/src/python/pants/backend/python/util_rules/python_sources.py b/src/python/pants/backend/python/util_rules/python_sources.py index 1d668d7d66f..2b4f7668cd2 100644 --- a/src/python/pants/backend/python/util_rules/python_sources.py +++ b/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( diff --git a/src/python/pants/backend/python/util_rules/python_sources_test.py b/src/python/pants/backend/python/util_rules/python_sources_test.py index 734044361d8..bcc153856df 100644 --- a/src/python/pants/backend/python/util_rules/python_sources_test.py +++ b/src/python/pants/backend/python/util_rules/python_sources_test.py @@ -8,6 +8,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 @@ -38,6 +39,7 @@ def rule_runner() -> RuleRunner: return RuleRunner( rules=[ *python_sources_rules(), + *additional_fields.rules(), *protobuf_rules(), QueryRule(PythonSourceFiles, [PythonSourceFilesRequest]), QueryRule(StrippedPythonSourceFiles, [PythonSourceFilesRequest]), @@ -215,17 +217,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")