diff --git a/crates/red_knot_python_semantic/resources/mdtest/import/basic.md b/crates/red_knot_python_semantic/resources/mdtest/import/basic.md index fb802bb215ebd..fc19ad815640a 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/import/basic.md +++ b/crates/red_knot_python_semantic/resources/mdtest/import/basic.md @@ -25,3 +25,69 @@ reveal_type(D) # revealed: Literal[C] ```py path=b.py class C: ... ``` + +## Nested + +```py +import a.b + +reveal_type(a.b.C) # revealed: Literal[C] +``` + +```py path=a/__init__.py +``` + +```py path=a/b.py +class C: ... +``` + +## Deeply nested + +```py +import a.b.c + +reveal_type(a.b.c.C) # revealed: Literal[C] +``` + +```py path=a/__init__.py +``` + +```py path=a/b/__init__.py +``` + +```py path=a/b/c.py +class C: ... +``` + +## Nested with rename + +```py +import a.b as b + +reveal_type(b.C) # revealed: Literal[C] +``` + +```py path=a/__init__.py +``` + +```py path=a/b.py +class C: ... +``` + +## Deeply nested with rename + +```py +import a.b.c as c + +reveal_type(c.C) # revealed: Literal[C] +``` + +```py path=a/__init__.py +``` + +```py path=a/b/__init__.py +``` + +```py path=a/b/c.py +class C: ... +``` diff --git a/crates/red_knot_python_semantic/resources/mdtest/import/tracking.md b/crates/red_knot_python_semantic/resources/mdtest/import/tracking.md new file mode 100644 index 0000000000000..6d10126c775fc --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/import/tracking.md @@ -0,0 +1,100 @@ +# Tracking imported modules + +These tests depend on how we track which modules have been imported. There are currently two +characteristics of our module tracking that can lead to inaccuracies: + +- Imports are tracked on a per-file basis. At runtime, importing a submodule in one file makes that + submodule globally available via any reference to the containing package. We will flag an error + if a file tries to access a submodule without there being an import of that submodule _in that + same file_. + + This is a purposeful decision, and not one we plan to change. If a module wants to re-export some + other module that it imports, there are ways to do that (tested below) that are blessed by the + typing spec and that are visible to our file-scoped import tracking. + +- Imports are tracked flow-insensitively: submodule accesses are allowed and resolved if that + submodule is imported _anywhere in the file_. This handles the common case where all imports are + grouped at the top of the file, and is easiest to implement. We might revisit this decision and + track submodule imports flow-sensitively, in which case we will have to update the assertions in + some of these tests. + +## Import submodule later in file + +This test highlights our flow-insensitive analysis, since we access the `a.b` submodule before it +has been imported. + +```py +import a + +# Would be an error with flow-sensitive tracking +reveal_type(a.b.C) # revealed: Literal[C] + +import a.b +``` + +```py path=a/__init__.py +``` + +```py path=a/b.py +class C: ... +``` + +## Rename a re-export + +This test highlights how import tracking is local to each file, but specifically to the file where a +containing module is first referenced. This allows the main module to see that `q.a` contains a +submodule `b`, even though `a.b` is never imported in the main module. + +```py +from q import a, b + +reveal_type(b) # revealed: +reveal_type(b.C) # revealed: Literal[C] + +reveal_type(a.b) # revealed: +reveal_type(a.b.C) # revealed: Literal[C] +``` + +```py path=a/__init__.py +``` + +```py path=a/b.py +class C: ... +``` + +```py path=q.py +import a as a +import a.b as b +``` + +## Attribute overrides submodule + +Technically, either a submodule or a non-module attribute could shadow the other, depending on the +ordering of when the submodule is loaded relative to the parent module's `__init__.py` file being +evaluated. We have chosen to always have the submodule take priority. (This matches pyright's +current behavior, and opposite of mypy's current behavior.) + +```py +import sub.b +import attr.b + +# In the Python interpreter, `attr.b` is Literal[1] +reveal_type(sub.b) # revealed: +reveal_type(attr.b) # revealed: +``` + +```py path=sub/__init__.py +b = 1 +``` + +```py path=sub/b.py +``` + +```py path=attr/__init__.py +from . import b as _ + +b = 1 +``` + +```py path=attr/b.py +``` diff --git a/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md b/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md index 56fd8233c3eeb..192b94b08c822 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md +++ b/crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md @@ -61,10 +61,8 @@ class B: ... ```py path=a/test.py import a.b -# TODO: no diagnostic -# error: [unresolved-attribute] def f(c: type[a.b.C]): - reveal_type(c) # revealed: @Todo(unsupported type[X] special form) + reveal_type(c) # revealed: type[C] ``` ```py path=a/__init__.py diff --git a/crates/red_knot_python_semantic/src/module_name.rs b/crates/red_knot_python_semantic/src/module_name.rs index 885c6adf9f4bf..3575b3192064f 100644 --- a/crates/red_knot_python_semantic/src/module_name.rs +++ b/crates/red_knot_python_semantic/src/module_name.rs @@ -186,6 +186,26 @@ impl ModuleName { self.0.push('.'); self.0.push_str(other); } + + /// Returns an iterator of this module name and all of its parent modules. + /// + /// # Examples + /// + /// ``` + /// use red_knot_python_semantic::ModuleName; + /// + /// assert_eq!( + /// ModuleName::new_static("foo.bar.baz").unwrap().ancestors().collect::>(), + /// vec![ + /// ModuleName::new_static("foo.bar.baz").unwrap(), + /// ModuleName::new_static("foo.bar").unwrap(), + /// ModuleName::new_static("foo").unwrap(), + /// ], + /// ); + /// ``` + pub fn ancestors(&self) -> impl Iterator { + std::iter::successors(Some(self.clone()), Self::parent) + } } impl Deref for ModuleName { diff --git a/crates/red_knot_python_semantic/src/module_resolver/module.rs b/crates/red_knot_python_semantic/src/module_resolver/module.rs index e2c1e939572cc..1891d3085d52c 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/module.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/module.rs @@ -7,7 +7,7 @@ use super::path::SearchPath; use crate::module_name::ModuleName; /// Representation of a Python module. -#[derive(Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq, Hash)] pub struct Module { inner: Arc, } @@ -61,7 +61,7 @@ impl std::fmt::Debug for Module { } } -#[derive(PartialEq, Eq)] +#[derive(PartialEq, Eq, Hash)] struct ModuleInner { name: ModuleName, kind: ModuleKind, diff --git a/crates/red_knot_python_semantic/src/semantic_index.rs b/crates/red_knot_python_semantic/src/semantic_index.rs index 2771ad301e380..73789f3048785 100644 --- a/crates/red_knot_python_semantic/src/semantic_index.rs +++ b/crates/red_knot_python_semantic/src/semantic_index.rs @@ -1,13 +1,14 @@ use std::iter::FusedIterator; use std::sync::Arc; -use rustc_hash::{FxBuildHasher, FxHashMap}; +use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; use salsa::plumbing::AsId; use ruff_db::files::File; use ruff_db::parsed::parsed_module; use ruff_index::{IndexSlice, IndexVec}; +use crate::module_name::ModuleName; use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey; use crate::semantic_index::ast_ids::AstIds; use crate::semantic_index::builder::SemanticIndexBuilder; @@ -60,6 +61,12 @@ pub(crate) fn symbol_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc(db: &'db dyn Db, file: File) -> Arc> { + semantic_index(db, file).imported_modules.clone() +} + /// Returns the use-def map for a specific `scope`. /// /// Using [`use_def_map`] over [`semantic_index`] has the advantage that @@ -116,6 +123,9 @@ pub(crate) struct SemanticIndex<'db> { /// changing a file invalidates all dependents. ast_ids: IndexVec, + /// The set of modules that are imported anywhere within this file. + imported_modules: Arc>, + /// Flags about the global scope (code usage impacting inference) has_future_annotations: bool, } diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 27e657aba959b..fdd6d11534195 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use except_handlers::TryNodeContextStackManager; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use ruff_db::files::File; use ruff_db::parsed::ParsedModule; @@ -12,6 +12,7 @@ use ruff_python_ast::visitor::{walk_expr, walk_pattern, walk_stmt, Visitor}; use ruff_python_ast::{BoolOp, Expr}; use crate::ast_node_ref::AstNodeRef; +use crate::module_name::ModuleName; use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey; use crate::semantic_index::ast_ids::AstIdsBuilder; use crate::semantic_index::definition::{ @@ -79,6 +80,7 @@ pub(super) struct SemanticIndexBuilder<'db> { scopes_by_expression: FxHashMap, definitions_by_node: FxHashMap>, expressions_by_node: FxHashMap>, + imported_modules: FxHashSet, } impl<'db> SemanticIndexBuilder<'db> { @@ -105,6 +107,8 @@ impl<'db> SemanticIndexBuilder<'db> { scopes_by_node: FxHashMap::default(), definitions_by_node: FxHashMap::default(), expressions_by_node: FxHashMap::default(), + + imported_modules: FxHashSet::default(), }; builder.push_scope_with_parent(NodeWithScopeRef::Module, None); @@ -558,6 +562,7 @@ impl<'db> SemanticIndexBuilder<'db> { scopes_by_expression: self.scopes_by_expression, scopes_by_node: self.scopes_by_node, use_def_maps, + imported_modules: Arc::new(self.imported_modules), has_future_annotations: self.has_future_annotations, } } @@ -661,6 +666,12 @@ where } ast::Stmt::Import(node) => { for alias in &node.names { + // Mark the imported module, and all of its parents, as being imported in this + // file. + if let Some(module_name) = ModuleName::new(&alias.name) { + self.imported_modules.extend(module_name.ancestors()); + } + let symbol_name = if let Some(asname) = &alias.asname { asname.id.clone() } else { diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index b714f3b0f9636..0a1c99513b3f5 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -14,13 +14,14 @@ pub(crate) use self::infer::{ infer_deferred_types, infer_definition_types, infer_expression_types, infer_scope_types, }; pub(crate) use self::signatures::Signature; -use crate::module_resolver::file_to_module; +use crate::module_name::ModuleName; +use crate::module_resolver::{file_to_module, resolve_module}; use crate::semantic_index::ast_ids::HasScopedExpressionId; use crate::semantic_index::definition::Definition; use crate::semantic_index::symbol::{self as symbol, ScopeId, ScopedSymbolId}; use crate::semantic_index::{ - global_scope, semantic_index, symbol_table, use_def_map, BindingWithConstraints, - BindingWithConstraintsIterator, DeclarationsIterator, + global_scope, imported_modules, semantic_index, symbol_table, use_def_map, + BindingWithConstraints, BindingWithConstraintsIterator, DeclarationsIterator, }; use crate::stdlib::{ builtins_symbol, core_module_symbol, typing_extensions_symbol, CoreStdlibModule, @@ -413,7 +414,7 @@ pub enum Type<'db> { /// A specific function object FunctionLiteral(FunctionType<'db>), /// A specific module object - ModuleLiteral(File), + ModuleLiteral(ModuleLiteralType<'db>), /// A specific class object ClassLiteral(ClassLiteralType<'db>), // The set of all class objects that are subclasses of the given class (C), spelled `type[C]`. @@ -475,15 +476,19 @@ impl<'db> Type<'db> { matches!(self, Type::ClassLiteral(..)) } - pub const fn into_module_literal(self) -> Option { + pub fn module_literal(db: &'db dyn Db, importing_file: File, submodule: Module) -> Self { + Self::ModuleLiteral(ModuleLiteralType::new(db, importing_file, submodule)) + } + + pub const fn into_module_literal(self) -> Option> { match self { - Type::ModuleLiteral(file) => Some(file), + Type::ModuleLiteral(module) => Some(module), _ => None, } } #[track_caller] - pub fn expect_module_literal(self) -> File { + pub fn expect_module_literal(self) -> ModuleLiteralType<'db> { self.into_module_literal() .expect("Expected a Type::ModuleLiteral variant") } @@ -1418,7 +1423,7 @@ impl<'db> Type<'db> { // TODO: attribute lookup on function type todo_type!().into() } - Type::ModuleLiteral(file) => { + Type::ModuleLiteral(module_ref) => { // `__dict__` is a very special member that is never overridden by module globals; // we should always look it up directly as an attribute on `types.ModuleType`, // never in the global scope of the module. @@ -1428,7 +1433,30 @@ impl<'db> Type<'db> { .member(db, "__dict__"); } - let global_lookup = symbol(db, global_scope(db, *file), name); + // If the file that originally imported the module has also imported a submodule + // named [name], then the result is (usually) that submodule, even if the module + // also defines a (non-module) symbol with that name. + // + // Note that technically, either the submodule or the non-module symbol could take + // priority, depending on the ordering of when the submodule is loaded relative to + // the parent module's `__init__.py` file being evaluated. That said, we have + // chosen to always have the submodule take priority. (This matches pyright's + // current behavior, and opposite of mypy's current behavior.) + if let Some(submodule_name) = ModuleName::new(name) { + let importing_file = module_ref.importing_file(db); + let imported_submodules = imported_modules(db, importing_file); + let mut full_submodule_name = module_ref.module(db).name().clone(); + full_submodule_name.extend(&submodule_name); + if imported_submodules.contains(&full_submodule_name) { + if let Some(submodule) = resolve_module(db, &full_submodule_name) { + let submodule_ty = Type::module_literal(db, importing_file, submodule); + return Symbol::Type(submodule_ty, Boundness::Bound); + } + } + } + + let global_lookup = + symbol(db, global_scope(db, module_ref.module(db).file()), name); // If it's unbound, check if it's present as an instance on `types.ModuleType` // or `builtins.object`. @@ -2860,6 +2888,18 @@ impl KnownFunction { } } +#[salsa::interned] +pub struct ModuleLiteralType<'db> { + /// The file in which this module was imported. + /// + /// We need this in order to know which submodules should be attached to it as attributes + /// (because the submodules were also imported in this file). + pub importing_file: File, + + /// The imported module. + pub module: Module, +} + /// Representation of a runtime class object. /// /// Does not in itself represent a type, @@ -3501,7 +3541,8 @@ pub(crate) mod tests { .class, ), Ty::StdlibModule(module) => { - Type::ModuleLiteral(resolve_module(db, &module.name()).unwrap().file()) + let module = resolve_module(db, &module.name()).unwrap(); + Type::module_literal(db, module.file(), module) } Ty::SliceLiteral(start, stop, step) => Type::SliceLiteral(SliceLiteralType::new( db, diff --git a/crates/red_knot_python_semantic/src/types/display.rs b/crates/red_knot_python_semantic/src/types/display.rs index 49a629aa9eef5..55c24a27f3004 100644 --- a/crates/red_knot_python_semantic/src/types/display.rs +++ b/crates/red_knot_python_semantic/src/types/display.rs @@ -79,8 +79,8 @@ impl Display for DisplayRepresentation<'_> { // `[Type::Todo]`'s display should be explicit that is not a valid display of // any other type Type::Todo(todo) => write!(f, "@Todo{todo}"), - Type::ModuleLiteral(file) => { - write!(f, "", file.path(self.db)) + Type::ModuleLiteral(module) => { + write!(f, "", module.module(self.db).name()) } // TODO functions and classes should display using a fully qualified name Type::ClassLiteral(ClassLiteralType { class }) => f.write_str(class.name(self.db)), diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 5c868163e3b69..e7d80e8495922 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -2116,22 +2116,55 @@ impl<'db> TypeInferenceBuilder<'db> { let ast::Alias { range: _, name, - asname: _, + asname, } = alias; - let module_ty = if let Some(module_name) = ModuleName::new(name) { - if let Some(module) = self.module_ty_from_name(&module_name) { - module + // The name of the module being imported + let Some(full_module_name) = ModuleName::new(name) else { + tracing::debug!("Failed to resolve import due to invalid syntax"); + self.add_declaration_with_binding( + alias.into(), + definition, + Type::Unknown, + Type::Unknown, + ); + return; + }; + + // Resolve the module being imported. + let Some(full_module_ty) = self.module_ty_from_name(&full_module_name) else { + self.diagnostics.add_unresolved_module(alias, 0, Some(name)); + self.add_declaration_with_binding( + alias.into(), + definition, + Type::Unknown, + Type::Unknown, + ); + return; + }; + + let binding_ty = if asname.is_some() { + // If we are renaming the imported module via an `as` clause, then we bind the resolved + // module's type to that name, even if that module is nested. + full_module_ty + } else if full_module_name.contains('.') { + // If there's no `as` clause and the imported module is nested, we're not going to bind + // the resolved module itself into the current scope; we're going to bind the top-most + // parent package of that module. + let topmost_parent_name = + ModuleName::new(full_module_name.components().next().unwrap()).unwrap(); + if let Some(topmost_parent_ty) = self.module_ty_from_name(&topmost_parent_name) { + topmost_parent_ty } else { - self.diagnostics.add_unresolved_module(alias, 0, Some(name)); Type::Unknown } } else { - tracing::debug!("Failed to resolve import due to invalid syntax"); - Type::Unknown + // If there's no `as` clause and the imported module isn't nested, then the imported + // module _is_ what we bind into the current scope. + full_module_ty }; - self.add_declaration_with_binding(alias.into(), definition, module_ty, module_ty); + self.add_declaration_with_binding(alias.into(), definition, binding_ty, binding_ty); } fn infer_import_from_statement(&mut self, import: &ast::StmtImportFrom) { @@ -2316,7 +2349,8 @@ impl<'db> TypeInferenceBuilder<'db> { } fn module_ty_from_name(&self, module_name: &ModuleName) -> Option> { - resolve_module(self.db, module_name).map(|module| Type::ModuleLiteral(module.file())) + resolve_module(self.db, module_name) + .map(|module| Type::module_literal(self.db, self.file, module)) } fn infer_decorator(&mut self, decorator: &ast::Decorator) -> Type<'db> {