feat(gazelle): Add python_generate_pyi_srcs directive

CHANGELOG.md

@@ -48,6 +48,30 @@ BEGIN_UNRELEASED_TEMPLATE
 END_UNRELEASED_TEMPLATE
 -->
 
+{#v0-0-0}
+## Unreleased
+
+[0.0.0]: https://github.com/bazel-contrib/rules_python/releases/tag/0.0.0
+
+{#v0-0-0-removed}
+### Removed
+
+* Nothing removed.
+{#v0-0-0-changed}
+### Changed
+* Nothing changed.
+
+{#v0-0-0-fixed}
+### Fixed
+* Nothing fixed.
+
+{#v0-0-0-added}
+### Added
+* (gazelle) A new directive `python_generate_pyi_deps` has been added. When
+  `true`, a py_* target's `pyi_srcs` attribute will be set if any `.pyi` files
+  that are associated with the target's `srcs` are present.
+  ([#3354](https://github.com/bazel-contrib/rules_python/issues/3354)).
+
 {#v1-7-0}
 ## [1.7.0] - 2025-10-11
 
@@ -1980,4 +2004,4 @@ Breaking changes:
 * (pip) Create all_data_requirements alias
 * Expose Python C headers through the toolchain.
 
-[0.24.0]: https://github.com/bazel-contrib/rules_python/releases/tag/0.24.0
+[0.24.0]: https://github.com/bazel-contrib/rules_python/releases/tag/0.24.0

gazelle/docs/directives.md

@@ -155,6 +155,12 @@ The Python-specific directives are:
   * Default: `false`
   * Allowed Values: `true`, `false`
 
+[`# gazelle:python_generate_pyi_srcs bool`](#python-generate-pyi-srcs)
+: Controls whether to generate a `pyi_srcs` attribute if a sibling `.pyi` file
+  is found. When `false` (default), the `pyi_srcs` attribute is not added.
+  * Default: `false`
+  * Allowed Values: `true`, `false`
+
 [`# gazelle:python_generate_proto bool`](#python-generate-proto)
 : Controls whether to generate a {bzl:obj}`py_proto_library` for each
   {bzl:obj}`proto_library` in the package. By default we load this rule from the
@@ -626,6 +632,28 @@ Detailed docs are not yet written.
 :::
 
 
+## `python_generate_pyi_deps`
+
+When `true`, include any sibling `.pyi` files in the `pyi_srcs` target attribute.
+
+For example, assume you have the following files:
+
+```
+foo.py
+foo.pyi
+```
+
+The generated target will be:
+
+```starlark
+py_library(
+    name = "foo",
+    srcs = ["foo.py"],
+    pyi_srcs = ["foo.pyi"],
+)
+```
+
+
 ## `python_generate_proto`
 
 When `# gazelle:python_generate_proto true`, Gazelle will generate one

gazelle/python/configure.go

@@ -70,6 +70,7 @@ func (py *Configurer) KnownDirectives() []string {
 		pythonconfig.LabelConvention,
 		pythonconfig.LabelNormalization,
 		pythonconfig.GeneratePyiDeps,
+		pythonconfig.GeneratePyiSrcs,
 		pythonconfig.ExperimentalAllowRelativeImports,
 		pythonconfig.GenerateProto,
 		pythonconfig.PythonResolveSiblingImports,
@@ -242,6 +243,12 @@ func (py *Configurer) Configure(c *config.Config, rel string, f *rule.File) {
 				log.Fatal(err)
 			}
 			config.SetGeneratePyiDeps(v)
+		case pythonconfig.GeneratePyiSrcs:
+			v, err := strconv.ParseBool(strings.TrimSpace(d.Value))
+			if err != nil {
+				log.Fatal(err)
+			}
+			config.SetGeneratePyiSrcs(v)
 		case pythonconfig.GenerateProto:
 			v, err := strconv.ParseBool(strings.TrimSpace(d.Value))
 			if err != nil {

gazelle/python/generate.go

@@ -250,6 +250,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 					srcs.Remove(name)
 				}
 			}
+
 			sort.Strings(mainFileNames)
 			for _, filename := range mainFileNames {
 				pyBinaryTargetName := strings.TrimSuffix(filepath.Base(filename), ".py")
@@ -259,9 +260,15 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 						fqTarget.String(), actualPyBinaryKind, err)
 					continue
 				}
+
+				// Add any sibling .pyi files to pyi_srcs
+				filenames := treeset.NewWith(godsutils.StringComparator, filename)
+				pyiSrcs, _ := getPyiFilenames(filenames, cfg.GeneratePyiSrcs(), args.Dir)
+
 				pyBinary := newTargetBuilder(pyBinaryKind, pyBinaryTargetName, pythonProjectRoot, args.Rel, pyFileNames, cfg.ResolveSiblingImports()).
 					addVisibility(visibility).
 					addSrc(filename).
+					addPyiSrcs(pyiSrcs).
 					addModuleDependencies(mainModules[filename]).
 					addResolvedDependencies(annotations.includeDeps).
 					generateImportsAttribute().
@@ -289,6 +296,9 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 			}
 		}
 
+		// Add any sibling .pyi files to pyi_srcs
+		pyiSrcs, _ := getPyiFilenames(srcs, cfg.GeneratePyiSrcs(), args.Dir)
+
 		// Check if a target with the same name we are generating already
 		// exists, and if it is of a different kind from the one we are
 		// generating. If so, we have to throw an error since Gazelle won't
@@ -304,6 +314,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 		pyLibrary := newTargetBuilder(pyLibraryKind, pyLibraryTargetName, pythonProjectRoot, args.Rel, pyFileNames, cfg.ResolveSiblingImports()).
 			addVisibility(visibility).
 			addSrcs(srcs).
+			addPyiSrcs(pyiSrcs).
 			addModuleDependencies(allDeps).
 			addResolvedDependencies(annotations.includeDeps).
 			generateImportsAttribute().
@@ -354,10 +365,15 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 			collisionErrors.Add(err)
 		}
 
+		// Add any sibling .pyi files to pyi_srcs
+		filenames := treeset.NewWith(godsutils.StringComparator, pyBinaryEntrypointFilename)
+		pyiSrcs, _ := getPyiFilenames(filenames, cfg.GeneratePyiSrcs(), args.Dir)
+
 		pyBinaryTarget := newTargetBuilder(pyBinaryKind, pyBinaryTargetName, pythonProjectRoot, args.Rel, pyFileNames, cfg.ResolveSiblingImports()).
 			setMain(pyBinaryEntrypointFilename).
 			addVisibility(visibility).
 			addSrc(pyBinaryEntrypointFilename).
+			addPyiSrcs(pyiSrcs).
 			addModuleDependencies(deps).
 			addResolvedDependencies(annotations.includeDeps).
 			setAnnotations(*annotations).
@@ -387,8 +403,13 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 			collisionErrors.Add(err)
 		}
 
+		// Add any sibling .pyi files to pyi_srcs
+		filenames := treeset.NewWith(godsutils.StringComparator, conftestFilename)
+		pyiSrcs, _ := getPyiFilenames(filenames, cfg.GeneratePyiSrcs(), args.Dir)
+
 		conftestTarget := newTargetBuilder(pyLibraryKind, conftestTargetname, pythonProjectRoot, args.Rel, pyFileNames, cfg.ResolveSiblingImports()).
 			addSrc(conftestFilename).
+			addPyiSrcs(pyiSrcs).
 			addModuleDependencies(deps).
 			addResolvedDependencies(annotations.includeDeps).
 			setAnnotations(*annotations).
@@ -419,8 +440,13 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 				fqTarget.String(), actualPyTestKind, err, pythonconfig.TestNamingConvention)
 			collisionErrors.Add(err)
 		}
+
+		// Add any sibling .pyi files to pyi_srcs
+		pyiSrcs, _ := getPyiFilenames(srcs, cfg.GeneratePyiSrcs(), args.Dir)
+
 		return newTargetBuilder(pyTestKind, pyTestTargetName, pythonProjectRoot, args.Rel, pyFileNames, cfg.ResolveSiblingImports()).
 			addSrcs(srcs).
+			addPyiSrcs(pyiSrcs).
 			addModuleDependencies(deps).
 			addResolvedDependencies(annotations.includeDeps).
 			setAnnotations(*annotations).
@@ -627,3 +653,23 @@ func generateProtoLibraries(args language.GenerateArgs, cfg *pythonconfig.Config
 	}
 
 }
+
+// getPyiFilenames returns a set of existing .pyi source file names for a given set of source
+// file names if GeneratePyiSrcs is set. Otherwise, returns an empty set.
+func getPyiFilenames(filenames *treeset.Set, generatePyiSrcs bool, basePath string) (*treeset.Set, error) {
+	pyiSrcs := treeset.NewWith(godsutils.StringComparator)
+	if generatePyiSrcs {
+		it := filenames.Iterator()
+		for it.Next() {
+			pyiFilename := it.Value().(string) + "i" // foo.py --> foo.pyi
+
+			_, err := os.Stat(filepath.Join(basePath, pyiFilename))
+			// If the file DNE or there's some other error, there's nothing to do.
+			if err == nil {
+				// pyi file exists, add it
+				pyiSrcs.Add(pyiFilename)
+			}
+		}
+	}
+	return pyiSrcs, nil
+}

gazelle/python/kinds.go

@@ -50,6 +50,7 @@ var pyKinds = map[string]rule.KindInfo{
 		ResolveAttrs: map[string]bool{
 			"deps":     true,
 			"pyi_deps": true,
+			"pyi_srcs": true,
 		},
 	},
 	pyLibraryKind: {
@@ -67,6 +68,7 @@ var pyKinds = map[string]rule.KindInfo{
 		ResolveAttrs: map[string]bool{
 			"deps":     true,
 			"pyi_deps": true,
+			"pyi_srcs": true,
 		},
 	},
 	pyProtoLibraryKind: {
@@ -90,6 +92,7 @@ var pyKinds = map[string]rule.KindInfo{
 		ResolveAttrs: map[string]bool{
 			"deps":     true,
 			"pyi_deps": true,
+			"pyi_srcs": true,
 		},
 	},
 }

gazelle/python/target.go

@@ -30,6 +30,7 @@ type targetBuilder struct {
 	pythonProjectRoot     string
 	bzlPackage            string
 	srcs                  *treeset.Set
+	pyiSrcs               *treeset.Set
 	siblingSrcs           *treeset.Set
 	deps                  *treeset.Set
 	resolvedDeps          *treeset.Set
@@ -49,6 +50,7 @@ func newTargetBuilder(kind, name, pythonProjectRoot, bzlPackage string, siblingS
 		pythonProjectRoot:     pythonProjectRoot,
 		bzlPackage:            bzlPackage,
 		srcs:                  treeset.NewWith(godsutils.StringComparator),
+		pyiSrcs:               treeset.NewWith(godsutils.StringComparator),
 		siblingSrcs:           siblingSrcs,
 		deps:                  treeset.NewWith(moduleComparator),
 		resolvedDeps:          treeset.NewWith(godsutils.StringComparator),
@@ -73,6 +75,21 @@ func (t *targetBuilder) addSrcs(srcs *treeset.Set) *targetBuilder {
 	return t
 }
 
+// addPyiSrc adds a single pyi_src to the target.
+func (t *targetBuilder) addPyiSrc(pyiSrc string) *targetBuilder {
+	t.pyiSrcs.Add(pyiSrc)
+	return t
+}
+
+// addPyiSrcs adds multiple pyi_srcs to the target.
+func (t *targetBuilder) addPyiSrcs(pyiSrcs *treeset.Set) *targetBuilder {
+	it := pyiSrcs.Iterator()
+	for it.Next() {
+		t.pyiSrcs.Add(it.Value().(string))
+	}
+	return t
+}
+
 // addModuleDependency adds a single module dep to the target.
 func (t *targetBuilder) addModuleDependency(dep Module) *targetBuilder {
 	fileName := dep.Name + ".py"
@@ -165,6 +182,9 @@ func (t *targetBuilder) build() *rule.Rule {
 	if !t.srcs.Empty() {
 		r.SetAttr("srcs", t.srcs.Values())
 	}
+	if !t.pyiSrcs.Empty() {
+		r.SetAttr("pyi_srcs", t.pyiSrcs.Values())
+	}
 	if !t.visibility.Empty() {
 		r.SetAttr("visibility", t.visibility.Values())
 	}

gazelle/python/testdata/directive_python_generate_pyi_srcs/BUILD.in

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

gazelle/python/testdata/directive_python_generate_pyi_srcs/BUILD.out

@@ -0,0 +1,36 @@
+load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
+
+# gazelle:python_generate_pyi_srcs true
+
+py_library(
+    name = "directive_python_generate_pyi_srcs",
+    srcs = [
+        "__init__.py",
+        "bar.py",
+        "baz.py",
+        "foo.py",
+    ],
+    pyi_srcs = [
+        "baz.pyi",
+        "foo.pyi",
+    ],
+    visibility = ["//:__subpackages__"],
+)
+
+py_binary(
+    name = "directive_python_generate_pyi_srcs_bin",
+    srcs = ["__main__.py"],
+    main = "__main__.py",
+    pyi_srcs = ["__main__.pyi"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_test(
+    name = "directive_python_generate_pyi_srcs_test",
+    srcs = [
+        "__test__.py",
+        "foo_test.py",
+    ],
+    main = "__test__.py",
+    pyi_srcs = ["foo_test.pyi"],
+)

gazelle/python/testdata/directive_python_generate_pyi_srcs/README.md

@@ -0,0 +1,13 @@
+# Directive: python_generate_pyi_srcs
+
+Test that the `python_generate_pyi_srcs` directive will add `pyi_srcs` to
+generated targets and that it can be toggled on/off on a per-package basis.
+
+The root of the test case asserts that the default generation mode (package)
+will compile multiple .pyi files into a single py_* target.
+
+The `per_file` directory asserts that the `file` generation mode will attach
+a single .pyi file to a given target.
+
+Lastly, the `per_file/turn_off` directory asserts that we can turn off the
+directive for subpackages. It continues with per-file generation mode.

gazelle/python/testdata/directive_python_generate_pyi_srcs/per_file/BUILD.in

@@ -0,0 +1,2 @@
+# gazelle:python_generate_pyi_srcs true
+# gazelle:python_generation_mode file

gazelle/python/testdata/directive_python_generate_pyi_srcs/per_file/BUILD.out

@@ -0,0 +1,30 @@
+load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
+
+# gazelle:python_generate_pyi_srcs true
+# gazelle:python_generation_mode file
+
+py_library(
+    name = "bar",
+    srcs = ["bar.py"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_library(
+    name = "foo",
+    srcs = ["foo.py"],
+    pyi_srcs = ["foo.pyi"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_binary(
+    name = "my_binary",
+    srcs = ["my_binary.py"],
+    pyi_srcs = ["my_binary.pyi"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_test(
+    name = "bar_test",
+    srcs = ["bar_test.py"],
+    pyi_srcs = ["bar_test.pyi"],
+)

gazelle/python/testdata/directive_python_generate_pyi_srcs/per_file/my_binary.py

@@ -0,0 +1,2 @@
+if __name__ == "__main__":
+    print("hey")

gazelle/python/testdata/directive_python_generate_pyi_srcs/per_file/turn_off/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_generate_pyi_srcs false

gazelle/python/testdata/directive_python_generate_pyi_srcs/per_file/turn_off/BUILD.out

@@ -0,0 +1,9 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+# gazelle:python_generate_pyi_srcs false
+
+py_library(
+    name = "foo",
+    srcs = ["foo.py"],
+    visibility = ["//:__subpackages__"],
+)