Generate imports for each proto src, as we should, and update tests accordingly

Author: shaldengeki

gazelle/python/generate.go

@@ -42,8 +42,8 @@ const (
 	pyTestEntrypointTargetname  = "__test__"
 	conftestFilename            = "conftest.py"
 	conftestTargetname          = "conftest"
-	protoKey                    = "_proto_target"
 	protoRelKey                 = "_proto_rel"
+	protoSrcsKey                = "_proto_srcs"
 )
 
 var (
@@ -575,13 +575,15 @@ func generateProtoLibraries(args language.GenerateArgs, cfg *pythonconfig.Config
 	// 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)
 
@@ -614,8 +616,9 @@ func generateProtoLibraries(args language.GenerateArgs, cfg *pythonconfig.Config
 			addVisibility(visibility).
 			addResolvedDependency(":" + protoRuleName).
 			generateImportsAttribute().build()
-		pyProtoLibrary.SetPrivateAttr(protoKey, protoRuleName)
+
 		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))

gazelle/python/resolve.go

@@ -134,17 +134,17 @@ func importsProtoLibrary(cfg *pythonconfig.Config, r *rule.Rule) []resolve.Impor
 
 	// Determine the root module and emit an import for that,
 	// i.e. for //foo:foo_py_pb2, we'd get foo.foo_pb2
-	protoRuleAttr := r.PrivateAttr(protoKey)
 	protoRelAttr := r.PrivateAttr(protoRelKey)
-	if protoRuleAttr == nil || protoRelAttr == nil {
+	protoSrcsAttr := r.PrivateAttr(protoSrcsKey)
+	if protoRelAttr == nil || protoSrcsAttr == nil {
 		return nil
 	}
 
-	protoRule := protoRuleAttr.(string)
-	generatedPbFileName := strings.TrimSuffix(protoRule, "_proto") + "_pb2.py"
 	protoRel := protoRelAttr.(string)
-
-	specs = append(specs, importSpecFromSrc(cfg.PythonProjectRoot(), protoRel, generatedPbFileName))
+	for _, protoSrc := range protoSrcsAttr.([]string) {
+		generatedPbFileName := strings.TrimSuffix(protoSrc, ".proto") + "_pb2.py"
+		specs = append(specs, importSpecFromSrc(cfg.PythonProjectRoot(), protoRel, generatedPbFileName))
+	}
 
 	return specs
 }

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

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

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

@@ -1 +1 @@
-# gazelle:resolve py test2_generates_using_resolve.foo.test2_generates_using_resolve_foo_pb2 //test2_generates_using_resolve/bar:bar_py_pb2
+# 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

@@ -1,6 +1,6 @@
 load("@rules_python//python:defs.bzl", "py_library")
 
-# gazelle:resolve py test2_generates_using_resolve.foo.test2_generates_using_resolve_foo_pb2 //test2_generates_using_resolve/bar:bar_py_pb2
+# 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",

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

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

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

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