feat(gazelle): Add "include_dep" Python comment annotation

gazelle/python/generate.go

@@ -233,7 +233,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 	collisionErrors := singlylinkedlist.New()
 
 	appendPyLibrary := func(srcs *treeset.Set, pyLibraryTargetName string) {
-		allDeps, mainModules, err := parser.parse(srcs)
+		allDeps, mainModules, annotations, err := parser.parse(srcs)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}
@@ -263,6 +263,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 					addVisibility(visibility).
 					addSrc(filename).
 					addModuleDependencies(mainModules[filename]).
+					addResolvedDependencies(annotations.includeDep).
 					generateImportsAttribute().build()
 				result.Gen = append(result.Gen, pyBinary)
 				result.Imports = append(result.Imports, pyBinary.PrivateAttr(config.GazelleImportsKey))
@@ -290,6 +291,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 			addVisibility(visibility).
 			addSrcs(srcs).
 			addModuleDependencies(allDeps).
+			addResolvedDependencies(annotations.includeDep).
 			generateImportsAttribute().
 			build()
 
@@ -314,7 +316,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 	}
 
 	if hasPyBinaryEntryPointFile {
-		deps, _, err := parser.parseSingle(pyBinaryEntrypointFilename)
+		deps, _, annotations, err := parser.parseSingle(pyBinaryEntrypointFilename)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}
@@ -338,6 +340,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 			addVisibility(visibility).
 			addSrc(pyBinaryEntrypointFilename).
 			addModuleDependencies(deps).
+			addResolvedDependencies(annotations.includeDep).
 			generateImportsAttribute()
 
 		pyBinary := pyBinaryTarget.build()
@@ -348,7 +351,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 
 	var conftest *rule.Rule
 	if hasConftestFile {
-		deps, _, err := parser.parseSingle(conftestFilename)
+		deps, _, annotations, err := parser.parseSingle(conftestFilename)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}
@@ -367,6 +370,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 		conftestTarget := newTargetBuilder(pyLibraryKind, conftestTargetname, pythonProjectRoot, args.Rel, pyFileNames).
 			addSrc(conftestFilename).
 			addModuleDependencies(deps).
+			addResolvedDependencies(annotations.includeDep).
 			addVisibility(visibility).
 			setTestonly().
 			generateImportsAttribute()
@@ -379,7 +383,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 
 	var pyTestTargets []*targetBuilder
 	newPyTestTargetBuilder := func(srcs *treeset.Set, pyTestTargetName string) *targetBuilder {
-		deps, _, err := parser.parse(srcs)
+		deps, _, annotations, err := parser.parse(srcs)
 		if err != nil {
 			log.Fatalf("ERROR: %v\n", err)
 		}
@@ -397,6 +401,7 @@ func (py *Python) GenerateRules(args language.GenerateArgs) language.GenerateRes
 		return newTargetBuilder(pyTestKind, pyTestTargetName, pythonProjectRoot, args.Rel, pyFileNames).
 			addSrcs(srcs).
 			addModuleDependencies(deps).
+			addResolvedDependencies(annotations.includeDep).
 			generateImportsAttribute()
 	}
 	if (hasPyTestEntryPointFile || hasPyTestEntryPointTarget || cfg.CoarseGrainedGeneration()) && !cfg.PerFileGeneration() {

gazelle/python/parser.go

@@ -101,15 +101,15 @@ func newPython3Parser(
 
 // parseSingle parses a single Python file and returns the extracted modules
 // from the import statements as well as the parsed comments.
-func (p *python3Parser) parseSingle(pyFilename string) (*treeset.Set, map[string]*treeset.Set, error) {
+func (p *python3Parser) parseSingle(pyFilename string) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {
 	pyFilenames := treeset.NewWith(godsutils.StringComparator)
 	pyFilenames.Add(pyFilename)
 	return p.parse(pyFilenames)
 }
 
 // parse parses multiple Python files and returns the extracted modules from
 // the import statements as well as the parsed comments.
-func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[string]*treeset.Set, error) {
+func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {
 	parserMutex.Lock()
 	defer parserMutex.Unlock()
 
@@ -122,28 +122,29 @@ func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[strin
 	}
 	encoder := json.NewEncoder(parserStdin)
 	if err := encoder.Encode(&req); err != nil {
-		return nil, nil, fmt.Errorf("failed to parse: %w", err)
+		return nil, nil, nil, fmt.Errorf("failed to parse: %w", err)
 	}
 
 	reader := bufio.NewReader(parserStdout)
 	data, err := reader.ReadBytes(0)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to parse: %w", err)
+		return nil, nil, nil, fmt.Errorf("failed to parse: %w", err)
 	}
 	data = data[:len(data)-1]
 	var allRes []parserResponse
 	if err := json.Unmarshal(data, &allRes); err != nil {
-		return nil, nil, fmt.Errorf("failed to parse: %w", err)
+		return nil, nil, nil, fmt.Errorf("failed to parse: %w", err)
 	}
 
 	mainModules := make(map[string]*treeset.Set, len(allRes))
+	allAnnotations := new(annotations)
 	for _, res := range allRes {
 		if res.HasMain {
 			mainModules[res.FileName] = treeset.NewWith(moduleComparator)
 		}
 		annotations, err := annotationsFromComments(res.Comments)
 		if err != nil {
-			return nil, nil, fmt.Errorf("failed to parse annotations: %w", err)
+			return nil, nil, nil, fmt.Errorf("failed to parse annotations: %w", err)
 		}
 
 		for _, m := range res.Modules {
@@ -164,9 +165,32 @@ func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[strin
 				mainModules[res.FileName].Add(m)
 			}
 		}
+
+		// Collect all annotations from each file into a single annotations struct.
+		for k, v := range annotations.ignore {
+			allAnnotations.ignore[k] = v
+		}
+		allAnnotations.includeDep = append(allAnnotations.includeDep, annotations.includeDep...)
 	}
 
-	return modules, mainModules, nil
+	allAnnotations.includeDep = removeDupesFromStringTreeSetSlice(allAnnotations.includeDep)
+
+	return modules, mainModules, allAnnotations, nil
+}
+
+// removeDupesFromStringTreeSetSlice takes a []string, makes a set out of the
+// elements, and then returns a new []string with all duplicates removed. Order
+// is preserved.
+func removeDupesFromStringTreeSetSlice(array []string) []string {
+	s := treeset.NewWith(godsutils.StringComparator)
+	for _, v := range array {
+		s.Add(v)
+	}
+	dedupe := make([]string, s.Size())
+	for i, v := range s.Values() {
+		dedupe[i] = fmt.Sprint(v)
+	}
+	return dedupe
 }
 
 // parserResponse represents a response returned by the parser.py for a given
@@ -211,7 +235,8 @@ const (
 	// The Gazelle annotation prefix.
 	annotationPrefix string = "gazelle:"
 	// The ignore annotation kind. E.g. '# gazelle:ignore <module_name>'.
-	annotationKindIgnore annotationKind = "ignore"
+	annotationKindIgnore     annotationKind = "ignore"
+	annotationKindIncludeDep annotationKind = "include_dep"
 )
 
 // comment represents a Python comment.
@@ -247,12 +272,15 @@ type annotation struct {
 type annotations struct {
 	// The parsed modules to be ignored by Gazelle.
 	ignore map[string]struct{}
+	// Labels that Gazelle should include as deps of the generated target.
+	includeDep []string
 }
 
 // annotationsFromComments returns all the annotations parsed out of the
 // comments of a Python module.
 func annotationsFromComments(comments []comment) (*annotations, error) {
 	ignore := make(map[string]struct{})
+	includeDep := []string{}
 	for _, comment := range comments {
 		annotation, err := comment.asAnnotation()
 		if err != nil {
@@ -269,10 +297,21 @@ func annotationsFromComments(comments []comment) (*annotations, error) {
 					ignore[m] = struct{}{}
 				}
 			}
+			if annotation.kind == annotationKindIncludeDep {
+				targets := strings.Split(annotation.value, ",")
+				for _, t := range targets {
+					if t == "" {
+						continue
+					}
+					t = strings.TrimSpace(t)
+					includeDep = append(includeDep, t)
+				}
+			}
 		}
 	}
 	return &annotations{
-		ignore: ignore,
+		ignore:     ignore,
+		includeDep: includeDep,
 	}, nil
 }
 

gazelle/python/target.go

@@ -99,6 +99,15 @@ func (t *targetBuilder) addResolvedDependency(dep string) *targetBuilder {
 	return t
 }
 
+// addResolvedDependencies adds multiple dependencies, that have already been
+// resolved or generated, to the target.
+func (t *targetBuilder) addResolvedDependencies(deps []string) *targetBuilder {
+	for _, dep := range deps {
+		t.addResolvedDependency(dep)
+	}
+	return t
+}
+
 // addVisibility adds visibility labels to the target.
 func (t *targetBuilder) addVisibility(visibility []string) *targetBuilder {
 	for _, item := range visibility {

gazelle/python/testdata/annotation_include_dep/BUILD.in

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

gazelle/python/testdata/annotation_include_dep/BUILD.out

@@ -0,0 +1,53 @@
+load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
+
+# gazelle:python_generation_mode file
+
+py_library(
+    name = "__init__",
+    srcs = ["__init__.py"],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        ":module1",
+        ":module2",
+        "//foo/bar:baz",
+        "//hello:world",
+        "@gazelle_python_test//foo",
+        "@star_wars//rebel_alliance/luke:skywalker",
+    ],
+)
+
+py_library(
+    name = "module1",
+    srcs = ["module1.py"],
+    visibility = ["//:__subpackages__"],
+)
+
+py_library(
+    name = "module2",
+    srcs = ["module2.py"],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        "//checking/py_binary/from/if:works",
+        "//foo:bar",
+    ],
+)
+
+py_binary(
+    name = "annotation_include_dep_bin",
+    srcs = ["__main__.py"],
+    main = "__main__.py",
+    visibility = ["//:__subpackages__"],
+    deps = [
+        ":module2",
+        "//checking/py_binary/from/__main__:works",
+    ],
+)
+
+py_test(
+    name = "module2_test",
+    srcs = ["module2_test.py"],
+    deps = [
+        ":module2",
+        "//checking/py_test/works:too",
+    ],
+)

gazelle/python/testdata/annotation_include_dep/README.md

@@ -0,0 +1,10 @@
+# Annotation: Include Dep
+
+Test that the Python gazelle annotation `# gazelle:include_dep` correctly adds dependences
+to the generated target even if those dependencies are not imported by the Python module.
+
+The root directory tests that all `py_*` targets will correctly include the additional
+dependencies.
+
+The `subpkg` directory tests that all `# gazlle:include_dep` annotations found in all source
+files are included in the generated target (such as during `generation_mode package`).

gazelle/python/testdata/annotation_include_dep/__init__.py

@@ -0,0 +1,9 @@
+import module1
+import foo  # third party package
+
+# gazelle:include_dep //foo/bar:baz
+# gazelle:include_dep //hello:world,@star_wars//rebel_alliance/luke:skywalker
+# gazelle:include_dep :module2
+
+del module1
+del foo

gazelle/python/testdata/annotation_include_dep/__main__.py

@@ -0,0 +1,7 @@
+# gazelle:include_dep //checking/py_binary/from/__main__:works
+# Check deduping
+# gazelle:include_dep //checking/py_binary/from/__main__:works
+
+import module2
+
+del module2

gazelle/python/testdata/annotation_include_dep/gazelle_python.yaml

@@ -0,0 +1,18 @@
+# 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:
+    foo: foo
+  pip_deps_repository_name: gazelle_python_test

gazelle/python/testdata/annotation_include_dep/module2.py

@@ -0,0 +1,5 @@
+# gazelle:include_dep //foo:bar
+
+if __name__ == "__main__":
+    # gazelle:include_dep //checking/py_binary/from/if:works
+    print("hello")

gazelle/python/testdata/annotation_include_dep/module2_test.py

@@ -0,0 +1,5 @@
+# gazelle:include_dep //checking/py_test/works:too
+
+import module2
+
+del module2

gazelle/python/testdata/annotation_include_dep/subpkg/BUILD.in

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

gazelle/python/testdata/annotation_include_dep/subpkg/BUILD.out

@@ -0,0 +1,29 @@
+load("@rules_python//python:defs.bzl", "py_library", "py_test")
+
+# gazelle:python_generation_mode package
+
+py_library(
+    name = "subpkg",
+    srcs = [
+        "__init__.py",
+        "module1.py",
+        "module2.py",
+        "module3.py",
+    ],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        ":nonexistant_target_from_include_dep_in_module3",
+        "//me_from_module1",
+        "//other/thing:from_include_dep_in_module2",
+        "//you_from_module1",
+    ],
+)
+
+py_test(
+    name = "module1_test",
+    srcs = ["module1_test.py"],
+    deps = [
+        ":subpkg",
+        "//:bagel_from_include_dep_in_module1_test",
+    ],
+)

gazelle/python/testdata/annotation_include_dep/subpkg/module1.py

@@ -0,0 +1,3 @@
+def hello():
+    # gazelle:include_dep //you_from_module1,//me_from_module1
+    pass

gazelle/python/testdata/annotation_include_dep/subpkg/module1_test.py

@@ -0,0 +1,5 @@
+# gazelle:include_dep //:bagel_from_include_dep_in_module1_test
+
+import module1
+
+del module1

gazelle/python/testdata/annotation_include_dep/subpkg/module2.py

@@ -0,0 +1,4 @@
+# gazelle:include_dep //other/thing:from_include_dep_in_module2
+import module1
+
+del module1

gazelle/python/testdata/annotation_include_dep/subpkg/module3.py

@@ -0,0 +1,3 @@
+def goodbye():
+    # gazelle:include_dep :nonexistant_target_from_include_dep_in_module3
+    pass