Allow protobuf_library targets to specify a python source root.

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

@@ -3,11 +3,24 @@
 
 from pants.backend.codegen.protobuf.target_types import ProtobufLibrary
 from pants.backend.python.target_types import PythonInterpreterCompatibility
+from pants.engine.target import StringField
 
 
 class ProtobufPythonInterpreterCompatibility(PythonInterpreterCompatibility):
     alias = "python_compatibility"
 
 
+class PythonSourceRootField(StringField):
+    """The source root to generate Python sources under.
+
+    If unspecified, the source root the protobuf_library is under will be used.
+    """
+
+    alias = "python_source_root"
+
+
 def rules():
-    return [ProtobufLibrary.register_plugin_field(ProtobufPythonInterpreterCompatibility)]
+    return [
+        ProtobufLibrary.register_plugin_field(ProtobufPythonInterpreterCompatibility),
+        ProtobufLibrary.register_plugin_field(PythonSourceRootField),
+    ]

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

@@ -1,7 +1,9 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+from pathlib import PurePath
 
 from pants.backend.codegen.protobuf.protoc import Protoc
+from pants.backend.codegen.protobuf.python.additional_fields import PythonSourceRootField
 from pants.backend.codegen.protobuf.target_types import ProtobufSources
 from pants.backend.python.target_types import PythonSources
 from pants.core.util_rules.determine_source_files import AllSourceFilesRequest, SourceFiles
@@ -105,11 +107,20 @@ async def generate_python_from_protobuf(
     )
 
     # We must do some path manipulation on the output digest for it to look like normal sources,
-    # including adding back the original source root.
+    # including adding back a source root.
+    py_source_root = request.protocol_target.get(PythonSourceRootField).value
+    if py_source_root:
+        # Verify that the python source root specified by the target is in fact a source root.
+        source_root_request = SourceRootRequest(PurePath(py_source_root))
+    else:
+        # The target didn't specify a python source root, so use the protobuf_library's source root.
+        source_root_request = SourceRootRequest.for_target(request.protocol_target)
+
     normalized_digest, source_root = await MultiGet(
         Get(Digest, RemovePrefix(result.output_digest, output_dir)),
-        Get(SourceRoot, SourceRootRequest, SourceRootRequest.for_target(request.protocol_target)),
+        Get(SourceRoot, SourceRootRequest, source_root_request),
     )
+
     source_root_restored = (
         await Get(Snapshot, AddPrefix(normalized_digest, source_root.path))
         if source_root.path != "."

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

@@ -4,18 +4,23 @@
 from textwrap import dedent
 from typing import List
 
+import pytest
+
+from pants.backend.codegen.protobuf.python import additional_fields
 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, strip_source_roots
 from pants.engine.addresses import Address
+from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.rules import RootRule
 from pants.engine.target import (
     GeneratedSources,
     HydratedSources,
     HydrateSourcesRequest,
     WrappedTarget,
 )
+from pants.source.source_root import NoSourceRootError
 from pants.testutil.engine.util import Params
 from pants.testutil.external_tool_test_base import ExternalToolTestBase
 from pants.testutil.option.util import create_options_bootstrapper
@@ -31,6 +36,7 @@ def rules(cls):
         return (
             *super().rules(),
             *protobuf_rules(),
+            *additional_fields.rules(),
             *determine_source_files.rules(),
             *strip_source_roots.rules(),
             RootRule(GeneratePythonFromProtobufRequest),
@@ -106,7 +112,8 @@ def test_generates_python(self) -> None:
             ),
         )
         self.add_to_build_file(
-            "src/protobuf/dir2", "protobuf_library(dependencies=['src/protobuf/dir1'])"
+            "src/protobuf/dir2",
+            "protobuf_library(dependencies=['src/protobuf/dir1'], python_source_root='src/python')",
         )
 
         # Test another source root.
@@ -126,7 +133,7 @@ def test_generates_python(self) -> None:
             "tests/protobuf/test_protos", "protobuf_library(dependencies=['src/protobuf/dir2'])"
         )
 
-        source_roots = ["/src/protobuf", "/tests/protobuf"]
+        source_roots = ["src/python", "/src/protobuf", "/tests/protobuf"]
         self.assert_files_generated(
             "src/protobuf/dir1",
             source_roots=source_roots,
@@ -135,15 +142,15 @@ def test_generates_python(self) -> None:
         self.assert_files_generated(
             "src/protobuf/dir2",
             source_roots=source_roots,
-            expected_files=["src/protobuf/dir2/f_pb2.py"],
+            expected_files=["src/python/dir2/f_pb2.py"],
         )
         self.assert_files_generated(
             "tests/protobuf/test_protos",
             source_roots=source_roots,
             expected_files=["tests/protobuf/test_protos/f_pb2.py"],
         )
 
-    def test_top_level_source_root(self) -> None:
+    def test_top_level_proto_root(self) -> None:
         self.create_file(
             "protos/f.proto",
             dedent(
@@ -158,3 +165,40 @@ def test_top_level_source_root(self) -> None:
         self.assert_files_generated(
             "protos", source_roots=["/"], expected_files=["protos/f_pb2.py"]
         )
+
+    def test_top_level_python_source_root(self) -> None:
+        self.create_file(
+            "src/proto/protos/f.proto",
+            dedent(
+                """\
+                syntax = "proto2";
+
+                package protos;
+                """
+            ),
+        )
+        self.add_to_build_file("src/proto/protos", "protobuf_library(python_source_root='.')")
+        self.assert_files_generated(
+            "src/proto/protos", source_roots=["/", "src/proto"], expected_files=["protos/f_pb2.py"]
+        )
+
+    def test_bad_python_source_root(self) -> None:
+        self.create_file(
+            "src/protobuf/dir1/f.proto",
+            dedent(
+                """\
+                syntax = "proto2";
+
+                package dir1;
+                """
+            ),
+        )
+        self.add_to_build_file(
+            "src/protobuf/dir1", "protobuf_library(python_source_root='notasourceroot')"
+        )
+        with pytest.raises(ExecutionError) as exc:
+            self.assert_files_generated(
+                "src/protobuf/dir1", source_roots=["src/protobuf"], expected_files=[]
+            )
+        assert len(exc.value.wrapped_exceptions) == 1
+        assert isinstance(exc.value.wrapped_exceptions[0], NoSourceRootError)