diff --git a/crates/red_knot_python_semantic/resources/mdtest/import/star.md b/crates/red_knot_python_semantic/resources/mdtest/import/star.md index 99aa39e74ef4f..1e3c0b56d831f 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/import/star.md +++ b/crates/red_knot_python_semantic/resources/mdtest/import/star.md @@ -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 diff --git a/crates/red_knot_python_semantic/src/semantic_index/re_exports.rs b/crates/red_knot_python_semantic/src/semantic_index/re_exports.rs index a6b1c6529ec66..f16819f32980f 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/re_exports.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/re_exports.rs @@ -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, + _value: &[Name], _count: u32, _file: File, -) -> salsa::CycleRecoveryAction> { +) -> salsa::CycleRecoveryAction> { salsa::CycleRecoveryAction::Iterate } -fn exports_cycle_initial(_db: &dyn Db, _file: File) -> FxHashSet { - 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 { +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, + 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, +}