Merge "Make protobufs respect pkg_path properly"
diff --git a/python/python.go b/python/python.go
index 8364169..daf7c14 100644
--- a/python/python.go
+++ b/python/python.go
@@ -120,6 +120,15 @@
 	// whether the binary is required to be built with embedded launcher for this actual_version.
 	// this is set by the python version mutator based on version-specific properties
 	Embedded_launcher *bool `blueprint:"mutated"`
+
+	Proto struct {
+		// Whether generated python protos should include the pkg_path in
+		// their import statements. This is a temporary flag to help transition to
+		// the new behavior where this is always true. It will be removed after all
+		// usages of protos with pkg_path have been updated. The default is currently
+		// false.
+		Respect_pkg_path *bool
+	}
 }
 
 type baseAttributes struct {
@@ -672,8 +681,26 @@
 		protoFlags := android.GetProtoFlags(ctx, &p.protoProperties)
 		protoFlags.OutTypeFlag = "--python_out"
 
+		// TODO(b/247578564): Change the default to true, and then eventually remove respect_pkg_path
+		protosRespectPkgPath := proptools.BoolDefault(p.properties.Proto.Respect_pkg_path, false)
+		pkgPathForProtos := pkgPath
+		if pkgPathForProtos != "" && protosRespectPkgPath {
+			pkgPathStagingDir := android.PathForModuleGen(ctx, "protos_staged_for_pkg_path")
+			rule := android.NewRuleBuilder(pctx, ctx)
+			var stagedProtoSrcs android.Paths
+			for _, srcFile := range protoSrcs {
+				stagedProtoSrc := pkgPathStagingDir.Join(ctx, pkgPath, srcFile.Rel())
+				rule.Command().Text("mkdir -p").Flag(filepath.Base(stagedProtoSrc.String()))
+				rule.Command().Text("cp -f").Input(srcFile).Output(stagedProtoSrc)
+				stagedProtoSrcs = append(stagedProtoSrcs, stagedProtoSrc)
+			}
+			rule.Build("stage_protos_for_pkg_path", "Stage protos for pkg_path")
+			protoSrcs = stagedProtoSrcs
+			pkgPathForProtos = ""
+		}
+
 		for _, srcFile := range protoSrcs {
-			zip := genProto(ctx, srcFile, protoFlags, pkgPath)
+			zip := genProto(ctx, srcFile, protoFlags, pkgPathForProtos)
 			zips = append(zips, zip)
 		}
 	}
diff --git a/python/tests/proto_pkg_path/Android.bp b/python/tests/proto_pkg_path/Android.bp
new file mode 100644
index 0000000..17afde2
--- /dev/null
+++ b/python/tests/proto_pkg_path/Android.bp
@@ -0,0 +1,13 @@
+python_test_host {
+    name: "py_proto_pkg_path_test",
+    main: "main.py",
+    srcs: [
+        "main.py",
+        "proto/*.proto",
+    ],
+    pkg_path: "mylib/subpackage",
+    proto: {
+        canonical_path_from_root: false,
+        respect_pkg_path: true,
+    },
+}
diff --git a/python/tests/proto_pkg_path/main.py b/python/tests/proto_pkg_path/main.py
new file mode 100644
index 0000000..c4acdde
--- /dev/null
+++ b/python/tests/proto_pkg_path/main.py
@@ -0,0 +1,18 @@
+import sys
+
+import unittest
+import mylib.subpackage.proto.test_pb2 as test_pb2
+import mylib.subpackage.proto.common_pb2 as common_pb2
+
+print(sys.path)
+
+class TestProtoWithPkgPath(unittest.TestCase):
+
+    def test_main(self):
+        x = test_pb2.MyMessage(name="foo",
+                               common = common_pb2.MyCommonMessage(common="common"))
+        self.assertEqual(x.name, "foo")
+        self.assertEqual(x.common.common, "common")
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/python/tests/proto_pkg_path/proto/common.proto b/python/tests/proto_pkg_path/proto/common.proto
new file mode 100644
index 0000000..b24b8ea
--- /dev/null
+++ b/python/tests/proto_pkg_path/proto/common.proto
@@ -0,0 +1,5 @@
+syntax = "proto3";
+
+message MyCommonMessage {
+  string common = 1;
+}
diff --git a/python/tests/proto_pkg_path/proto/test.proto b/python/tests/proto_pkg_path/proto/test.proto
new file mode 100644
index 0000000..55f3b17
--- /dev/null
+++ b/python/tests/proto_pkg_path/proto/test.proto
@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+import "mylib/subpackage/proto/common.proto";
+
+message MyMessage {
+  string name = 1;
+  MyCommonMessage common = 2;
+}