From cb847638af5a56969929ddeae57d407fc8ad98c0 Mon Sep 17 00:00:00 2001
From: Spandan Das <spandandas@google.com>
Date: Tue, 22 Aug 2023 19:24:07 +0000
Subject: [PATCH] Translate python_libray.pkg_path to proto.import_prefix

If a python_library uses a pkg_path foo/bar, then the proto srcs in
that libray need to import the dep .proto as foo/bar/proto.proto.

This behavior is restricted to python modules. To implement this is in
bp2build, this CL creates a new interface with a single method
`PkgPath`. Only python module structs implement this interface, and
this method is only available during bp2build

Test: Added a bp2build unit test
Test: TH

Change-Id: If8d207c0b321f75337a053795826b283a5eaaf46
---
 android/proto.go                           | 12 +++++++
 bp2build/python_library_conversion_test.go | 40 ++++++++++++++++++++++
 python/python.go                           |  8 +++++
 3 files changed, 60 insertions(+)

diff --git a/android/proto.go b/android/proto.go
index 6887900c4..b9365a70f 100644
--- a/android/proto.go
+++ b/android/proto.go
@@ -281,6 +281,13 @@ func Bp2buildProtoProperties(ctx Bp2buildMutatorContext, m *ModuleBase, srcs baz
 				}
 			}
 
+			if p, ok := m.module.(PkgPathInterface); ok && p.PkgPath(ctx) != nil {
+				// python_library with pkg_path
+				// proto_library for this module should have the pkg_path as the import_prefix
+				attrs.Import_prefix = p.PkgPath(ctx)
+				attrs.Strip_import_prefix = proptools.StringPtr("")
+			}
+
 			tags := ApexAvailableTagsWithoutTestApexes(ctx.(TopDownMutatorContext), ctx.Module())
 
 			moduleDir := ctx.ModuleDir()
@@ -333,6 +340,11 @@ func Bp2buildProtoProperties(ctx Bp2buildMutatorContext, m *ModuleBase, srcs baz
 	return info, true
 }
 
+// PkgPathInterface is used as a type assertion in bp2build to get pkg_path property of python_library_host
+type PkgPathInterface interface {
+	PkgPath(ctx BazelConversionContext) *string
+}
+
 var (
 	protoIncludeDirGeneratedSuffix = ".include_dir_bp2build_generated_proto"
 	protoIncludeDirsBp2buildKey    = NewOnceKey("protoIncludeDirsBp2build")
diff --git a/bp2build/python_library_conversion_test.go b/bp2build/python_library_conversion_test.go
index a53371d2b..595acd2e4 100644
--- a/bp2build/python_library_conversion_test.go
+++ b/bp2build/python_library_conversion_test.go
@@ -348,3 +348,43 @@ func TestPythonLibraryWithProtobufs(t *testing.T) {
 		},
 	})
 }
+
+func TestPythonLibraryWithProtobufsAndPkgPath(t *testing.T) {
+	t.Parallel()
+	runBp2BuildTestCaseWithPythonLibraries(t, Bp2buildTestCase{
+		Description: "test python_library protobuf with pkg_path",
+		Filesystem: map[string]string{
+			"dir/foo.proto": "",
+			"dir/bar.proto": "", // bar contains "import dir/foo.proto"
+			"dir/Android.bp": `
+python_library {
+  name: "foo",
+  pkg_path: "dir",
+  srcs: [
+    "foo.proto",
+    "bar.proto",
+  ],
+  bazel_module: {bp2build_available: true},
+}`,
+		},
+		Dir: "dir",
+		ExpectedBazelTargets: []string{
+			MakeBazelTarget("proto_library", "foo_proto", AttrNameToString{
+				"import_prefix":       `"dir"`,
+				"strip_import_prefix": `""`,
+				"srcs": `[
+        "foo.proto",
+        "bar.proto",
+    ]`,
+			}),
+			MakeBazelTarget("py_proto_library", "foo_py_proto", AttrNameToString{
+				"deps": `[":foo_proto"]`,
+			}),
+			MakeBazelTarget("py_library", "foo", AttrNameToString{
+				"srcs_version": `"PY3"`,
+				"imports":      `[".."]`,
+				"deps":         `[":foo_py_proto"]`,
+			}),
+		},
+	})
+}
diff --git a/python/python.go b/python/python.go
index 6c837a888..7d77ca772 100644
--- a/python/python.go
+++ b/python/python.go
@@ -197,6 +197,14 @@ func (p *PythonLibraryModule) getPkgPath() string {
 	return String(p.properties.Pkg_path)
 }
 
+// PkgPath is the "public" version of `getPkgPath` that is only available during bp2build
+func (p *PythonLibraryModule) PkgPath(ctx android.BazelConversionContext) *string {
+	if ctx.Config().BuildMode != android.Bp2build {
+		ctx.ModuleErrorf("PkgPath is only supported in bp2build mode")
+	}
+	return p.properties.Pkg_path
+}
+
 func (p *PythonLibraryModule) getBaseProperties() *BaseProperties {
 	return &p.properties
 }