[red-knot] Improve handling of visibility constraints in external modules when resolving * imports

Author: AlexWaygood

crates/red_knot_python_semantic/resources/mdtest/import/star.md

@@ -179,8 +179,12 @@ match 42:
         ...
     case object(foo=R):
         ...
+
+match 56:
     case object(S):
         ...
+
+match 12345:
     case T:
         ...
 
@@ -379,15 +383,21 @@ reveal_type(s)  # revealed: Unknown
 # error: [unresolved-reference]
 reveal_type(t)  # revealed: Unknown
 
-# TODO: these should all reveal `Unknown | int`.
+# TODO: these should all reveal `Unknown | int` and should not emit errors.
 # (We don't generally model elsewhere in red-knot that bindings from walruses
 # "leak" from comprehension scopes into outer scopes, but we should.)
 # See https://github.com/astral-sh/ruff/issues/16954
+# error: [unresolved-reference]
 reveal_type(g)  # revealed: Unknown
+# error: [unresolved-reference]
 reveal_type(i)  # revealed: Unknown
+# error: [unresolved-reference]
 reveal_type(k)  # revealed: Unknown
+# error: [unresolved-reference]
 reveal_type(m)  # revealed: Unknown
+# error: [unresolved-reference]
 reveal_type(o)  # revealed: Unknown
+# error: [unresolved-reference]
 reveal_type(q)  # revealed: Unknown
 ```
 
@@ -549,26 +559,54 @@ from a import *
 
 reveal_type(X)  # revealed: bool
 
-# TODO: should emit error: [unresolved-reference]
+# error: [unresolved-reference]
 reveal_type(Y)  # revealed: Unknown
 
-# TODO: The `*` import should not be considered a redefinition
+# The `*` import should not be considered a redefinition
 # of the global variable in this module, as the symbol in
 # the `a` module is in a branch that is statically known
 # to be dead code given the `python-version` configuration.
-# Thus this should reveal `Literal[True]`.
-reveal_type(Z)  # revealed: Unknown
+# Thus this correctly reveals `Literal[True]` and does not have an error emitted
+reveal_type(Z)  # revealed: Literal[True]
 
 if sys.version_info >= (3, 12):
     from aa import *
 
-    # TODO: should reveal `Never`
+    # TODO: should not error because this branch is unreachable;
+    # should reveal `Never` (see https://github.com/astral-sh/ruff/issues/15967)
+    # error: [unresolved-reference]
     reveal_type(AA)  # revealed: Unknown
 
 # error: [unresolved-reference]
 reveal_type(AA)  # revealed: Unknown
 ```
 
+`c.py`:
+
+```py
+def coinflip() -> bool:
+    return True
+
+if coinflip():
+    Z: str = "Z"
+```
+
+`d.py`:
+
+```py
+Z = True
+
+from a import *
+from a import *
+from a import *
+
+reveal_type(Z)  # revealed: Literal[True]
+
+from c import *
+
+reveal_type(Z)  # revealed: Literal[True] | str
+```
+
 ### Relative `*` imports
 
 Relative `*` imports are also supported by Python:
@@ -821,7 +859,7 @@ from a import *
 reveal_type(X)  # revealed: bool
 reveal_type(Y)  # revealed: bool
 
-# TODO: should error with [unresolved-reference]
+# error: [unresolved-reference]
 reveal_type(Z)  # revealed: Unknown
 ```
 
@@ -856,7 +894,7 @@ from a import *
 reveal_type(X)  # revealed: bool
 reveal_type(Y)  # revealed: bool
 
-# TODO should have an [unresolved-reference] diagnostic
+# error: [unresolved-reference]
 reveal_type(Z)  # revealed: Unknown
 ```
 

crates/red_knot_python_semantic/src/import_resolution.rs

@@ -0,0 +1,104 @@
+//! Utilities for resolving import statements that are used from the `types` submodule and the `symbol` submodule.
+
+use ruff_db::files::File;
+use ruff_python_ast as ast;
+
+use crate::{
+    db::Db,
+    module_name::{ModuleName, ModuleNameResolutionError},
+    module_resolver::resolve_module,
+    semantic_index::{definition::StarImportDefinitionKind, symbol::SymbolTable},
+    symbol::SymbolAndQualifiers,
+    types::Type,
+};
+
+pub(crate) fn resolve_star_import_definition<'db>(
+    db: &'db dyn Db,
+    file: File,
+    star_import_definition: &StarImportDefinitionKind,
+    symbol_table: &SymbolTable,
+) -> Result<SymbolAndQualifiers<'db>, UnresolvedImportFromError> {
+    let import_from = star_import_definition.import();
+    let alias = star_import_definition.alias();
+    let symbol_id = star_import_definition.symbol_id();
+
+    let (_, module_type) = resolve_import_from_module(db, file, import_from, alias)?;
+    let defined_name = symbol_table.symbol(symbol_id).name();
+    let imported_symbol = module_type.member(db, defined_name);
+
+    Ok(imported_symbol)
+}
+
+/// Resolve the [`ModuleName`], and the type of the module, being referred to by an
+/// [`ast::StmtImportFrom`] node. Emit a diagnostic if the module cannot be resolved.
+pub(crate) fn resolve_import_from_module<'db>(
+    db: &'db dyn Db,
+    file: File,
+    import_from: &ast::StmtImportFrom,
+    alias: &ast::Alias,
+) -> Result<(ModuleName, Type<'db>), UnresolvedImportFromError> {
+    let ast::StmtImportFrom { module, level, .. } = import_from;
+    let module = module.as_deref();
+
+    tracing::trace!(
+        "Resolving imported object `{}` from module `{}` into file `{}`",
+        alias.name,
+        format_import_from_module(*level, module),
+        file.path(db),
+    );
+
+    let module_name =
+        ModuleName::from_import_statement(db, file, import_from).map_err(|err| match err {
+            ModuleNameResolutionError::InvalidSyntax => {
+                tracing::debug!("Failed to resolve import due to invalid syntax");
+                UnresolvedImportFromError::InvalidSyntax
+            }
+            ModuleNameResolutionError::TooManyDots => {
+                tracing::debug!(
+                    "Relative module resolution `{}` failed: too many leading dots",
+                    format_import_from_module(*level, module),
+                );
+                UnresolvedImportFromError::UnresolvedModule
+            }
+            ModuleNameResolutionError::UnknownCurrentModule => {
+                tracing::debug!(
+                "Relative module resolution `{}` failed; could not resolve file `{}` to a module",
+                format_import_from_module(*level, module),
+                file.path(db)
+            );
+                UnresolvedImportFromError::UnresolvedModule
+            }
+        })?;
+
+    module_type_from_name(db, file, &module_name)
+        .map(|module_type| (module_name, module_type))
+        .ok_or(UnresolvedImportFromError::UnresolvedModule)
+}
+
+fn format_import_from_module(level: u32, module: Option<&str>) -> String {
+    format!(
+        "{}{}",
+        ".".repeat(level as usize),
+        module.unwrap_or_default()
+    )
+}
+
+pub(crate) fn module_type_from_name<'db>(
+    db: &'db dyn Db,
+    file: File,
+    module_name: &ModuleName,
+) -> Option<Type<'db>> {
+    resolve_module(db, module_name).map(|module| Type::module_literal(db, file, module))
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub(crate) enum UnresolvedImportFromError {
+    UnresolvedModule,
+    InvalidSyntax,
+}
+
+impl UnresolvedImportFromError {
+    pub(crate) fn is_invalid_syntax(self) -> bool {
+        matches!(self, UnresolvedImportFromError::InvalidSyntax)
+    }
+}

crates/red_knot_python_semantic/src/lib.rs

@@ -14,6 +14,7 @@ pub use site_packages::SysPrefixPathOrigin;
 
 pub mod ast_node_ref;
 mod db;
+mod import_resolution;
 pub mod lint;
 pub(crate) mod list;
 mod module_name;

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -398,6 +398,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         let kind = unsafe { definition_node.into_owned(self.module.clone()) };
         let category = kind.category(self.file.is_stub(self.db.upcast()));
         let is_reexported = kind.is_reexported();
+        let is_star_import = kind.is_star_import();
 
         let definition = Definition::new(
             self.db,
@@ -425,7 +426,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         let use_def = self.current_use_def_map_mut();
         match category {
             DefinitionCategory::DeclarationAndBinding => {
-                use_def.record_declaration_and_binding(symbol, definition);
+                use_def.record_declaration_and_binding(symbol, definition, is_star_import);
             }
             DefinitionCategory::Declaration => use_def.record_declaration(symbol, definition),
             DefinitionCategory::Binding => use_def.record_binding(symbol, definition),

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -559,6 +559,10 @@ impl DefinitionKind<'_> {
         }
     }
 
+    pub(crate) const fn is_star_import(&self) -> bool {
+        matches!(self, DefinitionKind::StarImport(_))
+    }
+
     /// Returns the [`TextRange`] of the definition target.
     ///
     /// A definition target would mainly be the node representing the symbol being defined i.e.,

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -321,6 +321,11 @@ pub(crate) struct UseDefMap<'db> {
     /// eager scope.
     eager_bindings: EagerBindings,
 
+    /// Snapshot of the [`SymbolState`] for a given symbol immediately prior
+    /// to a star-import [`Definition`] that appears as though it may have provided
+    /// a binding for that symbol.
+    bindings_at_star_imports: FxHashMap<Definition<'db>, SymbolState>,
+
     /// Whether or not the start of the scope is visible.
     /// This is used to check if the function can implicitly return `None`.
     /// For example:
@@ -352,6 +357,15 @@ impl<'db> UseDefMap<'db> {
         self.bindings_iterator(self.public_symbols[symbol].bindings())
     }
 
+    /// Iterate over the public bindings that were visible for a symbol immediately prior
+    /// to a star-import [`Definition`] that may have created a binding for that symbol.
+    pub(crate) fn public_bindings_prior_to_star_import(
+        &self,
+        star_import_definition: Definition<'db>,
+    ) -> BindingWithConstraintsIterator<'_, 'db> {
+        self.bindings_iterator(self.bindings_at_star_imports[&star_import_definition].bindings())
+    }
+
     pub(crate) fn eager_bindings(
         &self,
         eager_bindings: ScopedEagerBindingsId,
@@ -570,6 +584,11 @@ pub(super) struct UseDefMapBuilder<'db> {
     /// Snapshot of bindings in this scope that can be used to resolve a reference in a nested
     /// eager scope.
     eager_bindings: EagerBindings,
+
+    /// Snapshot of the [`SymbolState`] for a given symbol immediately prior
+    /// to a star-import [`Definition`] that appears as though it may have provided
+    /// a binding for that symbol.
+    bindings_at_star_imports: FxHashMap<Definition<'db>, SymbolState>,
 }
 
 impl Default for UseDefMapBuilder<'_> {
@@ -585,6 +604,7 @@ impl Default for UseDefMapBuilder<'_> {
             bindings_by_declaration: FxHashMap::default(),
             symbol_states: IndexVec::new(),
             eager_bindings: EagerBindings::default(),
+            bindings_at_star_imports: FxHashMap::default(),
         }
     }
 }
@@ -687,11 +707,16 @@ impl<'db> UseDefMapBuilder<'db> {
         &mut self,
         symbol: ScopedSymbolId,
         definition: Definition<'db>,
+        is_star_import: bool,
     ) {
         // We don't need to store anything in self.bindings_by_declaration or
         // self.declarations_by_binding.
         let def_id = self.all_definitions.push(Some(definition));
         let symbol_state = &mut self.symbol_states[symbol];
+        if is_star_import {
+            self.bindings_at_star_imports
+                .insert(definition, symbol_state.clone());
+        }
         symbol_state.record_declaration(def_id);
         symbol_state.record_binding(def_id, self.scope_start_visibility);
     }
@@ -796,6 +821,7 @@ impl<'db> UseDefMapBuilder<'db> {
         self.declarations_by_binding.shrink_to_fit();
         self.bindings_by_declaration.shrink_to_fit();
         self.eager_bindings.shrink_to_fit();
+        self.bindings_at_star_imports.shrink_to_fit();
 
         UseDefMap {
             all_definitions: self.all_definitions,
@@ -807,6 +833,7 @@ impl<'db> UseDefMapBuilder<'db> {
             declarations_by_binding: self.declarations_by_binding,
             bindings_by_declaration: self.bindings_by_declaration,
             eager_bindings: self.eager_bindings,
+            bindings_at_star_imports: self.bindings_at_star_imports,
             scope_start_visibility: self.scope_start_visibility,
         }
     }

crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs

@@ -313,7 +313,7 @@ impl SymbolBindings {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, salsa::Update)]
 pub(super) struct SymbolState {
     declarations: SymbolDeclarations,
     bindings: SymbolBindings,