feat(gazelle): Add type-checking only dependencies to pyi_deps

CHANGELOG.md

@@ -55,8 +55,11 @@ 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`.
+
+### Added
+* (gazelle) New directive `gazelle:python_generate_pyi_deps`; when `true`,
+  dependencies added to satisfy type-only imports (`if TYPE_CHECKING`) and type
+  stub packages are added to `pyi_deps` instead of `deps`.
 
 {#v0-0-0-fixed}
 ### Fixed

gazelle/README.md

@@ -220,6 +220,8 @@ Python-specific directives are as follows:
 | Defines the format of the distribution name in labels to third-party deps. Useful for using Gazelle plugin with other rules with different repository conventions (e.g. `rules_pycross`). Full label is always prepended with (pip) repository name, e.g. `@pip//numpy`.                        |
 | `# gazelle:python_label_normalization`                                                                                                                                                                                                                                                          | `snake_case` |
 | Controls how distribution names in labels to third-party deps are normalized. Useful for using Gazelle plugin with other rules with different label conventions (e.g. `rules_pycross` uses PEP-503). Can be "snake_case", "none", or "pep503".                                                  |
+| `# 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. |
 
 #### Directive: `python_root`:
 

gazelle/python/configure.go

@@ -68,6 +68,7 @@ func (py *Configurer) KnownDirectives() []string {
 		pythonconfig.TestFilePattern,
 		pythonconfig.LabelConvention,
 		pythonconfig.LabelNormalization,
+		pythonconfig.GeneratePyiDeps,
 	}
 }
 
@@ -222,6 +223,12 @@ func (py *Configurer) Configure(c *config.Config, rel string, f *rule.File) {
 			default:
 				config.SetLabelNormalization(pythonconfig.DefaultLabelNormalizationType)
 			}
+		case pythonconfig.GeneratePyiDeps:
+			v, err := strconv.ParseBool(strings.TrimSpace(d.Value))
+			if err != nil {
+				log.Fatal(err)
+			}
+			config.SetGeneratePyiDeps(v)
 		}
 	}
 

gazelle/python/resolve.go

@@ -155,10 +155,10 @@ func (py *Resolver) Resolve(
 	// other generators that generate py_* targets.
 	deps := treeset.NewWith(godsutils.StringComparator)
 	pyiDeps := treeset.NewWith(godsutils.StringComparator)
+	cfgs := c.Exts[languageName].(pythonconfig.Configs)
+	cfg := cfgs[from.Pkg]
 
 	if modulesRaw != nil {
-		cfgs := c.Exts[languageName].(pythonconfig.Configs)
-		cfg := cfgs[from.Pkg]
 		pythonProjectRoot := cfg.PythonProjectRoot()
 		modules := modulesRaw.(*treeset.Set)
 		it := modules.Iterator()
@@ -301,29 +301,40 @@ func (py *Resolver) Resolve(
 		}
 	}
 
-	addResolvedDepsAndSetAttr(r, deps, resolvedDepsKey, "deps")
-	addResolvedDepsAndSetAttr(r, pyiDeps, resolvedPyiDepsKey, "pyi_deps")
+	addResolvedDeps(r, deps)
+
+	if cfg.GeneratePyiDeps() {
+		if !deps.Empty() {
+			r.SetAttr("deps", convertDependencySetToExpr(deps))
+		}
+		if !pyiDeps.Empty() {
+			r.SetAttr("pyi_deps", convertDependencySetToExpr(pyiDeps))
+		}
+	} else {
+		// When generate_pyi_deps is false, merge both deps and pyiDeps into deps
+		combinedDeps := treeset.NewWith(godsutils.StringComparator)
+		combinedDeps.Add(deps.Values()...)
+		combinedDeps.Add(pyiDeps.Values()...)
+
+		if !combinedDeps.Empty() {
+			r.SetAttr("deps", convertDependencySetToExpr(combinedDeps))
+		}
+	}
 }
 
-// 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(
+// addResolvedDeps adds the pre-resolved dependencies from the rule's private attributes
+// to the provided deps set.
+func addResolvedDeps(
 	r *rule.Rule,
 	deps *treeset.Set,
-	resolvedDepsAttrName string,
-	depsAttrName string,
 ) {
-	resolvedDeps := r.PrivateAttr(resolvedDepsAttrName).(*treeset.Set)
+	resolvedDeps := r.PrivateAttr(resolvedDepsKey).(*treeset.Set)
 	if !resolvedDeps.Empty() {
 		it := resolvedDeps.Iterator()
 		for it.Next() {
 			deps.Add(it.Value())
 		}
 	}
-
-	if !deps.Empty() {
-		r.SetAttr(depsAttrName, convertDependencySetToExpr(deps))
-	}
 }
 
 // targetListFromResults returns a string with the human-readable list of

gazelle/python/target.go

@@ -15,12 +15,11 @@
 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.
@@ -32,9 +31,7 @@ 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
@@ -51,9 +48,7 @@ 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),
 	}
 }
@@ -84,13 +79,7 @@ func (t *targetBuilder) addModuleDependency(dep Module) *targetBuilder {
 		// dependency resolution easier
 		dep.Name = importSpecFromSrc(t.pythonProjectRoot, t.bzlPackage, fileName).Imp
 	}
-
-	// 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)
-	}
+	t.deps.Add(dep)
 	return t
 }
 
@@ -173,23 +162,12 @@ func (t *targetBuilder) build() *rule.Rule {
 	if t.imports != nil {
 		r.SetAttr("imports", t.imports)
 	}
-	if combinedDeps := t.combinedDeps(); !combinedDeps.Empty() {
-		r.SetPrivateAttr(config.GazelleImportsKey, combinedDeps)
+	if !t.deps.Empty() {
+		r.SetPrivateAttr(config.GazelleImportsKey, t.deps)
 	}
 	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.in

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

gazelle/python/testdata/add_type_stub_packages/BUILD.out

@@ -1,5 +1,7 @@
 load("@rules_python//python:defs.bzl", "py_binary")
 
+# gazelle:python_generate_pyi_deps true
+
 py_binary(
     name = "add_type_stub_packages_bin",
     srcs = ["__main__.py"],

gazelle/python/testdata/type_checking_imports/BUILD.in

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

gazelle/python/testdata/type_checking_imports/BUILD.out

@@ -1,6 +1,7 @@
 load("@rules_python//python:defs.bzl", "py_library")
 
 # gazelle:python_generation_mode file
+# gazelle:python_generate_pyi_deps true
 
 py_library(
     name = "bar",

gazelle/python/testdata/type_checking_imports_disabled/BUILD.in

@@ -0,0 +1,2 @@
+# gazelle:python_generation_mode file
+# gazelle:python_generate_pyi_deps false

gazelle/python/testdata/type_checking_imports_disabled/BUILD.out

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

gazelle/python/testdata/type_checking_imports_disabled/README.md

@@ -0,0 +1,3 @@
+# Type Checking Imports (disabled)
+
+See `type_checking_imports`; this is the same test case, but with the directive disabled.

gazelle/python/testdata/type_checking_imports_disabled/WORKSPACE

@@ -0,0 +1 @@
+workspace(name = "gazelle_python_test")

gazelle/python/testdata/type_checking_imports_disabled/bar.py

@@ -0,0 +1,9 @@
+from typing import TYPE_CHECKING
+
+# foo should be added as a pyi_deps, since it is only imported in a type-checking context, but baz should be
+# added as a deps.
+from baz import X
+
+if TYPE_CHECKING:
+    import baz
+    import foo

gazelle/python/testdata/type_checking_imports_disabled/baz.py

@@ -0,0 +1,15 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+X = 1

gazelle/python/testdata/type_checking_imports_disabled/foo.py

@@ -0,0 +1,20 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+import boto3
+
+if typing.TYPE_CHECKING:
+    from rest_framework import serializers

gazelle/python/testdata/type_checking_imports_disabled/gazelle_python.yaml

@@ -0,0 +1,19 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+manifest:
+  modules_mapping:
+    boto3: boto3
+    rest_framework: djangorestframework
+  pip_deps_repository_name: gazelle_python_test

gazelle/python/testdata/type_checking_imports_disabled/test.yaml

@@ -0,0 +1,15 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+---

gazelle/pythonconfig/pythonconfig.go

@@ -91,6 +91,10 @@ const (
 	// names of labels to third-party dependencies are normalized. Supported values
 	// are 'none', 'pep503' and 'snake_case' (default). See LabelNormalizationType.
 	LabelNormalization = "python_label_normalization"
+	// GeneratePyiDeps represents the directive that controls whether to generate
+	// separate pyi_deps attribute or merge type-checking dependencies into deps.
+	// Defaults to false for backward compatibility.
+	GeneratePyiDeps = "python_generate_pyi_deps"
 )
 
 // GenerationModeType represents one of the generation modes for the Python
@@ -177,6 +181,7 @@ type Config struct {
 	testFilePattern                           []string
 	labelConvention                           string
 	labelNormalization                        LabelNormalizationType
+	generatePyiDeps                           bool
 }
 
 type LabelNormalizationType int
@@ -212,6 +217,7 @@ func New(
 		testFilePattern:                           strings.Split(DefaultTestFilePatternString, ","),
 		labelConvention:                           DefaultLabelConvention,
 		labelNormalization:                        DefaultLabelNormalizationType,
+		generatePyiDeps:                           false,
 	}
 }
 
@@ -244,6 +250,7 @@ func (c *Config) NewChild() *Config {
 		testFilePattern:                           c.testFilePattern,
 		labelConvention:                           c.labelConvention,
 		labelNormalization:                        c.labelNormalization,
+		generatePyiDeps:                           c.generatePyiDeps,
 	}
 }
 
@@ -520,6 +527,18 @@ func (c *Config) LabelNormalization() LabelNormalizationType {
 	return c.labelNormalization
 }
 
+// SetGeneratePyiDeps sets whether pyi_deps attribute should be generated separately
+// or type-checking dependencies should be merged into the regular deps attribute.
+func (c *Config) SetGeneratePyiDeps(generatePyiDeps bool) {
+	c.generatePyiDeps = generatePyiDeps
+}
+
+// GeneratePyiDeps returns whether pyi_deps attribute should be generated separately
+// or type-checking dependencies should be merged into the regular deps attribute.
+func (c *Config) GeneratePyiDeps() bool {
+	return c.generatePyiDeps
+}
+
 // 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)