gazelle: Add support for generating pyi_deps

Author: amartani

CHANGELOG.md

@@ -55,6 +55,8 @@ END_UNRELEASED_TEMPLATE
 {#v0-0-0-changed}
 ### Changed
 * (gazelle) Types for exposed members of `python.ParserOutput` are now all public.
+* (gazelle) Dependencies added to satisfy type-only imports (`if TYPE_CHECKING`) and type stub packages are
+  now added to `pyi_deps` instead of `deps`.
 
 {#v0-0-0-fixed}
 ### Fixed

gazelle/python/file_parser.go

@@ -47,9 +47,10 @@ type ParserOutput struct {
 }
 
 type FileParser struct {
-	code        []byte
-	relFilepath string
-	output      ParserOutput
+	code                 []byte
+	relFilepath          string
+	output               ParserOutput
+	inTypeCheckingBlock  bool
 }
 
 func NewFileParser() *FileParser {
@@ -158,6 +159,7 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) bool {
 				continue
 			}
 			m.Filepath = p.relFilepath
+			m.TypeCheckingOnly = p.inTypeCheckingBlock
 			if strings.HasPrefix(m.Name, ".") {
 				continue
 			}
@@ -176,6 +178,7 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) bool {
 			m.Filepath = p.relFilepath
 			m.From = from
 			m.Name = fmt.Sprintf("%s.%s", from, m.Name)
+			m.TypeCheckingOnly = p.inTypeCheckingBlock
 			p.output.Modules = append(p.output.Modules, m)
 		}
 	} else {
@@ -200,10 +203,43 @@ func (p *FileParser) SetCodeAndFile(code []byte, relPackagePath, filename string
 	p.output.FileName = filename
 }
 
+// isTypeCheckingBlock returns true if the given node is an `if TYPE_CHECKING:` block.
+func (p *FileParser) isTypeCheckingBlock(node *sitter.Node) bool {
+	if node.Type() != sitterNodeTypeIfStatement || node.ChildCount() < 2 {
+		return false
+	}
+
+	condition := node.Child(1)
+
+	// Handle `if TYPE_CHECKING:`
+	if condition.Type() == sitterNodeTypeIdentifier && condition.Content(p.code) == "TYPE_CHECKING" {
+		return true
+	}
+
+	// Handle `if typing.TYPE_CHECKING:`
+	if condition.Type() == "attribute" && condition.ChildCount() >= 3 {
+		object := condition.Child(0)
+		attr := condition.Child(2)
+		if object.Type() == sitterNodeTypeIdentifier && object.Content(p.code) == "typing" &&
+			attr.Type() == sitterNodeTypeIdentifier && attr.Content(p.code) == "TYPE_CHECKING" {
+			return true
+		}
+	}
+
+	return false
+}
+
 func (p *FileParser) parse(ctx context.Context, node *sitter.Node) {
 	if node == nil {
 		return
 	}
+
+	// Check if this is a TYPE_CHECKING block
+	wasInTypeCheckingBlock := p.inTypeCheckingBlock
+	if p.isTypeCheckingBlock(node) {
+		p.inTypeCheckingBlock = true
+	}
+
 	for i := 0; i < int(node.ChildCount()); i++ {
 		if err := ctx.Err(); err != nil {
 			return
@@ -217,6 +253,9 @@ func (p *FileParser) parse(ctx context.Context, node *sitter.Node) {
 		}
 		p.parse(ctx, child)
 	}
+
+	// Restore the previous state
+	p.inTypeCheckingBlock = wasInTypeCheckingBlock
 }
 
 func (p *FileParser) Parse(ctx context.Context) (*ParserOutput, error) {

gazelle/python/file_parser_test.go

@@ -254,3 +254,40 @@ func TestParseFull(t *testing.T) {
 		FileName: "a.py",
 	}, *output)
 }
+
+func TestTypeCheckingImports(t *testing.T) {
+	code := `
+import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import boto3
+    from rest_framework import serializers
+
+def example_function():
+    _ = sys.version_info
+`
+	p := NewFileParser()
+	p.SetCodeAndFile([]byte(code), "", "test.py")
+
+	result, err := p.Parse(context.Background())
+	if err != nil {
+		t.Fatalf("Failed to parse: %v", err)
+	}
+
+	// Check that we found the expected modules
+	expectedModules := map[string]bool{
+		"sys": false,
+		"typing.TYPE_CHECKING": false,
+		"boto3": true,
+		"rest_framework.serializers": true,
+	}
+
+	for _, mod := range result.Modules {
+		if expected, exists := expectedModules[mod.Name]; exists {
+			if mod.TypeCheckingOnly != expected {
+				t.Errorf("Module %s: expected TypeCheckingOnly=%v, got %v", mod.Name, expected, mod.TypeCheckingOnly)
+			}
+		}
+	}
+}

gazelle/python/parser.go

@@ -158,6 +158,8 @@ type Module struct {
 	// If this was a from import, e.g. from foo import bar, From indicates the module
 	// from which it is imported.
 	From string `json:"from"`
+	// Whether this import is type-checking only (inside if TYPE_CHECKING block).
+	TypeCheckingOnly bool `json:"type_checking_only"`
 }
 
 // moduleComparator compares modules by name.

gazelle/python/resolve.go

@@ -39,6 +39,9 @@ const (
 	// resolvedDepsKey is the attribute key used to pass dependencies that don't
 	// need to be resolved by the dependency resolver in the Resolver step.
 	resolvedDepsKey = "_gazelle_python_resolved_deps"
+	// resolvedPyiDepsKey is the attribute key used to pass type-checking dependencies that don't
+	// need to be resolved by the dependency resolver in the Resolver step.
+	resolvedPyiDepsKey = "_gazelle_python_resolved_pyi_deps"
 )
 
 // Resolver satisfies the resolve.Resolver interface. It resolves dependencies
@@ -123,6 +126,16 @@ func (py *Resolver) Embeds(r *rule.Rule, from label.Label) []label.Label {
 	return make([]label.Label, 0)
 }
 
+// addDependency adds a dependency to either the regular deps or pyiDeps set based on
+// whether the module is type-checking only.
+func addDependency(dep string, mod Module, deps, pyiDeps *treeset.Set) {
+	if mod.TypeCheckingOnly {
+		pyiDeps.Add(dep)
+	} else {
+		deps.Add(dep)
+	}
+}
+
 // Resolve translates imported libraries for a given rule into Bazel
 // dependencies. Information about imported libraries is returned for each
 // rule generated by language.GenerateRules in
@@ -141,6 +154,8 @@ func (py *Resolver) Resolve(
 	// join with the main Gazelle binary with other rules. It may conflict with
 	// other generators that generate py_* targets.
 	deps := treeset.NewWith(godsutils.StringComparator)
+	pyiDeps := treeset.NewWith(godsutils.StringComparator)
+
 	if modulesRaw != nil {
 		cfgs := c.Exts[languageName].(pythonconfig.Configs)
 		cfg := cfgs[from.Pkg]
@@ -179,7 +194,7 @@ func (py *Resolver) Resolve(
 							override.Repo = ""
 						}
 						dep := override.Rel(from.Repo, from.Pkg).String()
-						deps.Add(dep)
+						addDependency(dep, mod, deps, pyiDeps)
 						if explainDependency == dep {
 							log.Printf("Explaining dependency (%s): "+
 								"in the target %q, the file %q imports %q at line %d, "+
@@ -190,7 +205,7 @@ func (py *Resolver) Resolve(
 					}
 				} else {
 					if dep, distributionName, ok := cfg.FindThirdPartyDependency(moduleName); ok {
-						deps.Add(dep)
+						addDependency(dep, mod, deps, pyiDeps)
 						// Add the type and stub dependencies if they exist.
 						modules := []string{
 							fmt.Sprintf("%s_stubs", strings.ToLower(distributionName)),
@@ -200,7 +215,8 @@ func (py *Resolver) Resolve(
 						}
 						for _, module := range modules {
 							if dep, _, ok := cfg.FindThirdPartyDependency(module); ok {
-								deps.Add(dep)
+								// Type stub packages always go to pyiDeps
+								pyiDeps.Add(dep)
 							}
 						}
 						if explainDependency == dep {
@@ -259,7 +275,7 @@ func (py *Resolver) Resolve(
 						}
 						matchLabel := filteredMatches[0].Label.Rel(from.Repo, from.Pkg)
 						dep := matchLabel.String()
-						deps.Add(dep)
+						addDependency(dep, mod, deps, pyiDeps)
 						if explainDependency == dep {
 							log.Printf("Explaining dependency (%s): "+
 								"in the target %q, the file %q imports %q at line %d, "+
@@ -284,15 +300,29 @@ func (py *Resolver) Resolve(
 			os.Exit(1)
 		}
 	}
-	resolvedDeps := r.PrivateAttr(resolvedDepsKey).(*treeset.Set)
+
+	addResolvedDepsAndSetAttr(r, deps, resolvedDepsKey, "deps")
+	addResolvedDepsAndSetAttr(r, pyiDeps, resolvedPyiDepsKey, "pyi_deps")
+}
+
+// addResolvedDepsAndSetAttr adds the pre-resolved dependencies from the rule's private attributes
+// to the provided deps set and sets the attribute on the rule.
+func addResolvedDepsAndSetAttr(
+	r *rule.Rule,
+	deps *treeset.Set,
+	resolvedDepsAttrName string,
+	depsAttrName string,
+) {
+	resolvedDeps := r.PrivateAttr(resolvedDepsAttrName).(*treeset.Set)
 	if !resolvedDeps.Empty() {
 		it := resolvedDeps.Iterator()
 		for it.Next() {
 			deps.Add(it.Value())
 		}
 	}
+
 	if !deps.Empty() {
-		r.SetAttr("deps", convertDependencySetToExpr(deps))
+		r.SetAttr(depsAttrName, convertDependencySetToExpr(deps))
 	}
 }
 

gazelle/python/target.go

@@ -15,11 +15,12 @@
 package python
 
 import (
+	"path/filepath"
+
 	"github.com/bazelbuild/bazel-gazelle/config"
 	"github.com/bazelbuild/bazel-gazelle/rule"
 	"github.com/emirpasic/gods/sets/treeset"
 	godsutils "github.com/emirpasic/gods/utils"
-	"path/filepath"
 )
 
 // targetBuilder builds targets to be generated by Gazelle.
@@ -31,7 +32,9 @@ type targetBuilder struct {
 	srcs              *treeset.Set
 	siblingSrcs       *treeset.Set
 	deps              *treeset.Set
+	pyiDeps           *treeset.Set
 	resolvedDeps      *treeset.Set
+	resolvedPyiDeps   *treeset.Set
 	visibility        *treeset.Set
 	main              *string
 	imports           []string
@@ -48,7 +51,9 @@ func newTargetBuilder(kind, name, pythonProjectRoot, bzlPackage string, siblingS
 		srcs:              treeset.NewWith(godsutils.StringComparator),
 		siblingSrcs:       siblingSrcs,
 		deps:              treeset.NewWith(moduleComparator),
+		pyiDeps:           treeset.NewWith(moduleComparator),
 		resolvedDeps:      treeset.NewWith(godsutils.StringComparator),
+		resolvedPyiDeps:   treeset.NewWith(godsutils.StringComparator),
 		visibility:        treeset.NewWith(godsutils.StringComparator),
 	}
 }
@@ -79,7 +84,13 @@ func (t *targetBuilder) addModuleDependency(dep Module) *targetBuilder {
 		// dependency resolution easier
 		dep.Name = importSpecFromSrc(t.pythonProjectRoot, t.bzlPackage, fileName).Imp
 	}
-	t.deps.Add(dep)
+
+	// Add to appropriate dependency set based on whether it's type-checking only
+	if dep.TypeCheckingOnly {
+		t.pyiDeps.Add(dep)
+	} else {
+		t.deps.Add(dep)
+	}
 	return t
 }
 
@@ -162,12 +173,23 @@ func (t *targetBuilder) build() *rule.Rule {
 	if t.imports != nil {
 		r.SetAttr("imports", t.imports)
 	}
-	if !t.deps.Empty() {
-		r.SetPrivateAttr(config.GazelleImportsKey, t.deps)
+	if combinedDeps := t.combinedDeps(); !combinedDeps.Empty() {
+		r.SetPrivateAttr(config.GazelleImportsKey, combinedDeps)
 	}
 	if t.testonly {
 		r.SetAttr("testonly", true)
 	}
 	r.SetPrivateAttr(resolvedDepsKey, t.resolvedDeps)
+	r.SetPrivateAttr(resolvedPyiDepsKey, t.resolvedPyiDeps)
 	return r
 }
+
+// Combine both regular and type-checking imports into a single set
+// for passing to the resolver. The resolver will distinguish them
+// based on the TypeCheckingOnly field.
+func (t *targetBuilder) combinedDeps() *treeset.Set {
+	combinedDeps := treeset.NewWith(moduleComparator)
+	combinedDeps.Add(t.pyiDeps.Values()...)
+	combinedDeps.Add(t.deps.Values()...)
+	return combinedDeps
+}

gazelle/python/testdata/add_type_stub_packages/BUILD.out

@@ -4,11 +4,13 @@ py_binary(
     name = "add_type_stub_packages_bin",
     srcs = ["__main__.py"],
     main = "__main__.py",
+    pyi_deps = [
+        "@gazelle_python_test//boto3_stubs",
+        "@gazelle_python_test//django_types",
+    ],
     visibility = ["//:__subpackages__"],
     deps = [
         "@gazelle_python_test//boto3",
-        "@gazelle_python_test//boto3_stubs",
         "@gazelle_python_test//django",
-        "@gazelle_python_test//django_types",
     ],
 )

gazelle/python/testdata/type_checking_imports/BUILD.in

@@ -0,0 +1 @@
+# gazelle:python_generation_mode file

gazelle/python/testdata/type_checking_imports/BUILD.out

@@ -0,0 +1,25 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+# gazelle:python_generation_mode file
+
+py_library(
+    name = "bar",
+    srcs = ["bar.py"],
+    pyi_deps = [":foo"],
+    visibility = ["//:__subpackages__"],
+    deps = [":baz"],
+)
+
+py_library(
+    name = "baz",
+    srcs = ["baz.py"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_library(
+    name = "foo",
+    srcs = ["foo.py"],
+    pyi_deps = ["@gazelle_python_test//djangorestframework"],
+    visibility = ["//:__subpackages__"],
+    deps = ["@gazelle_python_test//boto3"],
+)

gazelle/python/testdata/type_checking_imports/README.md

@@ -0,0 +1,5 @@
+# Type Checking Imports
+
+Test that the Python gazelle correctly handles type-only imports inside `if TYPE_CHECKING:` blocks.
+
+Type-only imports should be added to the `pyi_deps` attribute instead of the regular `deps` attribute.