[red-knot] Add some knowledge of `__all__` to `*`-import machinery (#17373)

This commit is contained in:
Alex Waygood 2025-04-15 12:56:40 +01:00 committed by GitHub
parent cf8dc60292
commit 312a487ea7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 189 additions and 44 deletions

View File

@ -626,6 +626,30 @@ reveal_type(X) # revealed: Unknown
reveal_type(Y) # revealed: bool
```
### An implicit import in a `.pyi` file later overridden by another assignment
`a.pyi`:
```pyi
X: bool = True
```
`b.pyi`:
```pyi
from a import X
X: bool = False
```
`c.py`:
```py
from b import *
reveal_type(X) # revealed: bool
```
## Visibility constraints
If an `importer` module contains a `from exporter import *` statement in its global namespace, the
@ -865,15 +889,10 @@ from exporter import *
reveal_type(X) # revealed: bool
# TODO none of these should error, should all reveal `bool`
# error: [unresolved-reference]
reveal_type(_private) # revealed: Unknown
# error: [unresolved-reference]
reveal_type(__protected) # revealed: Unknown
# error: [unresolved-reference]
reveal_type(__dunder__) # revealed: Unknown
# error: [unresolved-reference]
reveal_type(___thunder___) # revealed: Unknown
reveal_type(_private) # revealed: bool
reveal_type(__protected) # revealed: bool
reveal_type(__dunder__) # revealed: bool
reveal_type(___thunder___) # revealed: bool
# TODO: should emit [unresolved-reference] diagnostic & reveal `Unknown`
reveal_type(Y) # revealed: bool
@ -1072,6 +1091,44 @@ reveal_type(Y) # revealed: bool
reveal_type(Z) # revealed: Unknown
```
### `__all__` conditionally defined in a statically known branch (2)
The same example again, but with a different `python-version` set:
```toml
[environment]
python-version = "3.10"
```
`exporter.py`:
```py
import sys
X: bool = True
if sys.version_info >= (3, 11):
__all__ = ["X", "Y"]
Y: bool = True
else:
__all__ = ("Z",)
Z: bool = True
```
`importer.py`:
```py
from exporter import *
# TODO: should reveal `Unknown` and emit `[unresolved-reference]`
reveal_type(X) # revealed: bool
# error: [unresolved-reference]
reveal_type(Y) # revealed: Unknown
reveal_type(Z) # revealed: bool
```
### `__all__` conditionally mutated in a statically known branch
```toml
@ -1084,11 +1141,11 @@ python-version = "3.11"
```py
import sys
__all__ = ["X"]
__all__ = []
X: bool = True
if sys.version_info >= (3, 11):
__all__.append("Y")
__all__.extend(["X", "Y"])
Y: bool = True
else:
__all__.append("Z")
@ -1107,6 +1164,45 @@ reveal_type(Y) # revealed: bool
reveal_type(Z) # revealed: Unknown
```
### `__all__` conditionally mutated in a statically known branch (2)
The same example again, but with a different `python-version` set:
```toml
[environment]
python-version = "3.10"
```
`exporter.py`:
```py
import sys
__all__ = []
X: bool = True
if sys.version_info >= (3, 11):
__all__.extend(["X", "Y"])
Y: bool = True
else:
__all__.append("Z")
Z: bool = True
```
`importer.py`:
```py
from exporter import *
# TODO: should reveal `Unknown` & emit `[unresolved-reference]
reveal_type(X) # revealed: bool
# error: [unresolved-reference]
reveal_type(Y) # revealed: Unknown
reveal_type(Z) # revealed: bool
```
### Empty `__all__`
An empty `__all__` is valid, but a `*` import from a module with an empty `__all__` results in 0
@ -1166,6 +1262,7 @@ from b import *
# TODO: should not error, should reveal `bool`
# (`X` is re-exported from `b.pyi` due to presence in `__all__`)
# See https://github.com/astral-sh/ruff/issues/16159
#
# error: [unresolved-reference]
reveal_type(X) # revealed: Unknown

View File

@ -26,36 +26,37 @@ use ruff_python_ast::{
name::Name,
visitor::{walk_expr, walk_pattern, walk_stmt, Visitor},
};
use rustc_hash::FxHashSet;
use rustc_hash::FxHashMap;
use crate::{module_name::ModuleName, resolve_module, Db};
fn exports_cycle_recover(
_db: &dyn Db,
_value: &FxHashSet<Name>,
_value: &[Name],
_count: u32,
_file: File,
) -> salsa::CycleRecoveryAction<FxHashSet<Name>> {
) -> salsa::CycleRecoveryAction<Box<[Name]>> {
salsa::CycleRecoveryAction::Iterate
}
fn exports_cycle_initial(_db: &dyn Db, _file: File) -> FxHashSet<Name> {
FxHashSet::default()
fn exports_cycle_initial(_db: &dyn Db, _file: File) -> Box<[Name]> {
Box::default()
}
#[salsa::tracked(return_ref, cycle_fn=exports_cycle_recover, cycle_initial=exports_cycle_initial)]
pub(super) fn exported_names(db: &dyn Db, file: File) -> FxHashSet<Name> {
pub(super) fn exported_names(db: &dyn Db, file: File) -> Box<[Name]> {
let module = parsed_module(db.upcast(), file);
let mut finder = ExportFinder::new(db, file);
finder.visit_body(module.suite());
finder.exports
finder.resolve_exports()
}
struct ExportFinder<'db> {
db: &'db dyn Db,
file: File,
visiting_stub_file: bool,
exports: FxHashSet<Name>,
exports: FxHashMap<&'db Name, PossibleExportKind>,
dunder_all: DunderAll,
}
impl<'db> ExportFinder<'db> {
@ -64,30 +65,59 @@ impl<'db> ExportFinder<'db> {
db,
file,
visiting_stub_file: file.is_stub(db.upcast()),
exports: FxHashSet::default(),
exports: FxHashMap::default(),
dunder_all: DunderAll::NotPresent,
}
}
fn possibly_add_export(&mut self, name: &Name) {
if name.starts_with('_') {
return;
fn possibly_add_export(&mut self, export: &'db Name, kind: PossibleExportKind) {
self.exports.insert(export, kind);
if export == "__all__" {
self.dunder_all = DunderAll::Present;
}
}
fn resolve_exports(self) -> Box<[Name]> {
match self.dunder_all {
DunderAll::NotPresent => self
.exports
.into_iter()
.filter_map(|(name, kind)| {
if kind == PossibleExportKind::StubImportWithoutRedundantAlias {
return None;
}
if name.starts_with('_') {
return None;
}
Some(name.clone())
})
.collect(),
DunderAll::Present => self.exports.into_keys().cloned().collect(),
}
self.exports.insert(name.clone());
}
}
impl<'db> Visitor<'db> for ExportFinder<'db> {
fn visit_alias(&mut self, alias: &'db ast::Alias) {
let ast::Alias { name, asname, .. } = alias;
if self.visiting_stub_file {
// If the source is a stub, names defined by imports are only exported
// if they use the explicit `foo as foo` syntax:
if asname.as_ref().is_some_and(|asname| asname.id == name.id) {
self.possibly_add_export(&name.id);
}
let ast::Alias {
name,
asname,
range: _,
} = alias;
let name = &name.id;
let asname = asname.as_ref().map(|asname| &asname.id);
// If the source is a stub, names defined by imports are only exported
// if they use the explicit `foo as foo` syntax:
let kind = if self.visiting_stub_file && asname.is_none_or(|asname| asname != name) {
PossibleExportKind::StubImportWithoutRedundantAlias
} else {
self.possibly_add_export(&asname.as_ref().unwrap_or(name).id);
}
PossibleExportKind::Normal
};
self.possibly_add_export(asname.unwrap_or(name), kind);
}
fn visit_pattern(&mut self, pattern: &'db ast::Pattern) {
@ -106,7 +136,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
// all names with leading underscores, but this will not always be the case
// (in the future we will want to support modules with `__all__ = ['_']`).
if name != "_" {
self.possibly_add_export(&name.id);
self.possibly_add_export(&name.id, PossibleExportKind::Normal);
}
}
}
@ -120,12 +150,12 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
self.visit_pattern(pattern);
}
if let Some(rest) = rest {
self.possibly_add_export(&rest.id);
self.possibly_add_export(&rest.id, PossibleExportKind::Normal);
}
}
ast::Pattern::MatchStar(ast::PatternMatchStar { name, range: _ }) => {
if let Some(name) = name {
self.possibly_add_export(&name.id);
self.possibly_add_export(&name.id, PossibleExportKind::Normal);
}
}
ast::Pattern::MatchSequence(_)
@ -137,7 +167,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
}
}
fn visit_stmt(&mut self, stmt: &'db ruff_python_ast::Stmt) {
fn visit_stmt(&mut self, stmt: &'db ast::Stmt) {
match stmt {
ast::Stmt::ClassDef(ast::StmtClassDef {
name,
@ -147,7 +177,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
body: _, // We don't want to visit the body of the class
range: _,
}) => {
self.possibly_add_export(&name.id);
self.possibly_add_export(&name.id, PossibleExportKind::Normal);
for decorator in decorator_list {
self.visit_decorator(decorator);
}
@ -155,6 +185,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
self.visit_arguments(arguments);
}
}
ast::Stmt::FunctionDef(ast::StmtFunctionDef {
name,
decorator_list,
@ -165,7 +196,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
range: _,
is_async: _,
}) => {
self.possibly_add_export(&name.id);
self.possibly_add_export(&name.id, PossibleExportKind::Normal);
for decorator in decorator_list {
self.visit_decorator(decorator);
}
@ -174,6 +205,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
self.visit_expr(returns);
}
}
ast::Stmt::AnnAssign(ast::StmtAnnAssign {
target,
value,
@ -189,6 +221,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
self.visit_expr(value);
}
}
ast::Stmt::TypeAlias(ast::StmtTypeAlias {
name,
type_params: _,
@ -199,20 +232,22 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
// Neither walrus expressions nor statements cannot appear in type aliases;
// no need to recursively visit the `value` or `type_params`
}
ast::Stmt::ImportFrom(node) => {
let mut found_star = false;
for name in &node.names {
if &name.name.id == "*" {
if !found_star {
found_star = true;
self.exports.extend(
for export in
ModuleName::from_import_statement(self.db, self.file, node)
.ok()
.and_then(|module_name| resolve_module(self.db, &module_name))
.iter()
.flat_map(|module| exported_names(self.db, module.file()))
.cloned(),
);
{
self.possibly_add_export(export, PossibleExportKind::Normal);
}
}
} else {
self.visit_alias(name);
@ -248,7 +283,7 @@ impl<'db> Visitor<'db> for ExportFinder<'db> {
match expr {
ast::Expr::Name(ast::ExprName { id, ctx, range: _ }) => {
if ctx.is_store() {
self.possibly_add_export(id);
self.possibly_add_export(id, PossibleExportKind::Normal);
}
}
@ -325,7 +360,8 @@ impl<'db> Visitor<'db> for WalrusFinder<'_, 'db> {
range: _,
}) = &**target
{
self.export_finder.possibly_add_export(id);
self.export_finder
.possibly_add_export(id, PossibleExportKind::Normal);
}
}
@ -358,3 +394,15 @@ impl<'db> Visitor<'db> for WalrusFinder<'_, 'db> {
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PossibleExportKind {
Normal,
StubImportWithoutRedundantAlias,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DunderAll {
NotPresent,
Present,
}