Update the strip prefix logic for python libraries (#361)

* Fix imports directive for py_library when using stripPrefix

Signed-off-by: Vihang Mehta <vihang@gimletlabs.ai>

* Tests

Signed-off-by: Vihang Mehta <vihang@gimletlabs.ai>

---------

Signed-off-by: Vihang Mehta <vihang@gimletlabs.ai>

example/golden/testdata/strip_import_prefix/module_lib/util/nested/prefix/BUILD.in

@@ -0,0 +1 @@
+# gazelle:proto_strip_import_prefix /module_lib/util/nested

example/golden/testdata/strip_import_prefix/module_lib/util/nested/prefix/BUILD.out

@@ -0,0 +1,28 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@build_stack_rules_proto//rules/py:proto_py_library.bzl", "proto_py_library")
+load("@build_stack_rules_proto//rules:proto_compile.bzl", "proto_compile")
+
+# gazelle:proto_strip_import_prefix /module_lib/util/nested
+
+proto_library(
+    name = "prefix_test_proto",
+    srcs = ["test.proto"],
+    strip_import_prefix = "/module_lib/util/nested",
+    visibility = ["//visibility:public"],
+)
+
+proto_compile(
+    name = "prefix_test_python_compile",
+    output_mappings = ["test_pb2.py=/prefix/test_pb2.py"],
+    outputs = ["test_pb2.py"],
+    plugins = ["@build_stack_rules_proto//plugin/builtin:python"],
+    proto = "prefix_test_proto",
+)
+
+proto_py_library(
+    name = "prefix_test_py_library",
+    srcs = ["test_pb2.py"],
+    imports = [".."],
+    visibility = ["//visibility:public"],
+    deps = ["@com_google_protobuf//:protobuf_python"],
+)

example/golden/testdata/strip_import_prefix/module_lib/util/nested/prefix/test.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package prefix.test;
+
+message Data {
+    int64 id = 1;
+}

pkg/rule/rules_python/proto_py_library_test.go

@@ -64,12 +64,13 @@ proto_py_library(
 						Outputs: []string{"test_pb2.py"},
 					},
 				},
+                Rel: "com/foo/baz/qux/v1",
 			},
 			want: `
 proto_py_library(
     name = "test_py_library",
     srcs = ["test_pb2.py"],
-    imports = ["../.."],
+    imports = ["../../.."],
 )
 `,
 		},

pkg/rule/rules_python/py_library.go

@@ -1,6 +1,7 @@
 package rules_python
 
 import (
+	"path/filepath"
 	"strings"
 
 	"github.com/bazelbuild/bazel-gazelle/config"
@@ -71,17 +72,23 @@ func (s *PyLibrary) Visibility() []string {
 func (s *PyLibrary) ImportsAttr() (imps []string) {
 	// if we have a strip_import_prefix on the proto_library, the python search
 	// path should include the directory N parents above the current package,
-	// where N is the number of segments in an absolute strip_import_prefix
+	// where N is the number of segments needed to ascend to the prefix from
+	// the dir for the current rule.
 	if s.Config.Library.StripImportPrefix() == "" {
 		return
 	}
 	prefix := s.Config.Library.StripImportPrefix()
 	if !strings.HasPrefix(prefix, "/") {
 		return // deal with relative-imports at another time
 	}
+
 	prefix = strings.TrimPrefix(prefix, "/")
-	prefix = strings.TrimSuffix(prefix, "/")
-	parts := strings.Split(prefix, "/")
+	rel, err := filepath.Rel(prefix, s.Config.Rel)
+	if err != nil {
+		return // the prefix doesn't prefix the current path, shouldn't happen
+	}
+
+	parts := strings.Split(rel, "/")
 	for i := 0; i < len(parts); i++ {
 		parts[i] = ".."
 	}