feat(gazelle): Add python_generate_pyi_srcs directive (#3356)

Fixes #3354

This adds a new Gazelle directive `python_generate_pyi_srcs`
which accepts a boolean value `true` or `false`. It defaults
to `false` for backwards compatibility.

When `true`, the directive causes `.pyi` files, whose name matches
with a .py file found in `srcs`, to be added to the `pyi_srcs`
target attribute.

This helps with cases where manually-generated `.pyi` files are
present in the project and are needed for things such as rules_mypy.

Given the following files in a `my_dir` package:

```
BUILD.bazel
__init__.py
a.py
a.pyi
b.py
c.py
c.pyi
```

`# gazelle:python_generate_pyi_srcs true` will cause Gazelle to
generate:

```starlark
py_library(
    name = "my_dir",
    srcs = [
        "__init__.py",
        "a.py",
        "b.py",
        "c.py",
    ],
    pyi_srcs = [
        "a.py",
        "c.py",
    ],
)
```

Author: dougthor42

CHANGELOG.md

@@ -78,6 +78,10 @@ END_UNRELEASED_TEMPLATE
 * (binaries/tests) Build information is now included in binaries and tests.
   Use the `bazel_binary_info` module to access it. The {flag}`--stamp` flag will
   add {flag}`--workspace_status` information.
+* (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-8-1}
 ## [1.8.1] - 2026-01-20

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

@@ -273,6 +273,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")
@@ -282,9 +283,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().
@@ -312,6 +319,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
@@ -327,6 +337,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().
@@ -377,10 +388,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).
@@ -411,8 +427,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).
@@ -443,8 +464,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).
@@ -691,3 +717,25 @@ 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 {
+		return pyiSrcs, nil
+	}
+
+	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

@@ -51,6 +51,7 @@ var pyKinds = map[string]rule.KindInfo{
 		ResolveAttrs: map[string]bool{
 			"deps":     true,
 			"pyi_deps": true,
+			"pyi_srcs": true,
 		},
 	},
 	pyLibraryKind: {
@@ -68,6 +69,7 @@ var pyKinds = map[string]rule.KindInfo{
 		ResolveAttrs: map[string]bool{
 			"deps":     true,
 			"pyi_deps": true,
+			"pyi_srcs": true,
 		},
 	},
 	pyProtoLibraryKind: {
@@ -91,6 +93,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.