feat(gazelle): Gazelle plugin generates py_proto_library (#3057)

Fixes https://github.com/bazel-contrib/rules_python/issues/2994.

Please go over this with a fine-toothed comb! This is my first
contribution to `rules_python` / the gazelle plugin, and while I've
worked in Gazelle before, I'm pretty unfamiliar with the Python plugin's
architecture.

This adds support in the Gazelle plugin for generating
`py_proto_library` rules automatically, if there are any `proto_library`
rules detected in a given package. We do this via a new Gazelle
directive, `python_generate_proto`, which defaults to `true`, and
controls whether these rules are generated.

See the tests in `testdata/directive_python_generate_proto` for
examples.

By default, we source the `py_proto_library` rule from the `@protobuf`
repository. I think this the intended long-term home of the rule? Users
are expected to use `gazelle:map_kind` to change this if need be.

I haven't done anything here to support resolution of imports of
`py_proto_library`. I think this is worth landing first, to save folks
from having to maintain these by hand. But this should lay the
foundation for resolving that in
https://github.com/bazel-contrib/rules_python/issues/1703.

---------

Co-authored-by: Douglas Thor <dougthor42@users.noreply.github.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 834a2c1..e74f14b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -106,6 +106,8 @@
     * 3.12.11
     * 3.13.5
     * 3.14.0b3
+* (gazelle) New directive `gazelle:python_generate_proto`; when `true`,
+  Gazelle generates `py_proto_library` rules for `proto_library`. `false` by default.
 
 {#v0-0-0-removed}
 ### Removed
diff --git a/examples/bzlmod/py_proto_library/BUILD.bazel b/examples/bzlmod/py_proto_library/BUILD.bazel
index 969cb8e..daea410 100644
--- a/examples/bzlmod/py_proto_library/BUILD.bazel
+++ b/examples/bzlmod/py_proto_library/BUILD.bazel
@@ -6,7 +6,7 @@
     srcs = ["test.py"],
     main = "test.py",
     deps = [
-        "//py_proto_library/example.com/proto:pricetag_proto_py_pb2",
+        "//py_proto_library/example.com/proto:pricetag_py_pb2",
     ],
 )
 
@@ -14,7 +14,7 @@
     name = "message_test",
     srcs = ["message_test.py"],
     deps = [
-        "//py_proto_library/example.com/another_proto:message_proto_py_pb2",
+        "//py_proto_library/example.com/another_proto:message_py_pb2",
     ],
 )
 
diff --git a/examples/bzlmod/py_proto_library/example.com/another_proto/BUILD.bazel b/examples/bzlmod/py_proto_library/example.com/another_proto/BUILD.bazel
index 785d90d..29f08c2 100644
--- a/examples/bzlmod/py_proto_library/example.com/another_proto/BUILD.bazel
+++ b/examples/bzlmod/py_proto_library/example.com/another_proto/BUILD.bazel
@@ -2,7 +2,7 @@
 load("@rules_python//python:proto.bzl", "py_proto_library")
 
 py_proto_library(
-    name = "message_proto_py_pb2",
+    name = "message_py_pb2",
     visibility = ["//visibility:public"],
     deps = [":message_proto"],
 )
diff --git a/examples/bzlmod/py_proto_library/example.com/proto/BUILD.bazel b/examples/bzlmod/py_proto_library/example.com/proto/BUILD.bazel
index 72af672..1f8e8f2 100644
--- a/examples/bzlmod/py_proto_library/example.com/proto/BUILD.bazel
+++ b/examples/bzlmod/py_proto_library/example.com/proto/BUILD.bazel
@@ -2,7 +2,7 @@
 load("@rules_python//python:proto.bzl", "py_proto_library")
 
 py_proto_library(
-    name = "pricetag_proto_py_pb2",
+    name = "pricetag_py_pb2",
     visibility = ["//visibility:public"],
     deps = [":pricetag_proto"],
 )
diff --git a/examples/py_proto_library/BUILD.bazel b/examples/py_proto_library/BUILD.bazel
index d782fb2..b57c528 100644
--- a/examples/py_proto_library/BUILD.bazel
+++ b/examples/py_proto_library/BUILD.bazel
@@ -5,7 +5,7 @@
     srcs = ["test.py"],
     main = "test.py",
     deps = [
-        "//example.com/proto:pricetag_proto_py_pb2",
+        "//example.com/proto:pricetag_py_pb2",
     ],
 )
 
@@ -13,6 +13,6 @@
     name = "message_test",
     srcs = ["message_test.py"],
     deps = [
-        "//example.com/another_proto:message_proto_py_pb2",
+        "//example.com/another_proto:message_py_pb2",
     ],
 )
diff --git a/examples/py_proto_library/example.com/another_proto/BUILD.bazel b/examples/py_proto_library/example.com/another_proto/BUILD.bazel
index 3d84155..55e83a2 100644
--- a/examples/py_proto_library/example.com/another_proto/BUILD.bazel
+++ b/examples/py_proto_library/example.com/another_proto/BUILD.bazel
@@ -2,7 +2,7 @@
 load("@rules_python//python:proto.bzl", "py_proto_library")
 
 py_proto_library(
-    name = "message_proto_py_pb2",
+    name = "message_py_pb2",
     visibility = ["//visibility:public"],
     deps = [":message_proto"],
 )
diff --git a/examples/py_proto_library/example.com/proto/BUILD.bazel b/examples/py_proto_library/example.com/proto/BUILD.bazel
index f84454f..fdf2e6f 100644
--- a/examples/py_proto_library/example.com/proto/BUILD.bazel
+++ b/examples/py_proto_library/example.com/proto/BUILD.bazel
@@ -2,7 +2,7 @@
 load("@rules_python//python:proto.bzl", "py_proto_library")
 
 py_proto_library(
-    name = "pricetag_proto_py_pb2",
+    name = "pricetag_py_pb2",
     visibility = ["//visibility:public"],
     deps = [":pricetag_proto"],
 )
diff --git a/gazelle/README.md b/gazelle/README.md
index 3dc8e12..35a1e4f 100644
--- a/gazelle/README.md
+++ b/gazelle/README.md
@@ -224,6 +224,8 @@
 | Controls whether Gazelle resolves dependencies for import statements that use paths relative to the current package. Can be "true" or "false".|
 | `# gazelle:python_generate_pyi_deps`                                                                                                                                                                                                                                                            | `false` |
 | Controls whether to generate a separate `pyi_deps` attribute for type-checking dependencies or merge them into the regular `deps` attribute. When `false` (default), type-checking dependencies are merged into `deps` for backward compatibility. When `true`, generates separate `pyi_deps`. Imports in blocks with the format `if typing.TYPE_CHECKING:`/`if TYPE_CHECKING:` and type-only stub packages (eg. boto3-stubs) are recognized as type-checking dependencies. |
+| [`# gazelle:python_generate_proto`](#directive-python_generate_proto)                                                                                                                                                                                                                                                                | `false` |
+| Controls whether to generate a `py_proto_library` for each `proto_library` in the package. By default we load this rule from the `@protobuf` repository; use `gazelle:map_kind` if you need to load this from somewhere else. |
 
 #### Directive: `python_root`:
 
@@ -484,6 +486,41 @@
 )
 ```
 
+#### Directive: `python_generate_proto`:
+
+When `# gazelle:python_generate_proto true`, Gazelle will generate one
+`py_proto_library` for each `proto_library`, generating Python clients for
+protobuf in each package. By default this is turned off. Gazelle will also
+generate a load statement for the `py_proto_library` - attempting to detect
+the configured name for the `@protobuf` / `@com_google_protobuf` repo in your
+`MODULE.bazel`, and otherwise falling back to `@com_google_protobuf` for
+compatibility with `WORKSPACE`.
+
+For example, in a package with `# gazelle:python_generate_proto true` and a
+`foo.proto`, if you have both the proto extension and the Python extension
+loaded into Gazelle, you'll get something like:
+
+```starlark
+load("@protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
+```
+
+When `false`, Gazelle will ignore any `py_proto_library`, including previously-generated or hand-created rules.
+
 ### Annotations
 
 *Annotations* refer to comments found _within Python files_ that configure how
diff --git a/gazelle/python/BUILD.bazel b/gazelle/python/BUILD.bazel
index 8e8216d..1a7c54f 100644
--- a/gazelle/python/BUILD.bazel
+++ b/gazelle/python/BUILD.bazel
@@ -34,6 +34,7 @@
         "@bazel_gazelle//config:go_default_library",
         "@bazel_gazelle//label:go_default_library",
         "@bazel_gazelle//language:go_default_library",
+        "@bazel_gazelle//language/proto:go_default_library",
         "@bazel_gazelle//repo:go_default_library",
         "@bazel_gazelle//resolve:go_default_library",
         "@bazel_gazelle//rule:go_default_library",
@@ -91,7 +92,10 @@
 
 gazelle_binary(
     name = "gazelle_binary",
-    languages = [":python"],
+    languages = [
+        "@bazel_gazelle//language/proto",
+        ":python",
+    ],
     visibility = ["//visibility:public"],
 )
 
diff --git a/gazelle/python/configure.go b/gazelle/python/configure.go
index db80fc1..7131be2 100644
--- a/gazelle/python/configure.go
+++ b/gazelle/python/configure.go
@@ -70,6 +70,7 @@
 		pythonconfig.LabelNormalization,
 		pythonconfig.GeneratePyiDeps,
 		pythonconfig.ExperimentalAllowRelativeImports,
+		pythonconfig.GenerateProto,
 	}
 }
 
@@ -237,6 +238,12 @@
 				log.Fatal(err)
 			}
 			config.SetGeneratePyiDeps(v)
+		case pythonconfig.GenerateProto:
+			v, err := strconv.ParseBool(strings.TrimSpace(d.Value))
+			if err != nil {
+				log.Fatal(err)
+			}
+			config.SetGenerateProto(v)
 		}
 	}
 
diff --git a/gazelle/python/generate.go b/gazelle/python/generate.go
index c1edec4..3437435 100644
--- a/gazelle/python/generate.go
+++ b/gazelle/python/generate.go
@@ -226,6 +226,10 @@
 	var result language.GenerateResult
 	result.Gen = make([]*rule.Rule, 0)
 
+	if cfg.GenerateProto() {
+		generateProtoLibraries(args, pythonProjectRoot, visibility, &result)
+	}
+
 	collisionErrors := singlylinkedlist.New()
 
 	appendPyLibrary := func(srcs *treeset.Set, pyLibraryTargetName string) {
@@ -551,3 +555,51 @@
 	}
 	return nil
 }
+
+func generateProtoLibraries(args language.GenerateArgs, pythonProjectRoot string, visibility []string, res *language.GenerateResult) {
+	// First, enumerate all the proto_library in this package.
+	var protoRuleNames []string
+	for _, r := range args.OtherGen {
+		if r.Kind() != "proto_library" {
+			continue
+		}
+		protoRuleNames = append(protoRuleNames, r.Name())
+	}
+	sort.Strings(protoRuleNames)
+
+	// Next, enumerate all the pre-existing py_proto_library in this package, so we can delete unnecessary rules later.
+	pyProtoRules := map[string]bool{}
+	if args.File != nil {
+		for _, r := range args.File.Rules {
+			if r.Kind() == "py_proto_library" {
+				pyProtoRules[r.Name()] = false
+			}
+		}
+	}
+
+	emptySiblings := treeset.Set{}
+	// Generate a py_proto_library for each proto_library.
+	for _, protoRuleName := range protoRuleNames {
+		pyProtoLibraryName := strings.TrimSuffix(protoRuleName, "_proto") + "_py_pb2"
+		pyProtoLibrary := newTargetBuilder(pyProtoLibraryKind, pyProtoLibraryName, pythonProjectRoot, args.Rel, &emptySiblings).
+			addVisibility(visibility).
+			addResolvedDependency(":" + protoRuleName).
+			generateImportsAttribute().build()
+
+		res.Gen = append(res.Gen, pyProtoLibrary)
+		res.Imports = append(res.Imports, pyProtoLibrary.PrivateAttr(config.GazelleImportsKey))
+		pyProtoRules[pyProtoLibrary.Name()] = true
+
+	}
+
+	// Finally, emit an empty rule for each pre-existing py_proto_library that we didn't already generate.
+	for ruleName, generated := range pyProtoRules {
+		if generated {
+			continue
+		}
+
+		emptyRule := newTargetBuilder(pyProtoLibraryKind, ruleName, pythonProjectRoot, args.Rel, &emptySiblings).build()
+		res.Empty = append(res.Empty, emptyRule)
+	}
+
+}
diff --git a/gazelle/python/kinds.go b/gazelle/python/kinds.go
index ff3f6ce..a4ce572 100644
--- a/gazelle/python/kinds.go
+++ b/gazelle/python/kinds.go
@@ -15,13 +15,16 @@
 package python
 
 import (
+	"fmt"
+
 	"github.com/bazelbuild/bazel-gazelle/rule"
 )
 
 const (
-	pyBinaryKind  = "py_binary"
-	pyLibraryKind = "py_library"
-	pyTestKind    = "py_test"
+	pyBinaryKind       = "py_binary"
+	pyLibraryKind      = "py_library"
+	pyProtoLibraryKind = "py_proto_library"
+	pyTestKind         = "py_test"
 )
 
 // Kinds returns a map that maps rule names (kinds) and information on how to
@@ -32,7 +35,7 @@
 
 var pyKinds = map[string]rule.KindInfo{
 	pyBinaryKind: {
-		MatchAny: false,
+		MatchAny:   false,
 		MatchAttrs: []string{"srcs"},
 		NonEmptyAttrs: map[string]bool{
 			"deps":    true,
@@ -45,7 +48,7 @@
 			"srcs": true,
 		},
 		ResolveAttrs: map[string]bool{
-			"deps": true,
+			"deps":     true,
 			"pyi_deps": true,
 		},
 	},
@@ -62,10 +65,16 @@
 			"srcs": true,
 		},
 		ResolveAttrs: map[string]bool{
-			"deps": true,
+			"deps":     true,
 			"pyi_deps": true,
 		},
 	},
+	pyProtoLibraryKind: {
+		NonEmptyAttrs: map[string]bool{
+			"deps": true,
+		},
+		ResolveAttrs: map[string]bool{"deps": true},
+	},
 	pyTestKind: {
 		MatchAny: false,
 		NonEmptyAttrs: map[string]bool{
@@ -79,26 +88,43 @@
 			"srcs": true,
 		},
 		ResolveAttrs: map[string]bool{
-			"deps": true,
+			"deps":     true,
 			"pyi_deps": true,
 		},
 	},
 }
 
+func (py *Python) Loads() []rule.LoadInfo {
+	panic("ApparentLoads should be called instead")
+}
+
 // Loads returns .bzl files and symbols they define. Every rule generated by
 // GenerateRules, now or in the past, should be loadable from one of these
 // files.
-func (py *Python) Loads() []rule.LoadInfo {
-	return pyLoads
+func (py *Python) ApparentLoads(moduleToApparentName func(string) string) []rule.LoadInfo {
+	return apparentLoads(moduleToApparentName)
 }
 
-var pyLoads = []rule.LoadInfo{
-	{
-		Name: "@rules_python//python:defs.bzl",
-		Symbols: []string{
-			pyBinaryKind,
-			pyLibraryKind,
-			pyTestKind,
+func apparentLoads(moduleToApparentName func(string) string) []rule.LoadInfo {
+	protobuf := moduleToApparentName("protobuf")
+	if protobuf == "" {
+		protobuf = "com_google_protobuf"
+	}
+
+	return []rule.LoadInfo{
+		{
+			Name: "@rules_python//python:defs.bzl",
+			Symbols: []string{
+				pyBinaryKind,
+				pyLibraryKind,
+				pyTestKind,
+			},
 		},
-	},
+		{
+			Name: fmt.Sprintf("@%s//bazel:py_proto_library.bzl", protobuf),
+			Symbols: []string{
+				pyProtoLibraryKind,
+			},
+		},
+	}
 }
diff --git a/gazelle/python/testdata/directive_python_generate_proto/README.md b/gazelle/python/testdata/directive_python_generate_proto/README.md
new file mode 100644
index 0000000..54261f4
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/README.md
@@ -0,0 +1,9 @@
+# Directive: `python_generate_proto`
+
+This test case asserts that the `# gazelle:python_generate_proto` directive
+correctly:
+
+1.  Uses the default value when `python_generate_proto` is not set.
+2.  Generates (or not) `py_proto_library` when `python_generate_proto` is set, based on whether a proto is present.
+
+[gh-2994]: https://github.com/bazel-contrib/rules_python/issues/2994
diff --git a/gazelle/python/testdata/directive_python_generate_proto/WORKSPACE b/gazelle/python/testdata/directive_python_generate_proto/WORKSPACE
new file mode 100644
index 0000000..faff6af
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/WORKSPACE
@@ -0,0 +1 @@
+# This is a Bazel workspace for the Gazelle test data.
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test.yaml b/gazelle/python/testdata/directive_python_generate_proto/test.yaml
new file mode 100644
index 0000000..36dd656
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test.yaml
@@ -0,0 +1,3 @@
+---
+expect:
+  exit_code: 0
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/BUILD.in
new file mode 100644
index 0000000..9784aaf
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/BUILD.in
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# python_generate_proto is not set, so py_proto_library is not generated.
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/BUILD.out
new file mode 100644
index 0000000..9784aaf
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/BUILD.out
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# python_generate_proto is not set, so py_proto_library is not generated.
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/foo.proto b/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/foo.proto
new file mode 100644
index 0000000..fe2af27
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test1_default_with_proto/foo.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo;
+
+message Foo {
+    string bar = 1;
+}
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test2_default_without_proto/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test2_default_without_proto/BUILD.in
new file mode 100644
index 0000000..0a869d0
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test2_default_without_proto/BUILD.in
@@ -0,0 +1 @@
+# python_generate_proto is not set, so py_proto_library is not generated.
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test2_default_without_proto/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test2_default_without_proto/BUILD.out
new file mode 100644
index 0000000..0a869d0
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test2_default_without_proto/BUILD.out
@@ -0,0 +1 @@
+# python_generate_proto is not set, so py_proto_library is not generated.
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/BUILD.in
new file mode 100644
index 0000000..62fd4be
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/BUILD.in
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto false
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/BUILD.out
new file mode 100644
index 0000000..62fd4be
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/BUILD.out
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto false
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/foo.proto b/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/foo.proto
new file mode 100644
index 0000000..022e29a
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test3_disabled_with_proto/foo.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo.bar;
+
+message Foo {
+    string bar = 1;
+}
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test4_disabled_without_proto/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test4_disabled_without_proto/BUILD.in
new file mode 100644
index 0000000..b283b5f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test4_disabled_without_proto/BUILD.in
@@ -0,0 +1 @@
+# gazelle:python_generate_proto false
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test4_disabled_without_proto/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test4_disabled_without_proto/BUILD.out
new file mode 100644
index 0000000..b283b5f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test4_disabled_without_proto/BUILD.out
@@ -0,0 +1 @@
+# gazelle:python_generate_proto false
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/BUILD.in
new file mode 100644
index 0000000..4713404
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/BUILD.in
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/BUILD.out
new file mode 100644
index 0000000..686252f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/BUILD.out
@@ -0,0 +1,16 @@
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/foo.proto b/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/foo.proto
new file mode 100644
index 0000000..fe2af27
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test5_enabled_with_proto/foo.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo;
+
+message Foo {
+    string bar = 1;
+}
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test6_enabled_without_proto/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test6_enabled_without_proto/BUILD.in
new file mode 100644
index 0000000..ce3eec6
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test6_enabled_without_proto/BUILD.in
@@ -0,0 +1 @@
+# gazelle:python_generate_proto true
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test6_enabled_without_proto/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test6_enabled_without_proto/BUILD.out
new file mode 100644
index 0000000..ce3eec6
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test6_enabled_without_proto/BUILD.out
@@ -0,0 +1 @@
+# gazelle:python_generate_proto true
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test7_removes_when_unnecessary/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test7_removes_when_unnecessary/BUILD.in
new file mode 100644
index 0000000..686252f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test7_removes_when_unnecessary/BUILD.in
@@ -0,0 +1,16 @@
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test7_removes_when_unnecessary/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test7_removes_when_unnecessary/BUILD.out
new file mode 100644
index 0000000..ce3eec6
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test7_removes_when_unnecessary/BUILD.out
@@ -0,0 +1 @@
+# gazelle:python_generate_proto true
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/BUILD.in
new file mode 100644
index 0000000..f14ed4f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/BUILD.in
@@ -0,0 +1,16 @@
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto false
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/BUILD.out
new file mode 100644
index 0000000..f14ed4f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/BUILD.out
@@ -0,0 +1,16 @@
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto false
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/foo.proto b/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/foo.proto
new file mode 100644
index 0000000..022e29a
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto/test8_disabled_ignores_py_proto_library/foo.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo.bar;
+
+message Foo {
+    string bar = 1;
+}
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/BUILD.in
new file mode 100644
index 0000000..4713404
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/BUILD.in
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/BUILD.out
new file mode 100644
index 0000000..dab84a6
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/BUILD.out
@@ -0,0 +1,16 @@
+load("@protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/MODULE.bazel b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/MODULE.bazel
new file mode 100644
index 0000000..66d64af
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/MODULE.bazel
@@ -0,0 +1 @@
+bazel_dep(name = "protobuf", version = "29.3")
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/README.md b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/README.md
new file mode 100644
index 0000000..2d91ccf
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/README.md
@@ -0,0 +1,6 @@
+# Directive: `python_generate_proto`
+
+This test case asserts that the `# gazelle:python_generate_proto` directive
+correctly reads the name of the protobuf repository when bzlmod is being used.
+
+[gh-2994]: https://github.com/bazel-contrib/rules_python/issues/2994
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/WORKSPACE b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/WORKSPACE
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/WORKSPACE
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/foo.proto b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/foo.proto
new file mode 100644
index 0000000..fe2af27
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/foo.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo;
+
+message Foo {
+    string bar = 1;
+}
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/test.yaml b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/test.yaml
new file mode 100644
index 0000000..36dd656
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf/test.yaml
@@ -0,0 +1,3 @@
+---
+expect:
+  exit_code: 0
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/BUILD.in b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/BUILD.in
new file mode 100644
index 0000000..4713404
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/BUILD.in
@@ -0,0 +1,9 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/BUILD.out b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/BUILD.out
new file mode 100644
index 0000000..686252f
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/BUILD.out
@@ -0,0 +1,16 @@
+load("@com_google_protobuf//bazel:py_proto_library.bzl", "py_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+
+# gazelle:python_generate_proto true
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    visibility = ["//:__subpackages__"],
+    deps = [":foo_proto"],
+)
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/MODULE.bazel b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/MODULE.bazel
new file mode 100644
index 0000000..9ab4c17
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/MODULE.bazel
@@ -0,0 +1 @@
+bazel_dep(name = "protobuf", version = "29.3", repo_name = "com_google_protobuf")
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/README.md b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/README.md
new file mode 100644
index 0000000..7900d49
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/README.md
@@ -0,0 +1,7 @@
+# Directive: `python_generate_proto`
+
+This test case asserts that the `# gazelle:python_generate_proto` directive
+correctly reads the name of the protobuf repository when bzlmod is being used,
+but the repository is renamed.
+
+[gh-2994]: https://github.com/bazel-contrib/rules_python/issues/2994
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/WORKSPACE b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/WORKSPACE
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/WORKSPACE
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/foo.proto b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/foo.proto
new file mode 100644
index 0000000..fe2af27
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/foo.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo;
+
+message Foo {
+    string bar = 1;
+}
diff --git a/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/test.yaml b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/test.yaml
new file mode 100644
index 0000000..36dd656
--- /dev/null
+++ b/gazelle/python/testdata/directive_python_generate_proto_bzlmod_protobuf_renamed/test.yaml
@@ -0,0 +1,3 @@
+---
+expect:
+  exit_code: 0
diff --git a/gazelle/pythonconfig/pythonconfig.go b/gazelle/pythonconfig/pythonconfig.go
index 8bf79cb..b76e1f9 100644
--- a/gazelle/pythonconfig/pythonconfig.go
+++ b/gazelle/pythonconfig/pythonconfig.go
@@ -98,6 +98,9 @@
 	// separate pyi_deps attribute or merge type-checking dependencies into deps.
 	// Defaults to false for backward compatibility.
 	GeneratePyiDeps = "python_generate_pyi_deps"
+	// GenerateProto represents the directive that controls whether to generate
+	// python_generate_proto targets.
+	GenerateProto = "python_generate_proto"
 )
 
 // GenerationModeType represents one of the generation modes for the Python
@@ -186,6 +189,7 @@
 	labelNormalization                        LabelNormalizationType
 	experimentalAllowRelativeImports          bool
 	generatePyiDeps                           bool
+	generateProto                             bool
 }
 
 type LabelNormalizationType int
@@ -223,6 +227,7 @@
 		labelNormalization:                        DefaultLabelNormalizationType,
 		experimentalAllowRelativeImports:          false,
 		generatePyiDeps:                           false,
+		generateProto:                             false,
 	}
 }
 
@@ -257,6 +262,7 @@
 		labelNormalization:                        c.labelNormalization,
 		experimentalAllowRelativeImports:          c.experimentalAllowRelativeImports,
 		generatePyiDeps:                           c.generatePyiDeps,
+		generateProto:                             c.generateProto,
 	}
 }
 
@@ -555,6 +561,16 @@
 	return c.generatePyiDeps
 }
 
+// SetGenerateProto sets whether py_proto_library should be generated for proto_library.
+func (c *Config) SetGenerateProto(generateProto bool) {
+	c.generateProto = generateProto
+}
+
+// GenerateProto returns whether py_proto_library should be generated for proto_library.
+func (c *Config) GenerateProto() bool {
+	return c.generateProto
+}
+
 // FormatThirdPartyDependency returns a label to a third-party dependency performing all formating and normalization.
 func (c *Config) FormatThirdPartyDependency(repositoryName string, distributionName string) label.Label {
 	conventionalDistributionName := strings.ReplaceAll(c.labelConvention, distributionNameLabelConventionSubstitution, distributionName)