diff --git a/crates/ty_python_semantic/src/semantic_index.rs b/crates/ty_python_semantic/src/semantic_index.rs index a38b0a7ded..9963ef23af 100644 --- a/crates/ty_python_semantic/src/semantic_index.rs +++ b/crates/ty_python_semantic/src/semantic_index.rs @@ -79,7 +79,10 @@ pub(crate) fn place_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc(db: &'db dyn Db, file: File) -> Arc> { +pub(crate) fn imported_modules<'db>( + db: &'db dyn Db, + file: File, +) -> Arc> { semantic_index(db, file).imported_modules.clone() } @@ -246,7 +249,7 @@ pub(crate) struct SemanticIndex<'db> { ast_ids: IndexVec, /// The set of modules that are imported anywhere within this file. - imported_modules: Arc>, + imported_modules: Arc>, /// Flags about the global scope (code usage impacting inference) has_future_annotations: bool, @@ -583,6 +586,12 @@ impl<'db> SemanticIndex<'db> { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq, get_size2::GetSize)] +pub(crate) enum ImportKind { + Import, + ImportFrom, +} + pub(crate) struct AncestorsIter<'a> { scopes: &'a IndexSlice, next_id: Option, diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 373c80f7ee..3591bd66d2 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -47,7 +47,7 @@ use crate::semantic_index::symbol::{ScopedSymbolId, Symbol}; use crate::semantic_index::use_def::{ EnclosingSnapshotKey, FlowSnapshot, ScopedEnclosingSnapshotId, UseDefMapBuilder, }; -use crate::semantic_index::{ExpressionsScopeMap, SemanticIndex, VisibleAncestorsIter}; +use crate::semantic_index::{ExpressionsScopeMap, ImportKind, SemanticIndex, VisibleAncestorsIter}; use crate::semantic_model::HasTrackedScope; use crate::unpack::{EvaluationMode, Unpack, UnpackKind, UnpackPosition, UnpackValue}; use crate::{Db, Program}; @@ -110,7 +110,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> { scopes_by_expression: ExpressionsScopeMapBuilder, definitions_by_node: FxHashMap>, expressions_by_node: FxHashMap>, - imported_modules: FxHashSet, + imported_modules: FxHashMap, seen_submodule_imports: FxHashSet, /// Hashset of all [`FileScopeId`]s that correspond to [generator functions]. /// @@ -150,7 +150,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { expressions_by_node: FxHashMap::default(), seen_submodule_imports: FxHashSet::default(), - imported_modules: FxHashSet::default(), + imported_modules: FxHashMap::default(), generator_functions: FxHashSet::default(), enclosing_snapshots: FxHashMap::default(), @@ -1474,7 +1474,11 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { // 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()); + self.imported_modules.extend( + module_name + .ancestors() + .zip(std::iter::repeat(ImportKind::Import)), + ); } let (symbol_name, is_reexported) = if let Some(asname) = &alias.asname { diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index e9311547d7..e94d69454a 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -42,7 +42,7 @@ use crate::place::{ use crate::semantic_index::definition::{Definition, DefinitionKind}; use crate::semantic_index::place::ScopedPlaceId; use crate::semantic_index::scope::ScopeId; -use crate::semantic_index::{imported_modules, place_table, semantic_index}; +use crate::semantic_index::{ImportKind, imported_modules, place_table, semantic_index}; use crate::suppression::check_suppressions; use crate::types::bound_super::BoundSuperType; use crate::types::builder::RecursivelyDefined; @@ -13236,12 +13236,22 @@ impl<'db> ModuleLiteralType<'db> { /// /// We instead prefer handling most other import effects as definitions in the scope of /// the current file (i.e. [`crate::semantic_index::definition::ImportFromDefinitionNodeRef`]). - fn available_submodule_attributes(&self, db: &'db dyn Db) -> impl Iterator { + fn available_submodule_attributes( + &self, + db: &'db dyn Db, + ) -> impl Iterator { self.importing_file(db) .into_iter() .flat_map(|file| imported_modules(db, file)) - .filter_map(|submodule_name| submodule_name.relative_to(self.module(db).name(db))) - .filter_map(|relative_submodule| relative_submodule.components().next().map(Name::from)) + .filter_map(|(submodule_name, kind)| { + Some((submodule_name.relative_to(self.module(db).name(db))?, kind)) + }) + .filter_map(|(relative_submodule, kind)| { + relative_submodule + .components() + .next() + .map(|module| (Name::from(module), *kind)) + }) } fn resolve_submodule(self, db: &'db dyn Db, name: &str) -> Option> { @@ -13285,19 +13295,27 @@ impl<'db> ModuleLiteralType<'db> { .member(db, "__dict__"); } - // 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, but is the opposite of mypy's current behavior.) - if self.available_submodule_attributes(db).contains(name) { - if let Some(submodule) = self.resolve_submodule(db, name) { - return Place::bound(submodule).into(); - } + let mut submodule_type = None; + + let available_submodule_kind = self + .available_submodule_attributes(db) + .find_map(|(attr, kind)| (attr == name).then_some(kind)); + + if available_submodule_kind.is_some() { + submodule_type = self.resolve_submodule(db, name); + } + + // if we're in a module `foo` and `foo` contains `import a.b`, + // and the package `a` has a submodule `b`, we assume that the + // attribute access `a.b` inside `foo` will resolve to the submodule + // `a.b` *even if* `a/__init__.py` also defines a symbol `b` (e.g. `b = 42`). + // This is a heuristic, but it's almost certainly what will actually happen + // at runtime. However, if `foo` only contains `from a.b import , + // we prioritise the `b` attribute in `a/__init__.py` over the submodule `a.b`. + if available_submodule_kind == Some(ImportKind::Import) + && let Some(submodule) = submodule_type + { + return Place::bound(submodule).into(); } let place_and_qualifiers = self @@ -13306,12 +13324,16 @@ impl<'db> ModuleLiteralType<'db> { .map(|file| imported_symbol(db, file, name, None)) .unwrap_or_default(); - // If the normal lookup failed, try to call the module's `__getattr__` function - if place_and_qualifiers.place.is_undefined() { - return self.try_module_getattr(db, name); + if !place_and_qualifiers.is_undefined() { + return place_and_qualifiers; } - place_and_qualifiers + if let Some(submodule) = submodule_type { + return Place::bound(submodule).into(); + } + + // If the normal lookup failed, try to call the module's `__getattr__` function + self.try_module_getattr(db, name) } } diff --git a/crates/ty_python_semantic/src/types/list_members.rs b/crates/ty_python_semantic/src/types/list_members.rs index 4e4a32c294..5c538a5631 100644 --- a/crates/ty_python_semantic/src/types/list_members.rs +++ b/crates/ty_python_semantic/src/types/list_members.rs @@ -371,7 +371,7 @@ impl<'db> AllMembers<'db> { self.members .extend(literal.available_submodule_attributes(db).filter_map( - |submodule_name| { + |(submodule_name, _)| { let ty = literal.resolve_submodule(db, &submodule_name)?; let name = submodule_name.clone(); Some(Member { name, ty })