diff --git a/crates/ty_ide/src/goto_declaration.rs b/crates/ty_ide/src/goto_declaration.rs index a2f147b2d3..e7fdb54be4 100644 --- a/crates/ty_ide/src/goto_declaration.rs +++ b/crates/ty_ide/src/goto_declaration.rs @@ -2622,16 +2622,12 @@ def ab(a: int, *, c: int): ... ) .build(); - // TODO(submodule-imports): this should only highlight `subpkg` in the import statement - // This happens because DefinitionKind::ImportFromSubmodule claims the entire ImportFrom node, - // which is correct but unhelpful. Unfortunately even if it only claimed the LHS identifier it - // would highlight `subpkg.submod` which is strictly better but still isn't what we want. assert_snapshot!(test.goto_declaration(), @r" info[goto-declaration]: Declaration - --> mypackage/__init__.py:2:1 + --> mypackage/__init__.py:2:7 | 2 | from .subpkg.submod import val - | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | ^^^^^^ 3 | 4 | x = subpkg | @@ -2846,25 +2842,26 @@ def ab(a: int, *, c: int): ... ) .build(); - // TODO(submodule-imports): Ok this one is FASCINATING and it's kinda right but confusing! + // TODO(submodule-imports): Ok this one is FASCINATING but definitely wrong! // // So there's 3 relevant definitions here: // - // * `subpkg: int = 10` in the other file is in fact the original definition + // * `subpkg: int = 10` in the other file is in fact the original definition. + // Including it here is accurate and possibly useful? // // * the LHS `subpkg` in the import is an instance of `subpkg = ...` // because it's a `DefinitionKind::ImportFromSubmodle`. - // This is the span that covers the entire import. + // Including it here is Pedantically Correct but Unhelpful. // // * `the RHS `subpkg` in the import is a second instance of `subpkg = ...` - // that *immediately* overwrites the `ImportFromSubmodule`'s definition - // This span seemingly doesn't appear at all!? Is it getting hidden by the LHS span? + // that *immediately* overwrites the `ImportFromSubmodule`'s definition. + // This is the most important one and doesn't show up at all! Sadness! assert_snapshot!(test.goto_declaration(), @r" info[goto-declaration]: Declaration - --> mypackage/__init__.py:2:1 + --> mypackage/__init__.py:2:7 | 2 | from .subpkg import subpkg - | ^^^^^^^^^^^^^^^^^^^^^^^^^^ + | ^^^^^^ 3 | 4 | x = subpkg | diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 373c80f7ee..0e1f9cfabc 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -1529,7 +1529,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { // Record whether this is equivalent to `from . import ...` is_self_import = module_name == thispackage; - if node.module.is_some() + if let Some(module_node) = &node.module && let Some(relative_submodule) = module_name.relative_to(&thispackage) && let Some(direct_submodule) = relative_submodule.components().next() && !self.seen_submodule_imports.contains(direct_submodule) @@ -1540,9 +1540,23 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { let direct_submodule_name = Name::new(direct_submodule); let symbol = self.add_symbol(direct_submodule_name); + + let module_index = if node.level == 0 { + // "whatever.thispackage.x.y" we want `x` + thispackage.components().count() + } else { + // ".x.y" we want `x` (level 1 => index 0) + // "..x.y" we want `y` (level 2 => index 1) + // (The Identifier doesn't include the prefix dots) + node.level as usize - 1 + }; self.add_definition( symbol.into(), - ImportFromSubmoduleDefinitionNodeRef { node }, + ImportFromSubmoduleDefinitionNodeRef { + node, + module: module_node, + module_index, + }, ); } } diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index 2659e75493..c0e4c8e4d7 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -3,7 +3,7 @@ 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_text_size::{Ranged, TextRange}; +use ruff_text_size::{Ranged, TextLen, TextRange, TextSize}; use crate::Db; use crate::ast_node_ref::AstNodeRef; @@ -374,6 +374,8 @@ pub(crate) struct ImportFromDefinitionNodeRef<'ast> { #[derive(Copy, Clone, Debug)] pub(crate) struct ImportFromSubmoduleDefinitionNodeRef<'ast> { pub(crate) node: &'ast ast::StmtImportFrom, + pub(crate) module: &'ast ast::Identifier, + pub(crate) module_index: usize, } #[derive(Copy, Clone, Debug)] @@ -456,8 +458,12 @@ impl<'db> DefinitionNodeRef<'_, 'db> { }), DefinitionNodeRef::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node, + module, + module_index, }) => DefinitionKind::ImportFromSubmodule(ImportFromSubmoduleDefinitionKind { node: AstNodeRef::new(parsed, node), + module: AstNodeRef::new(parsed, module), + module_index, }), DefinitionNodeRef::ImportStar(star_import) => { let StarImportDefinitionNodeRef { node, symbol_id } = star_import; @@ -584,7 +590,9 @@ impl<'db> DefinitionNodeRef<'_, 'db> { alias_index, is_reexported: _, }) => (&node.names[alias_index]).into(), - Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef { node }) => 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 @@ -755,7 +763,7 @@ impl DefinitionKind<'_> { match self { DefinitionKind::Import(import) => import.alias(module).range(), DefinitionKind::ImportFrom(import) => import.alias(module).range(), - DefinitionKind::ImportFromSubmodule(import) => import.import(module).range(), + DefinitionKind::ImportFromSubmodule(import) => import.target_range(module), DefinitionKind::StarImport(import) => import.alias(module).range(), DefinitionKind::Function(function) => function.node(module).name.range(), DefinitionKind::Class(class) => class.node(module).name.range(), @@ -793,7 +801,7 @@ impl DefinitionKind<'_> { match self { DefinitionKind::Import(import) => import.alias(module).range(), DefinitionKind::ImportFrom(import) => import.alias(module).range(), - DefinitionKind::ImportFromSubmodule(import) => import.import(module).range(), + DefinitionKind::ImportFromSubmodule(import) => import.module(module).range(), DefinitionKind::StarImport(import) => import.import(module).range(), DefinitionKind::Function(function) => function.node(module).range(), DefinitionKind::Class(class) => class.node(module).range(), @@ -1033,12 +1041,37 @@ impl ImportFromDefinitionKind { #[derive(Clone, Debug, get_size2::GetSize)] pub struct ImportFromSubmoduleDefinitionKind { node: AstNodeRef, + module: AstNodeRef, + module_index: usize, } impl ImportFromSubmoduleDefinitionKind { pub fn import<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::StmtImportFrom { self.node.node(module) } + + pub fn module<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::Identifier { + self.module.node(module) + } + + pub fn target_range(&self, module: &ParsedModuleRef) -> TextRange { + let module_ident = self.module(module); + let module_str = module_ident.as_str(); + let Some(component_str) = module_str.split('.').nth(self.module_index) else { + // This shouldn't happen but just in case, provide a safe default + return module_ident.range(); + }; + let base_addr = module_str.as_ptr().addr(); + let component_addr = component_str.as_ptr().addr(); + let offset = base_addr.saturating_sub(component_addr); + let Ok(offset_size) = TextSize::try_from(offset) else { + // This shouldn't happen but just in case, provide a safe default + return module_ident.range(); + }; + let start = module_ident.start().saturating_add(offset_size); + let end = start.saturating_add(component_str.text_len()); + TextRange::new(start, end) + } } #[derive(Clone, Debug, get_size2::GetSize)]