Skip to content

Commit e5ef69b

Browse files
amartanidougthor42
amartani
and
dougthor42
authored
feat(gazelle): Add type-checking only dependencies to pyi_deps (bazel-contrib#3014)
bazel-contrib#2538 added the attribute `pyi_deps` to python rules, intended to be used for dependencies that are only used for type-checking purposes. This PR adds a new directive, `gazelle:python_generate_pyi_deps`, which, when enabled: - When a dependency is added only to satisfy type-checking only imports (in a `if TYPE_CHECKING:` block), the dependency is added to `pyi_deps` instead of `deps`; - Third-party stub packages (eg. `boto3-stubs`) are now added to `pyi_deps` instead of `deps`. --------- Co-authored-by: Douglas Thor <dougthor42@users.noreply.github.com>
1 parent 4ec1e80 commit e5ef69b

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+667
-20
lines changed

‎CHANGELOG.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ END_UNRELEASED_TEMPLATE
7070
* (pypi) To configure the environment for `requirements.txt` evaluation, use the newly added
7171
developer preview of the `pip.default` tag class. Only `rules_python` and root modules can use
7272
this feature. You can also configure custom `config_settings` using `pip.default`.
73+
* (gazelle) New directive `gazelle:python_generate_pyi_deps`; when `true`,
74+
dependencies added to satisfy type-only imports (`if TYPE_CHECKING`) and type
75+
stub packages are added to `pyi_deps` instead of `deps`.
7376

7477
{#v0-0-0-removed}
7578
### Removed

‎gazelle/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,8 @@ Python-specific directives are as follows:
222222
| 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". |
223223
| `# gazelle:experimental_allow_relative_imports` | `false` |
224224
| Controls whether Gazelle resolves dependencies for import statements that use paths relative to the current package. Can be "true" or "false".|
225+
| `# gazelle:python_generate_pyi_deps` | `false` |
226+
| 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. |
225227

226228
#### Directive: `python_root`:
227229

‎gazelle/python/configure.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ func (py *Configurer) KnownDirectives() []string {
6868
pythonconfig.TestFilePattern,
6969
pythonconfig.LabelConvention,
7070
pythonconfig.LabelNormalization,
71+
pythonconfig.GeneratePyiDeps,
7172
pythonconfig.ExperimentalAllowRelativeImports,
7273
}
7374
}
@@ -230,6 +231,12 @@ func (py *Configurer) Configure(c *config.Config, rel string, f *rule.File) {
230231
pythonconfig.ExperimentalAllowRelativeImports, rel, d.Value)
231232
}
232233
config.SetExperimentalAllowRelativeImports(v)
234+
case pythonconfig.GeneratePyiDeps:
235+
v, err := strconv.ParseBool(strings.TrimSpace(d.Value))
236+
if err != nil {
237+
log.Fatal(err)
238+
}
239+
config.SetGeneratePyiDeps(v)
233240
}
234241
}
235242

‎gazelle/python/file_parser.go

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,10 @@ type ParserOutput struct {
4747
}
4848

4949
type FileParser struct {
50-
code []byte
51-
relFilepath string
52-
output ParserOutput
50+
code []byte
51+
relFilepath string
52+
output ParserOutput
53+
inTypeCheckingBlock bool
5354
}
5455

5556
func NewFileParser() *FileParser {
@@ -158,6 +159,7 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) bool {
158159
continue
159160
}
160161
m.Filepath = p.relFilepath
162+
m.TypeCheckingOnly = p.inTypeCheckingBlock
161163
if strings.HasPrefix(m.Name, ".") {
162164
continue
163165
}
@@ -178,6 +180,7 @@ func (p *FileParser) parseImportStatements(node *sitter.Node) bool {
178180
m.Filepath = p.relFilepath
179181
m.From = from
180182
m.Name = fmt.Sprintf("%s.%s", from, m.Name)
183+
m.TypeCheckingOnly = p.inTypeCheckingBlock
181184
p.output.Modules = append(p.output.Modules, m)
182185
}
183186
} else {
@@ -202,10 +205,43 @@ func (p *FileParser) SetCodeAndFile(code []byte, relPackagePath, filename string
202205
p.output.FileName = filename
203206
}
204207

208+
// isTypeCheckingBlock returns true if the given node is an `if TYPE_CHECKING:` block.
209+
func (p *FileParser) isTypeCheckingBlock(node *sitter.Node) bool {
210+
if node.Type() != sitterNodeTypeIfStatement || node.ChildCount() < 2 {
211+
return false
212+
}
213+
214+
condition := node.Child(1)
215+
216+
// Handle `if TYPE_CHECKING:`
217+
if condition.Type() == sitterNodeTypeIdentifier && condition.Content(p.code) == "TYPE_CHECKING" {
218+
return true
219+
}
220+
221+
// Handle `if typing.TYPE_CHECKING:`
222+
if condition.Type() == "attribute" && condition.ChildCount() >= 3 {
223+
object := condition.Child(0)
224+
attr := condition.Child(2)
225+
if object.Type() == sitterNodeTypeIdentifier && object.Content(p.code) == "typing" &&
226+
attr.Type() == sitterNodeTypeIdentifier && attr.Content(p.code) == "TYPE_CHECKING" {
227+
return true
228+
}
229+
}
230+
231+
return false
232+
}
233+
205234
func (p *FileParser) parse(ctx context.Context, node *sitter.Node) {
206235
if node == nil {
207236
return
208237
}
238+
239+
// Check if this is a TYPE_CHECKING block
240+
wasInTypeCheckingBlock := p.inTypeCheckingBlock
241+
if p.isTypeCheckingBlock(node) {
242+
p.inTypeCheckingBlock = true
243+
}
244+
209245
for i := 0; i < int(node.ChildCount()); i++ {
210246
if err := ctx.Err(); err != nil {
211247
return
@@ -219,6 +255,9 @@ func (p *FileParser) parse(ctx context.Context, node *sitter.Node) {
219255
}
220256
p.parse(ctx, child)
221257
}
258+
259+
// Restore the previous state
260+
p.inTypeCheckingBlock = wasInTypeCheckingBlock
222261
}
223262

224263
func (p *FileParser) Parse(ctx context.Context) (*ParserOutput, error) {

‎gazelle/python/file_parser_test.go

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,3 +254,40 @@ func TestParseFull(t *testing.T) {
254254
FileName: "a.py",
255255
}, *output)
256256
}
257+
258+
func TestTypeCheckingImports(t *testing.T) {
259+
code := `
260+
import sys
261+
from typing import TYPE_CHECKING
262+
263+
if TYPE_CHECKING:
264+
import boto3
265+
from rest_framework import serializers
266+
267+
def example_function():
268+
_ = sys.version_info
269+
`
270+
p := NewFileParser()
271+
p.SetCodeAndFile([]byte(code), "", "test.py")
272+
273+
result, err := p.Parse(context.Background())
274+
if err != nil {
275+
t.Fatalf("Failed to parse: %v", err)
276+
}
277+
278+
// Check that we found the expected modules
279+
expectedModules := map[string]bool{
280+
"sys": false,
281+
"typing.TYPE_CHECKING": false,
282+
"boto3": true,
283+
"rest_framework.serializers": true,
284+
}
285+
286+
for _, mod := range result.Modules {
287+
if expected, exists := expectedModules[mod.Name]; exists {
288+
if mod.TypeCheckingOnly != expected {
289+
t.Errorf("Module %s: expected TypeCheckingOnly=%v, got %v", mod.Name, expected, mod.TypeCheckingOnly)
290+
}
291+
}
292+
}
293+
}

‎gazelle/python/parser.go

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,9 +112,9 @@ func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[strin
112112
continue
113113
}
114114

115-
modules.Add(m)
115+
addModuleToTreeSet(modules, m)
116116
if res.HasMain {
117-
mainModules[res.FileName].Add(m)
117+
addModuleToTreeSet(mainModules[res.FileName], m)
118118
}
119119
}
120120

@@ -158,13 +158,24 @@ type Module struct {
158158
// If this was a from import, e.g. from foo import bar, From indicates the module
159159
// from which it is imported.
160160
From string `json:"from"`
161+
// Whether this import is type-checking only (inside if TYPE_CHECKING block).
162+
TypeCheckingOnly bool `json:"type_checking_only"`
161163
}
162164

163165
// moduleComparator compares modules by name.
164166
func moduleComparator(a, b interface{}) int {
165167
return godsutils.StringComparator(a.(Module).Name, b.(Module).Name)
166168
}
167169

170+
// addModuleToTreeSet adds a module to a treeset.Set, ensuring that a TypeCheckingOnly=false module is
171+
// prefered over a TypeCheckingOnly=true module.
172+
func addModuleToTreeSet(set *treeset.Set, mod Module) {
173+
if mod.TypeCheckingOnly && set.Contains(mod) {
174+
return
175+
}
176+
set.Add(mod)
177+
}
178+
168179
// annotationKind represents Gazelle annotation kinds.
169180
type annotationKind string
170181

‎gazelle/python/resolve.go

Lines changed: 47 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,16 @@ func (py *Resolver) Embeds(r *rule.Rule, from label.Label) []label.Label {
123123
return make([]label.Label, 0)
124124
}
125125

126+
// addDependency adds a dependency to either the regular deps or pyiDeps set based on
127+
// whether the module is type-checking only.
128+
func addDependency(dep string, mod Module, deps, pyiDeps *treeset.Set) {
129+
if mod.TypeCheckingOnly {
130+
pyiDeps.Add(dep)
131+
} else {
132+
deps.Add(dep)
133+
}
134+
}
135+
126136
// Resolve translates imported libraries for a given rule into Bazel
127137
// dependencies. Information about imported libraries is returned for each
128138
// rule generated by language.GenerateRules in
@@ -141,9 +151,11 @@ func (py *Resolver) Resolve(
141151
// join with the main Gazelle binary with other rules. It may conflict with
142152
// other generators that generate py_* targets.
143153
deps := treeset.NewWith(godsutils.StringComparator)
154+
pyiDeps := treeset.NewWith(godsutils.StringComparator)
155+
cfgs := c.Exts[languageName].(pythonconfig.Configs)
156+
cfg := cfgs[from.Pkg]
157+
144158
if modulesRaw != nil {
145-
cfgs := c.Exts[languageName].(pythonconfig.Configs)
146-
cfg := cfgs[from.Pkg]
147159
pythonProjectRoot := cfg.PythonProjectRoot()
148160
modules := modulesRaw.(*treeset.Set)
149161
it := modules.Iterator()
@@ -228,7 +240,7 @@ func (py *Resolver) Resolve(
228240
override.Repo = ""
229241
}
230242
dep := override.Rel(from.Repo, from.Pkg).String()
231-
deps.Add(dep)
243+
addDependency(dep, mod, deps, pyiDeps)
232244
if explainDependency == dep {
233245
log.Printf("Explaining dependency (%s): "+
234246
"in the target %q, the file %q imports %q at line %d, "+
@@ -239,7 +251,7 @@ func (py *Resolver) Resolve(
239251
}
240252
} else {
241253
if dep, distributionName, ok := cfg.FindThirdPartyDependency(moduleName); ok {
242-
deps.Add(dep)
254+
addDependency(dep, mod, deps, pyiDeps)
243255
// Add the type and stub dependencies if they exist.
244256
modules := []string{
245257
fmt.Sprintf("%s_stubs", strings.ToLower(distributionName)),
@@ -249,7 +261,8 @@ func (py *Resolver) Resolve(
249261
}
250262
for _, module := range modules {
251263
if dep, _, ok := cfg.FindThirdPartyDependency(module); ok {
252-
deps.Add(dep)
264+
// Type stub packages always go to pyiDeps
265+
pyiDeps.Add(dep)
253266
}
254267
}
255268
if explainDependency == dep {
@@ -308,7 +321,7 @@ func (py *Resolver) Resolve(
308321
}
309322
matchLabel := filteredMatches[0].Label.Rel(from.Repo, from.Pkg)
310323
dep := matchLabel.String()
311-
deps.Add(dep)
324+
addDependency(dep, mod, deps, pyiDeps)
312325
if explainDependency == dep {
313326
log.Printf("Explaining dependency (%s): "+
314327
"in the target %q, the file %q imports %q at line %d, "+
@@ -333,16 +346,41 @@ func (py *Resolver) Resolve(
333346
os.Exit(1)
334347
}
335348
}
349+
350+
addResolvedDeps(r, deps)
351+
352+
if cfg.GeneratePyiDeps() {
353+
if !deps.Empty() {
354+
r.SetAttr("deps", convertDependencySetToExpr(deps))
355+
}
356+
if !pyiDeps.Empty() {
357+
r.SetAttr("pyi_deps", convertDependencySetToExpr(pyiDeps))
358+
}
359+
} else {
360+
// When generate_pyi_deps is false, merge both deps and pyiDeps into deps
361+
combinedDeps := treeset.NewWith(godsutils.StringComparator)
362+
combinedDeps.Add(deps.Values()...)
363+
combinedDeps.Add(pyiDeps.Values()...)
364+
365+
if !combinedDeps.Empty() {
366+
r.SetAttr("deps", convertDependencySetToExpr(combinedDeps))
367+
}
368+
}
369+
}
370+
371+
// addResolvedDeps adds the pre-resolved dependencies from the rule's private attributes
372+
// to the provided deps set.
373+
func addResolvedDeps(
374+
r *rule.Rule,
375+
deps *treeset.Set,
376+
) {
336377
resolvedDeps := r.PrivateAttr(resolvedDepsKey).(*treeset.Set)
337378
if !resolvedDeps.Empty() {
338379
it := resolvedDeps.Iterator()
339380
for it.Next() {
340381
deps.Add(it.Value())
341382
}
342383
}
343-
if !deps.Empty() {
344-
r.SetAttr("deps", convertDependencySetToExpr(deps))
345-
}
346384
}
347385

348386
// targetListFromResults returns a string with the human-readable list of

‎gazelle/python/target.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,12 @@
1515
package python
1616

1717
import (
18+
"path/filepath"
19+
1820
"github.com/bazelbuild/bazel-gazelle/config"
1921
"github.com/bazelbuild/bazel-gazelle/rule"
2022
"github.com/emirpasic/gods/sets/treeset"
2123
godsutils "github.com/emirpasic/gods/utils"
22-
"path/filepath"
2324
)
2425

2526
// targetBuilder builds targets to be generated by Gazelle.
@@ -79,7 +80,8 @@ func (t *targetBuilder) addModuleDependency(dep Module) *targetBuilder {
7980
// dependency resolution easier
8081
dep.Name = importSpecFromSrc(t.pythonProjectRoot, t.bzlPackage, fileName).Imp
8182
}
82-
t.deps.Add(dep)
83+
84+
addModuleToTreeSet(t.deps, dep)
8385
return t
8486
}
8587

‎gazelle/python/testdata/add_type_stub_packages/BUILD.in

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
# gazelle:python_generate_pyi_deps true

‎gazelle/python/testdata/add_type_stub_packages/BUILD.out

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,18 @@
11
load("@rules_python//python:defs.bzl", "py_binary")
22

3+
# gazelle:python_generate_pyi_deps true
4+
35
py_binary(
46
name = "add_type_stub_packages_bin",
57
srcs = ["__main__.py"],
68
main = "__main__.py",
9+
pyi_deps = [
10+
"@gazelle_python_test//boto3_stubs",
11+
"@gazelle_python_test//django_types",
12+
],
713
visibility = ["//:__subpackages__"],
814
deps = [
915
"@gazelle_python_test//boto3",
10-
"@gazelle_python_test//boto3_stubs",
1116
"@gazelle_python_test//django",
12-
"@gazelle_python_test//django_types",
1317
],
1418
)

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy