Add tests for project / package mode; make the exported treesets deterministic.

Author: amartani

gazelle/python/parser.go

@@ -112,9 +112,9 @@ func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[strin
 				continue
 			}
 
-			modules.Add(m)
+			addModuleToTreeSet(modules, m)
 			if res.HasMain {
-				mainModules[res.FileName].Add(m)
+				addModuleToTreeSet(mainModules[res.FileName], m)
 			}
 		}
 
@@ -167,6 +167,15 @@ func moduleComparator(a, b interface{}) int {
 	return godsutils.StringComparator(a.(Module).Name, b.(Module).Name)
 }
 
+// addModuleToTreeSet adds a module to a treeset.Set, ensuring that a TypeCheckingOnly=false module is
+// prefered over a TypeCheckingOnly=true module.
+func addModuleToTreeSet(set *treeset.Set, mod Module) {
+	if mod.TypeCheckingOnly && set.Contains(mod) {
+		return
+	}
+	set.Add(mod)
+}
+
 // annotationKind represents Gazelle annotation kinds.
 type annotationKind string
 

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,6 +32,7 @@ type targetBuilder struct {
 	srcs              *treeset.Set
 	siblingSrcs       *treeset.Set
 	deps              *treeset.Set
+	pyiDeps           *treeset.Set
 	resolvedDeps      *treeset.Set
 	visibility        *treeset.Set
 	main              *string
@@ -48,6 +50,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),
 		visibility:        treeset.NewWith(godsutils.StringComparator),
 	}
@@ -79,7 +82,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 +171,24 @@ 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)
 	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)
+	// If an import is in both pyi_deps and deps, the one in deps will override the one in pyi_deps, resulting
+	// in the resolver properly adding to deps instead of pyi_deps.
+	combinedDeps.Add(t.pyiDeps.Values()...)
+	combinedDeps.Add(t.deps.Values()...)
+	return combinedDeps
+}

gazelle/python/testdata/type_checking_imports_package/BUILD.in

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

gazelle/python/testdata/type_checking_imports_package/BUILD.out

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

gazelle/python/testdata/type_checking_imports_package/README.md

@@ -0,0 +1,3 @@
+# Type Checking Imports (package mode)
+
+See `type_checking_imports`; this is the same test case, but using the package generation mode.

gazelle/python/testdata/type_checking_imports_package/WORKSPACE

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

gazelle/python/testdata/type_checking_imports_package/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_package/baz.py

@@ -0,0 +1,23 @@
+# 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.
+
+
+# While this format is not official, it is supported by most type checkers and
+# is used in the wild to avoid importing the typing module.
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    # Both boto3 and boto3_stubs should be added to pyi_deps.
+    import boto3
+
+X = 1

gazelle/python/testdata/type_checking_imports_package/foo.py

@@ -0,0 +1,21 @@
+# 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
+
+# boto3 should be added to deps. boto3_stubs and djangorestframework should be added to pyi_deps.
+import boto3
+
+if typing.TYPE_CHECKING:
+    from rest_framework import serializers

gazelle/python/testdata/type_checking_imports_package/gazelle_python.yaml

@@ -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.
+
+manifest:
+  modules_mapping:
+    boto3: boto3
+    boto3_stubs: boto3_stubs
+    rest_framework: djangorestframework
+  pip_deps_repository_name: gazelle_python_test