feat(gazelle): Gazelle resolves module-level imports of py_proto_library

CHANGELOG.md

@@ -122,6 +122,9 @@ END_UNRELEASED_TEMPLATE
   Gazelle generates `py_proto_library` rules for `proto_library`. `false` by default.
 * (gazelle) New directive `gazelle:python_proto_naming_convention`; controls
   naming of `py_proto_library` rules.
+* (gazelle) Gazelle now resolves dependencies for `py_proto_library`
+  module-level imports, i.e. `import some.package.foo_pb2`. Imports of
+  messages/enums/services inside modules are not yet supported.
 
 {#v0-0-0-removed}
 ### Removed

gazelle/README.md

@@ -288,7 +288,7 @@ py_proto_library(
 
 The default naming convention is `$proto_name$_pb2_py`, so by default in the above example Gazelle would generate `foo_pb2_py`. Any pre-existing rules are left in place and not renamed.
 
-Note that the Python library will always be imported as `foo_pb2` in Python code, regardless of the naming convention. Also note that Gazelle is currently not able to map said imports, e.g. `import foo_pb2`, to fill in `py_proto_library` targets as dependencies of other rules. See [this issue](https://github.com/bazel-contrib/rules_python/issues/1703).
+Note that the Python library will always be imported as `foo_pb2` in Python code, regardless of the naming convention.
 
 #### Directive: `python_default_visibility`:
 

gazelle/python/generate.go

@@ -42,6 +42,8 @@ const (
 	pyTestEntrypointTargetname  = "__test__"
 	conftestFilename            = "conftest.py"
 	conftestTargetname          = "conftest"
+	protoRelKey                 = "_proto_rel"
+	protoSrcsKey                = "_proto_srcs"
 )
 
 var (
@@ -572,11 +574,16 @@ func ensureNoCollision(file *rule.File, targetName, kind string) error {
 func generateProtoLibraries(args language.GenerateArgs, cfg *pythonconfig.Config, pythonProjectRoot string, visibility []string, res *language.GenerateResult) {
 	// First, enumerate all the proto_library in this package.
 	var protoRuleNames []string
+	protoRel := map[string]string{}
+	protoSrcs := map[string][]string{}
+
 	for _, r := range args.OtherGen {
 		if r.Kind() != "proto_library" {
 			continue
 		}
 		protoRuleNames = append(protoRuleNames, r.Name())
+		protoRel[r.Name()] = args.Rel
+		protoSrcs[r.Name()] = r.AttrStrings("srcs")
 	}
 	sort.Strings(protoRuleNames)
 
@@ -610,6 +617,9 @@ func generateProtoLibraries(args language.GenerateArgs, cfg *pythonconfig.Config
 			addResolvedDependency(":" + protoRuleName).
 			generateImportsAttribute().build()
 
+		pyProtoLibrary.SetPrivateAttr(protoRelKey, protoRel[protoRuleName])
+		pyProtoLibrary.SetPrivateAttr(protoSrcsKey, protoSrcs[protoRuleName])
+
 		res.Gen = append(res.Gen, pyProtoLibrary)
 		res.Imports = append(res.Imports, pyProtoLibrary.PrivateAttr(config.GazelleImportsKey))
 		pyProtoRules[pyProtoLibrary.Name()] = true

gazelle/python/resolve.go

@@ -57,7 +57,18 @@ func (*Resolver) Name() string { return languageName }
 func (py *Resolver) Imports(c *config.Config, r *rule.Rule, f *rule.File) []resolve.ImportSpec {
 	cfgs := c.Exts[languageName].(pythonconfig.Configs)
 	cfg := cfgs[f.Pkg]
+
 	srcs := r.AttrStrings("srcs")
+	if srcs != nil {
+		return importsSrcLibrary(cfg, srcs, f)
+	} else if isProtoLibrary(r) {
+		return importsProtoLibrary(cfg, r)
+	}
+
+	return nil
+}
+
+func importsSrcLibrary(cfg *pythonconfig.Config, srcs []string, f *rule.File) []resolve.ImportSpec {
 	provides := make([]resolve.ImportSpec, 0, len(srcs)+1)
 	for _, src := range srcs {
 		ext := filepath.Ext(src)
@@ -114,6 +125,30 @@ func importSpecFromSrc(pythonProjectRoot, bzlPkg, src string) resolve.ImportSpec
 	}
 }
 
+func isProtoLibrary(r *rule.Rule) bool {
+	return r.Kind() == pyProtoLibraryKind
+}
+
+func importsProtoLibrary(cfg *pythonconfig.Config, r *rule.Rule) []resolve.ImportSpec {
+	specs := []resolve.ImportSpec{}
+
+	// Determine the root module and emit an import for that,
+	// i.e. for //foo:foo_py_pb2, we'd get foo.foo_pb2
+	protoRelAttr := r.PrivateAttr(protoRelKey)
+	protoSrcsAttr := r.PrivateAttr(protoSrcsKey)
+	if protoRelAttr == nil || protoSrcsAttr == nil {
+		return nil
+	}
+
+	protoRel := protoRelAttr.(string)
+	for _, protoSrc := range protoSrcsAttr.([]string) {
+		generatedPbFileName := strings.TrimSuffix(protoSrc, ".proto") + "_pb2.py"
+		specs = append(specs, importSpecFromSrc(cfg.PythonProjectRoot(), protoRel, generatedPbFileName))
+	}
+
+	return specs
+}
+
 // Embeds returns a list of labels of rules that the given rule embeds. If
 // a rule is embedded by another importable rule of the same language, only
 // the embedding rule will be indexed. The embedding rule will inherit
@@ -210,9 +245,9 @@ func (py *Resolver) Resolve(
 					baseParts = pkgParts[:len(pkgParts)-(relativeDepth-1)]
 				}
 				// Build absolute module path
-				absParts := append([]string{}, baseParts...)       // base path
-				absParts = append(absParts, fromParts...)          // subpath from 'from'
-				absParts = append(absParts, imported)              // actual imported symbol
+				absParts := append([]string{}, baseParts...) // base path
+				absParts = append(absParts, fromParts...)    // subpath from 'from'
+				absParts = append(absParts, imported)        // actual imported symbol
 
 				moduleName = strings.Join(absParts, ".")
 			}

gazelle/python/testdata/directive_python_proto_naming_convention/README.md

@@ -6,4 +6,4 @@ correctly:
 1.  Has no effect on pre-existing `py_proto_library` when `gazelle:python_generate_proto` is disabled.
 2.  Uses the default value when proto generation is on and `python_proto_naming_convention` is not set.
 3.  Uses the provided naming convention when proto generation is on and `python_proto_naming_convention` is set.
-4.  With a pre-existing `py_proto_library` not following a given naming convention, keeps it intact and does not rename it.
+4.  With a pre-existing `py_proto_library` not following a given naming convention, keeps it intact and does not rename it.

gazelle/python/testdata/resolves_proto_imports/README.md

@@ -0,0 +1,8 @@
+# Resolves proto imports
+
+This test asserts that Gazelle can resolve imports from `py_proto_library` targets:
+
+1.  Generates a dependency in the default case.
+2.  Uses `gazelle:resolve` to generate dependencies.
+3.  Uses `python_proto_naming_convention` to generate dependencies.
+4.  Generates a correct dependency for a proto_library with multiple srcs.

gazelle/python/testdata/resolves_proto_imports/test1_generates_dependency/BUILD.out

@@ -0,0 +1,8 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+py_library(
+    name = "test1_generates_dependency",
+    srcs = ["bar.py"],
+    visibility = ["//:__subpackages__"],
+    deps = ["//test1_generates_dependency/foo:test1_generates_dependency_foo_py_pb2"],
+)

gazelle/python/testdata/resolves_proto_imports/test1_generates_dependency/bar.py

@@ -0,0 +1,3 @@
+import test1_generates_dependency.foo.foo_pb2
+
+x = foo_pb2.Foo()

gazelle/python/testdata/resolves_proto_imports/test1_generates_dependency/foo/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_generate_proto true

gazelle/python/testdata/resolves_proto_imports/test1_generates_dependency/foo/BUILD.out

@@ -0,0 +1,16 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "test1_generates_dependency_foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//visibility:public"],
+)
+
+py_proto_library(
+    name = "test1_generates_dependency_foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":test1_generates_dependency_foo_proto"],
+)

gazelle/python/testdata/resolves_proto_imports/test1_generates_dependency/foo/foo.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package test1_generates_dependency.foo;
+
+message Foo {
+    bool bar = 1;
+}

gazelle/python/testdata/resolves_proto_imports/test2_generates_using_resolve/BUILD.in

@@ -0,0 +1 @@
+# gazelle:resolve py test2_generates_using_resolve.bar.bar_pb2 //test2_generates_using_resolve/bar:bar_py_pb2

gazelle/python/testdata/resolves_proto_imports/test2_generates_using_resolve/BUILD.out

@@ -0,0 +1,10 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+# gazelle:resolve py test2_generates_using_resolve.bar.bar_pb2 //test2_generates_using_resolve/bar:bar_py_pb2
+
+py_library(
+    name = "test2_generates_using_resolve",
+    srcs = ["baz.py"],
+    visibility = ["//:__subpackages__"],
+    deps = ["//test2_generates_using_resolve/bar:bar_py_pb2"],
+)

gazelle/python/testdata/resolves_proto_imports/test2_generates_using_resolve/bar/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_generate_proto true

gazelle/python/testdata/resolves_proto_imports/test2_generates_using_resolve/bar/BUILD.out

@@ -0,0 +1,16 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "test2_generates_using_resolve_bar_proto",
+    srcs = ["bar.proto"],
+    visibility = ["//visibility:public"],
+)
+
+py_proto_library(
+    name = "test2_generates_using_resolve_bar_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":test2_generates_using_resolve_bar_proto"],
+)

gazelle/python/testdata/resolves_proto_imports/test2_generates_using_resolve/bar/bar.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package test2_generates_using_resolve.bar;
+
+message Bar {
+    bool bar = 1;
+}

gazelle/python/testdata/resolves_proto_imports/test2_generates_using_resolve/baz.py

@@ -0,0 +1,3 @@
+import test2_generates_using_resolve.bar.bar_pb2
+
+x = bar_pb2.Bar()

gazelle/python/testdata/resolves_proto_imports/test3_uses_naming_convention/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_proto_naming_convention some_$proto_name$_value

gazelle/python/testdata/resolves_proto_imports/test3_uses_naming_convention/BUILD.out

@@ -0,0 +1,10 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+# gazelle:python_proto_naming_convention some_$proto_name$_value
+
+py_library(
+    name = "test3_uses_naming_convention",
+    srcs = ["baz.py"],
+    visibility = ["//:__subpackages__"],
+    deps = ["//test3_uses_naming_convention/bar:some_test3_uses_naming_convention_bar_value"],
+)

gazelle/python/testdata/resolves_proto_imports/test3_uses_naming_convention/bar/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_generate_proto true

gazelle/python/testdata/resolves_proto_imports/test3_uses_naming_convention/bar/BUILD.out

@@ -0,0 +1,16 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "test3_uses_naming_convention_bar_proto",
+    srcs = ["bar.proto"],
+    visibility = ["//visibility:public"],
+)
+
+py_proto_library(
+    name = "some_test3_uses_naming_convention_bar_value",
+    visibility = ["//:__subpackages__"],
+    deps = [":test3_uses_naming_convention_bar_proto"],
+)

gazelle/python/testdata/resolves_proto_imports/test3_uses_naming_convention/bar/bar.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package test3_uses_naming_convention.bar;
+
+message Bar {
+    bool bar = 1;
+}

gazelle/python/testdata/resolves_proto_imports/test3_uses_naming_convention/baz.py

@@ -0,0 +1,3 @@
+import test3_uses_naming_convention.bar.bar_pb2
+
+x = bar_pb2.Bar()

gazelle/python/testdata/resolves_proto_imports/test4_generates_imports_for_multiple_proto_srcs/BUILD.out

@@ -0,0 +1,8 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+py_library(
+    name = "test4_generates_imports_for_multiple_proto_srcs",
+    srcs = ["bar.py"],
+    visibility = ["//:__subpackages__"],
+    deps = ["//test4_generates_imports_for_multiple_proto_srcs/foo:test4_generates_imports_for_multiple_proto_srcs_foo_py_pb2"],
+)

gazelle/python/testdata/resolves_proto_imports/test4_generates_imports_for_multiple_proto_srcs/bar.py

@@ -0,0 +1,5 @@
+import test4_generates_imports_for_multiple_proto_srcs.foo.foo_pb2
+import test4_generates_imports_for_multiple_proto_srcs.foo.bar_pb2
+
+x = foo_pb2.Foo()
+y = bar_pb2.Bar()

gazelle/python/testdata/resolves_proto_imports/test4_generates_imports_for_multiple_proto_srcs/foo/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_generate_proto true

gazelle/python/testdata/resolves_proto_imports/test4_generates_imports_for_multiple_proto_srcs/foo/BUILD.out

@@ -0,0 +1,19 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "test4_generates_imports_for_multiple_proto_srcs_foo_proto",
+    srcs = [
+        "bar.proto",
+        "foo.proto",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+py_proto_library(
+    name = "test4_generates_imports_for_multiple_proto_srcs_foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":test4_generates_imports_for_multiple_proto_srcs_foo_proto"],
+)

gazelle/python/testdata/resolves_proto_imports/test4_generates_imports_for_multiple_proto_srcs/foo/bar.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package test4_generates_imports_for_multiple_proto_srcs.foo;
+
+message Bar {
+    bool bar = 1;
+}

gazelle/python/testdata/resolves_proto_imports/test4_generates_imports_for_multiple_proto_srcs/foo/foo.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package test4_generates_imports_for_multiple_proto_srcs.foo;
+
+message Foo {
+    bool bar = 1;
+}