[wip] Begin implementation of module-level py_proto_library import support

Author: shaldengeki

gazelle/python/generate.go

@@ -42,6 +42,8 @@ const (
 	pyTestEntrypointTargetname  = "__test__"
 	conftestFilename            = "conftest.py"
 	conftestTargetname          = "conftest"
+	protoKey                    = "_proto_target"
+	protoRelKey                 = "_proto_rel"
 )
 
 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
+	protoRules := map[string]*rule.Rule{}
+	protoRel := map[string]string{}
+
 	for _, r := range args.OtherGen {
 		if r.Kind() != "proto_library" {
 			continue
 		}
 		protoRuleNames = append(protoRuleNames, r.Name())
+		protoRules[r.Name()] = r
+		protoRel[r.Name()] = args.Rel
 	}
 	sort.Strings(protoRuleNames)
 
@@ -609,6 +616,8 @@ func generateProtoLibraries(args language.GenerateArgs, cfg *pythonconfig.Config
 			addVisibility(visibility).
 			addResolvedDependency(":" + protoRuleName).
 			generateImportsAttribute().build()
+		pyProtoLibrary.SetPrivateAttr(protoKey, protoRuleName)
+		pyProtoLibrary.SetPrivateAttr(protoRelKey, protoRel[protoRuleName])
 
 		res.Gen = append(res.Gen, pyProtoLibrary)
 		res.Imports = append(res.Imports, pyProtoLibrary.PrivateAttr(config.GazelleImportsKey))

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, f)
+	}
+
+	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,42 @@ 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, f *rule.File) []resolve.ImportSpec {
+	specs := []resolve.ImportSpec{}
+
+	// First, 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 {
+		return specs
+	}
+
+	protoRule := protoRuleAttr.(string)
+	generatedPbFileName := strings.TrimSuffix(protoRule, "_proto") + "_pb2.py"
+	protoRel := protoRelAttr.(string)
+
+	specs = append(specs, importSpecFromSrc(cfg.PythonProjectRoot(), protoRel, generatedPbFileName))
+
+	// TODO: use parsed proto FileInfo to enumerate importable constants, like messages,
+	// and emit ImportSpec for them
+	// protoPkg := r.PrivateAttr(proto.PackageKey).(proto.Package)
+	// for _, protoFileInfo := range protoPkg.Files {
+	// 	for _, svc := range protoFileInfo.Services {
+	// 		specs = append(specs, resolve.ImportSpec{
+	// 			Lang: languageName,
+	// 			Imp:  fmt.Sprintf("%s.%s", rootPath, svc),
+	// 		})
+	// 	}
+	// Repeat for Messages, Enums
+
+	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 +257,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, ".")
 			}
@@ -282,6 +329,9 @@ func (py *Resolver) Resolve(
 							// Check if the imported module is part of the standard library.
 							if isStdModule(Module{Name: moduleName}) {
 								continue MODULES_LOOP
+							} else if r.Kind() == pyProtoLibraryKind {
+								// For py_proto_library, fall back to guessing the label based on the proto_library rule name.
+								matches = py.resolveProtoFallback(cfg)
 							} else if cfg.ValidateImportStatements() {
 								err := fmt.Errorf(
 									"%[1]q, line %[2]d: %[3]q is an invalid dependency: possible solutions:\n"+
@@ -372,6 +422,11 @@ func (py *Resolver) Resolve(
 	}
 }
 
+func (*Resolver) resolveProtoFallback(c *pythonconfig.Config) []resolve.FindResult {
+	// TODO
+	return []resolve.FindResult{}
+}
+
 // addResolvedDeps adds the pre-resolved dependencies from the rule's private attributes
 // to the provided deps set.
 func addResolvedDeps(

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/BUILD.in

@@ -0,0 +1,9 @@
+# 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.
+
+[gh-1703]: https://github.com/bazel-contrib/rules_python/issues/1703

gazelle/python/testdata/resolves_proto_imports/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"],
+)