From bd8812127daa556bd86fa81c9a79f5f49a2feaa8 Mon Sep 17 00:00:00 2001 From: Aria Desires Date: Tue, 11 Nov 2025 13:04:42 -0500 Subject: [PATCH] [ty] support absolute `from` imports introducing local submodules in `__init__.py` files (#21372) By resolving `.` and the LHS of the from import during semantic indexing, we can check if the LHS is a submodule of `.`, and handle `from whatever.thispackage.x.y import z` exactly like we do `from .x.y import z`. Fixes https://github.com/astral-sh/ty/issues/1484 --- .../mdtest/import/nonstandard_conventions.md | 9 +- crates/ty_python_semantic/src/module_name.rs | 11 +++ .../src/semantic_index/builder.rs | 22 +++-- .../src/semantic_index/definition.rs | 14 +-- .../src/types/infer/builder.rs | 91 ++++++++++--------- 5 files changed, 79 insertions(+), 68 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/import/nonstandard_conventions.md b/crates/ty_python_semantic/resources/mdtest/import/nonstandard_conventions.md index 120ea0746d..44163c17b4 100644 --- a/crates/ty_python_semantic/resources/mdtest/import/nonstandard_conventions.md +++ b/crates/ty_python_semantic/resources/mdtest/import/nonstandard_conventions.md @@ -333,7 +333,7 @@ reveal_type(mypackage.nested.X) # revealed: Unknown ### In Non-Stub -`from mypackage.submodule import nested` in an `__init__.py` only creates `nested`. +`from mypackage.submodule import nested` in an `__init__.py` creates both `submodule` and `nested`. `mypackage/__init__.py`: @@ -357,12 +357,11 @@ X: int = 42 ```py import mypackage +reveal_type(mypackage.submodule) # revealed: # TODO: this would be nice to support -# error: "has no member `submodule`" -reveal_type(mypackage.submodule) # revealed: Unknown -# error: "has no member `submodule`" +# error: "has no member `nested`" reveal_type(mypackage.submodule.nested) # revealed: Unknown -# error: "has no member `submodule`" +# error: "has no member `nested`" reveal_type(mypackage.submodule.nested.X) # revealed: Unknown reveal_type(mypackage.nested) # revealed: reveal_type(mypackage.nested.X) # revealed: int diff --git a/crates/ty_python_semantic/src/module_name.rs b/crates/ty_python_semantic/src/module_name.rs index e1aa4509bd..b257d0a6df 100644 --- a/crates/ty_python_semantic/src/module_name.rs +++ b/crates/ty_python_semantic/src/module_name.rs @@ -295,6 +295,7 @@ impl ModuleName { Self::from_identifier_parts(db, importing_file, module.as_deref(), *level) } + /// Computes the absolute module name from the LHS components of `from LHS import RHS` pub(crate) fn from_identifier_parts( db: &dyn Db, importing_file: File, @@ -309,6 +310,16 @@ impl ModuleName { .ok_or(ModuleNameResolutionError::InvalidSyntax) } } + + /// Computes the absolute module name for the package this file belongs to. + /// + /// i.e. this resolves `.` + pub(crate) fn package_for_file( + db: &dyn Db, + importing_file: File, + ) -> Result { + Self::from_identifier_parts(db, importing_file, None, 1) + } } impl Deref for ModuleName { diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 8029a775fe..d7784c2cf1 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -1451,7 +1451,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { // If we see: // - // * `from .x.y import z` (must be relative!) + // * `from .x.y import z` (or `from whatever.thispackage.x.y`) // * And we are in an `__init__.py(i)` (hereafter `thispackage`) // * And this is the first time we've seen `from .x` in this module // * And we're in the global scope @@ -1465,14 +1465,18 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { // reasons but it works well for most practical purposes. In particular it's nice // that `x` can be freely overwritten, and that we don't assume that an import // in one function is visible in another function. - // - // TODO: Also support `from thispackage.x.y import z`? - if self.current_scope() == FileScopeId::global() - && node.level == 1 - && let Some(submodule) = &node.module - && let Some(parsed_submodule) = ModuleName::new(submodule.as_str()) - && let Some(direct_submodule) = parsed_submodule.components().next() + if node.module.is_some() + && self.current_scope().is_global() && self.file.is_package(self.db) + && let Ok(module_name) = ModuleName::from_identifier_parts( + self.db, + self.file, + node.module.as_deref(), + node.level, + ) + && let Ok(thispackage) = ModuleName::package_for_file(self.db, self.file) + && let Some(relative_submodule) = module_name.relative_to(&thispackage) + && let Some(direct_submodule) = relative_submodule.components().next() && !self.seen_submodule_imports.contains(direct_submodule) { self.seen_submodule_imports @@ -1482,7 +1486,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { let symbol = self.add_symbol(direct_submodule_name); self.add_definition( symbol.into(), - ImportFromSubmoduleDefinitionNodeRef { node, submodule }, + ImportFromSubmoduleDefinitionNodeRef { node }, ); } diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index db5d519560..85a7ff6aed 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -3,7 +3,6 @@ use std::ops::Deref; use ruff_db::files::{File, FileRange}; use ruff_db::parsed::{ParsedModuleRef, parsed_module}; use ruff_python_ast as ast; -use ruff_python_ast::name::Name; use ruff_text_size::{Ranged, TextRange}; use crate::Db; @@ -368,7 +367,6 @@ pub(crate) struct ImportFromDefinitionNodeRef<'ast> { #[derive(Copy, Clone, Debug)] pub(crate) struct ImportFromSubmoduleDefinitionNodeRef<'ast> { pub(crate) node: &'ast ast::StmtImportFrom, - pub(crate) submodule: &'ast ast::Identifier, } #[derive(Copy, Clone, Debug)] pub(crate) struct AssignmentDefinitionNodeRef<'ast, 'db> { @@ -450,10 +448,8 @@ impl<'db> DefinitionNodeRef<'_, 'db> { }), DefinitionNodeRef::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node, - submodule, }) => DefinitionKind::ImportFromSubmodule(ImportFromSubmoduleDefinitionKind { node: AstNodeRef::new(parsed, node), - submodule: submodule.as_str().into(), }), DefinitionNodeRef::ImportStar(star_import) => { let StarImportDefinitionNodeRef { node, symbol_id } = star_import; @@ -580,10 +576,7 @@ impl<'db> DefinitionNodeRef<'_, 'db> { alias_index, is_reexported: _, }) => (&node.names[alias_index]).into(), - Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { - node, - submodule: _, - }) => node.into(), + Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node }) => node.into(), // INVARIANT: for an invalid-syntax statement such as `from foo import *, bar, *`, // we only create a `StarImportDefinitionKind` for the *first* `*` alias in the names list. Self::ImportStar(StarImportDefinitionNodeRef { node, symbol_id: _ }) => node @@ -1021,17 +1014,12 @@ impl ImportFromDefinitionKind { #[derive(Clone, Debug, get_size2::GetSize)] pub struct ImportFromSubmoduleDefinitionKind { node: AstNodeRef, - submodule: Name, } impl ImportFromSubmoduleDefinitionKind { pub fn import<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::StmtImportFrom { self.node.node(module) } - - pub(crate) fn submodule(&self) -> &Name { - &self.submodule - } } #[derive(Clone, Debug, get_size2::GetSize)] diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index b4f7a42099..e775ca1993 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -4,7 +4,6 @@ use itertools::{Either, Itertools}; use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity}; use ruff_db::files::File; use ruff_db::parsed::ParsedModuleRef; -use ruff_python_ast::name::Name; use ruff_python_ast::visitor::{Visitor, walk_expr}; use ruff_python_ast::{ self as ast, AnyNodeRef, ExprContext, HasNodeIndex, NodeIndex, PythonVersion, @@ -1218,7 +1217,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { DefinitionKind::ImportFromSubmodule(import_from) => { self.infer_import_from_submodule_definition( import_from.import(self.module()), - import_from.submodule(), definition, ); } @@ -5901,51 +5899,64 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - /// Infer the implicit local definition `x = ` that - /// `from .x.y import z` can introduce in an `__init__.py(i)`. + /// Infer the implicit local definition `x = ` that + /// `from .x.y import z` or `from whatever.thispackage.x.y` can introduce in `__init__.py(i)`. /// /// For the definition `z`, see [`TypeInferenceBuilder::infer_import_from_definition`]. + /// + /// The runtime semantic of this kind of statement is to introduce a variable in the global + /// scope of this module *the first time it's imported in the entire program*. This + /// implementation just blindly introduces a local variable wherever the `from..import` is + /// (if the imports actually resolve). + /// + /// That gap between the semantics and implementation are currently the responsibility of the + /// code that actually creates these kinds of Definitions (so blindly introducing a local + /// is all we need to be doing here). fn infer_import_from_submodule_definition( &mut self, import_from: &ast::StmtImportFrom, - submodule: &Name, definition: Definition<'db>, ) { - // The runtime semantic of this kind of statement is to introduce a variable in the global - // scope of this module, so we do just that. (Actually we introduce a local variable, but - // this type of Definition is only created when a `from..import` is in global scope.) - - // Get this package's module by resolving `.` - let Ok(module_name) = ModuleName::from_identifier_parts(self.db(), self.file(), None, 1) - else { + // Get this package's absolute module name by resolving `.`, and make sure it exists + let Ok(thispackage_name) = ModuleName::package_for_file(self.db(), self.file()) else { + self.add_binding(import_from.into(), definition, |_, _| Type::unknown()); + return; + }; + let Some(module) = resolve_module(self.db(), &thispackage_name) else { self.add_binding(import_from.into(), definition, |_, _| Type::unknown()); return; }; - let Some(module) = resolve_module(self.db(), &module_name) else { + // We have `from whatever.thispackage.x.y ...` or `from .x.y ...` + // and we want to extract `x` (to ultimately construct `whatever.thispackage.x`): + + // First we normalize to `whatever.thispackage.x.y` + let Some(final_part) = ModuleName::from_identifier_parts( + self.db(), + self.file(), + import_from.module.as_deref(), + import_from.level, + ) + .ok() + // `whatever.thispackage.x.y` => `x.y` + .and_then(|submodule_name| submodule_name.relative_to(&thispackage_name)) + // `x.y` => `x` + .and_then(|relative_submodule_name| { + relative_submodule_name + .components() + .next() + .and_then(ModuleName::new) + }) else { self.add_binding(import_from.into(), definition, |_, _| Type::unknown()); return; }; - // Now construct the submodule `.x` - assert!( - !submodule.is_empty(), - "ImportFromSubmoduleDefinitionKind constructed with empty module" - ); - let name = submodule - .split_once('.') - .map(|(first, _)| first) - .unwrap_or(submodule.as_str()); - let full_submodule_name = ModuleName::new(name).map(|final_part| { - let mut ret = module_name.clone(); - ret.extend(&final_part); - ret - }); - // And try to import it - if let Some(submodule_type) = full_submodule_name - .as_ref() - .and_then(|submodule_name| self.module_type_from_name(submodule_name)) - { + // `x` => `whatever.thispackage.x` + let mut full_submodule_name = thispackage_name.clone(); + full_submodule_name.extend(&final_part); + + // Try to actually resolve the import `whatever.thispackage.x` + if let Some(submodule_type) = self.module_type_from_name(&full_submodule_name) { // Success, introduce a binding! // // We explicitly don't introduce a *declaration* because it's actual ok @@ -5970,17 +5981,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { }; let diagnostic = builder.into_diagnostic(format_args!( - "Module `{module_name}` has no submodule `{name}`" + "Module `{thispackage_name}` has no submodule `{final_part}`" )); - if let Some(full_submodule_name) = full_submodule_name { - hint_if_stdlib_submodule_exists_on_other_versions( - self.db(), - diagnostic, - &full_submodule_name, - module, - ); - } + hint_if_stdlib_submodule_exists_on_other_versions( + self.db(), + diagnostic, + &full_submodule_name, + module, + ); } fn infer_return_statement(&mut self, ret: &ast::StmtReturn) {