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

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`
+reveal_type(_private)  # revealed: bool
-# error: [unresolved-reference]
+reveal_type(__protected)  # revealed: bool
-reveal_type(_private)  # revealed: Unknown
+reveal_type(__dunder__)  # revealed: bool
-# error: [unresolved-reference]
+reveal_type(___thunder___)  # revealed: bool
-reveal_type(__protected)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(__dunder__)  # revealed: Unknown
-# error: [unresolved-reference]
-reveal_type(___thunder___)  # revealed: Unknown
 
 # 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

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<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> {
+fn exports_cycle_initial(_db: &dyn Db, _file: File) -> Box<[Name]> {
-    FxHashSet::default()
+    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) {
+    fn possibly_add_export(&mut self, export: &'db Name, kind: PossibleExportKind) {
-        if name.starts_with('_') {
+        self.exports.insert(export, kind);
-            return;
+
+        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;
+        let ast::Alias {
-        if self.visiting_stub_file {
+            name,
-            // If the source is a stub, names defined by imports are only exported
+            asname,
-            // if they use the explicit `foo as foo` syntax:
+            range: _,
-            if asname.as_ref().is_some_and(|asname| asname.id == name.id) {
+        } = alias;
-                self.possibly_add_export(&name.id);
+
-            }
+        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,
+}