From 0acc273286aaac6fc9a65e4acf42f3f2ebbf9e1f Mon Sep 17 00:00:00 2001 From: UnboundVariable Date: Sat, 19 Jul 2025 11:22:07 -0700 Subject: [PATCH] [ty] Implemented "go to definition" support for import statements (#19428) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR extends the "go to declaration" and "go to definition" functionality to support import statements — both standard imports and "from" import forms. --------- Co-authored-by: UnboundVariable --- crates/ty_ide/src/goto.rs | 247 ++++++++++++++++-- crates/ty_ide/src/goto_declaration.rs | 223 ++++++++++++++++ crates/ty_python_semantic/src/lib.rs | 6 +- crates/ty_python_semantic/src/types.rs | 3 +- .../src/types/ide_support.rs | 21 +- 5 files changed, 469 insertions(+), 31 deletions(-) diff --git a/crates/ty_ide/src/goto.rs b/crates/ty_ide/src/goto.rs index c7b235b1c2..134e7dd140 100644 --- a/crates/ty_ide/src/goto.rs +++ b/crates/ty_ide/src/goto.rs @@ -10,22 +10,52 @@ use ruff_python_parser::TokenKind; use ruff_text_size::{Ranged, TextRange, TextSize}; use ty_python_semantic::types::Type; use ty_python_semantic::types::definitions_for_keyword_argument; -use ty_python_semantic::{HasType, SemanticModel, definitions_for_name}; +use ty_python_semantic::{ + HasType, SemanticModel, definitions_for_imported_symbol, definitions_for_name, +}; -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Debug)] pub(crate) enum GotoTarget<'a> { Expression(ast::ExprRef<'a>), FunctionDef(&'a ast::StmtFunctionDef), ClassDef(&'a ast::StmtClassDef), Parameter(&'a ast::Parameter), - Alias(&'a ast::Alias), - /// Go to on the module name of an import from + /// Multi-part module names + /// Handles both `import foo.bar` and `from foo.bar import baz` cases /// ```py - /// from foo import bar - /// ^^^ + /// import foo.bar + /// ^^^ + /// from foo.bar import baz + /// ^^^ /// ``` - ImportedModule(&'a ast::StmtImportFrom), + ImportModuleComponent { + module_name: String, + component_index: usize, + component_range: TextRange, + }, + + /// Import alias in standard import statement + /// ```py + /// import foo.bar as baz + /// ^^^ + /// ``` + ImportModuleAlias { + alias: &'a ast::Alias, + }, + + /// Import alias in from import statement + /// ```py + /// from foo import bar as baz + /// ^^^ + /// from foo import bar as baz + /// ^^^ + /// ``` + ImportSymbolAlias { + alias: &'a ast::Alias, + range: TextRange, + import_from: &'a ast::StmtImportFrom, + }, /// Go to on the exception handler variable /// ```py @@ -112,25 +142,22 @@ pub(crate) enum GotoTarget<'a> { } impl GotoTarget<'_> { - pub(crate) fn inferred_type<'db>(self, model: &SemanticModel<'db>) -> Option> { + pub(crate) fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Option> { let ty = match self { GotoTarget::Expression(expression) => expression.inferred_type(model), GotoTarget::FunctionDef(function) => function.inferred_type(model), GotoTarget::ClassDef(class) => class.inferred_type(model), GotoTarget::Parameter(parameter) => parameter.inferred_type(model), - GotoTarget::Alias(alias) => alias.inferred_type(model), + GotoTarget::ImportSymbolAlias { alias, .. } => alias.inferred_type(model), + GotoTarget::ImportModuleAlias { alias } => alias.inferred_type(model), GotoTarget::ExceptVariable(except) => except.inferred_type(model), - GotoTarget::KeywordArgument { keyword, .. } => { - // TODO: Pyright resolves the declared type of the matching parameter. This seems more accurate - // than using the inferred value. - keyword.value.inferred_type(model) - } + GotoTarget::KeywordArgument { keyword, .. } => keyword.value.inferred_type(model), // TODO: Support identifier targets GotoTarget::PatternMatchRest(_) | GotoTarget::PatternKeywordArgument(_) | GotoTarget::PatternMatchStarName(_) | GotoTarget::PatternMatchAsName(_) - | GotoTarget::ImportedModule(_) + | GotoTarget::ImportModuleComponent { .. } | GotoTarget::TypeParamTypeVarName(_) | GotoTarget::TypeParamParamSpecName(_) | GotoTarget::TypeParamTypeVarTupleName(_) @@ -145,7 +172,7 @@ impl GotoTarget<'_> { /// If a stub mapper is provided, definitions from stub files will be mapped to /// their corresponding source file implementations. pub(crate) fn get_definition_targets( - self, + &self, file: ruff_db::files::File, db: &dyn crate::Db, stub_mapper: Option<&StubMapper>, @@ -196,11 +223,41 @@ impl GotoTarget<'_> { })) } - // For imports, find the symbol being imported - GotoTarget::Alias(_alias) => { - // For aliases, we don't have the ExprName node, so we can't get the scope - // For now, return None. In the future, we could look up the imported symbol - None + // For import aliases (offset within 'y' or 'z' in "from x import y as z") + GotoTarget::ImportSymbolAlias { + alias, import_from, .. + } => { + // Handle both original names and alias names in `from x import y as z` statements + let symbol_name = alias.name.as_str(); + let definitions = + definitions_for_imported_symbol(db, file, import_from, symbol_name); + + definitions_to_navigation_targets(db, stub_mapper, definitions) + } + + GotoTarget::ImportModuleComponent { + module_name, + component_index, + .. + } => { + // Handle both `import foo.bar` and `from foo.bar import baz` where offset is within module component + let components: Vec<&str> = module_name.split('.').collect(); + + // Build the module name up to and including the component containing the offset + let target_module_name = components[..=*component_index].join("."); + + // Try to resolve the module + resolve_module_to_navigation_target(db, &target_module_name) + } + + // Handle import aliases (offset within 'z' in "import x.y as z") + GotoTarget::ImportModuleAlias { alias } => { + // For import aliases, navigate to the module being aliased + // This only applies to regular import statements like "import x.y as z" + let full_module_name = alias.name.as_str(); + + // Try to resolve the module + resolve_module_to_navigation_target(db, full_module_name) } // Handle keyword arguments in call expressions @@ -213,8 +270,6 @@ impl GotoTarget<'_> { definitions_to_navigation_targets(db, stub_mapper, definitions) } - // TODO: Handle multi-part module names in import statements - // TODO: Handle imported symbol in y in `from x import y as z` statement // TODO: Handle string literals that map to TypedDict fields _ => None, } @@ -228,8 +283,11 @@ impl Ranged for GotoTarget<'_> { GotoTarget::FunctionDef(function) => function.name.range, GotoTarget::ClassDef(class) => class.name.range, GotoTarget::Parameter(parameter) => parameter.name.range, - GotoTarget::Alias(alias) => alias.name.range, - GotoTarget::ImportedModule(module) => module.module.as_ref().unwrap().range, + GotoTarget::ImportSymbolAlias { range, .. } => *range, + GotoTarget::ImportModuleComponent { + component_range, .. + } => *component_range, + GotoTarget::ImportModuleAlias { alias } => alias.asname.as_ref().unwrap().range, GotoTarget::ExceptVariable(except) => except.name.as_ref().unwrap().range, GotoTarget::KeywordArgument { keyword, .. } => keyword.arg.as_ref().unwrap().range, GotoTarget::PatternMatchRest(rest) => rest.rest.as_ref().unwrap().range, @@ -324,8 +382,89 @@ pub(crate) fn find_goto_target( Some(AnyNodeRef::StmtFunctionDef(function)) => Some(GotoTarget::FunctionDef(function)), Some(AnyNodeRef::StmtClassDef(class)) => Some(GotoTarget::ClassDef(class)), Some(AnyNodeRef::Parameter(parameter)) => Some(GotoTarget::Parameter(parameter)), - Some(AnyNodeRef::Alias(alias)) => Some(GotoTarget::Alias(alias)), - Some(AnyNodeRef::StmtImportFrom(from)) => Some(GotoTarget::ImportedModule(from)), + Some(AnyNodeRef::Alias(alias)) => { + // Find the containing import statement to determine the type + let import_stmt = covering_node.ancestors().find(|node| { + matches!( + node, + AnyNodeRef::StmtImport(_) | AnyNodeRef::StmtImportFrom(_) + ) + }); + + match import_stmt { + Some(AnyNodeRef::StmtImport(_)) => { + // Regular import statement like "import x.y as z" + + // Is the offset within the alias name (asname) part? + if let Some(asname) = &alias.asname { + if asname.range.contains_inclusive(offset) { + return Some(GotoTarget::ImportModuleAlias { alias }); + } + } + + // Is the offset in the module name part? + if alias.name.range.contains_inclusive(offset) { + let full_name = alias.name.as_str(); + + if let Some((component_index, component_range)) = + find_module_component(full_name, alias.name.range.start(), offset) + { + return Some(GotoTarget::ImportModuleComponent { + module_name: full_name.to_string(), + component_index, + component_range, + }); + } + } + + None + } + Some(AnyNodeRef::StmtImportFrom(import_from)) => { + // From import statement like "from x import y as z" + + // Is the offset within the alias name (asname) part? + if let Some(asname) = &alias.asname { + if asname.range.contains_inclusive(offset) { + return Some(GotoTarget::ImportSymbolAlias { + alias, + range: asname.range, + import_from, + }); + } + } + + // Is the offset in the original name part? + if alias.name.range.contains_inclusive(offset) { + return Some(GotoTarget::ImportSymbolAlias { + alias, + range: alias.name.range, + import_from, + }); + } + + None + } + _ => None, + } + } + Some(AnyNodeRef::StmtImportFrom(from)) => { + // Handle offset within module name in from import statements + if let Some(module_expr) = &from.module { + let full_module_name = module_expr.to_string(); + + if let Some((component_index, component_range)) = + find_module_component(&full_module_name, module_expr.range.start(), offset) + { + return Some(GotoTarget::ImportModuleComponent { + module_name: full_module_name, + component_index, + component_range, + }); + } + } + + None + } Some(AnyNodeRef::ExceptHandlerExceptHandler(handler)) => { Some(GotoTarget::ExceptVariable(handler)) } @@ -376,3 +515,57 @@ pub(crate) fn find_goto_target( node => node.as_expr_ref().map(GotoTarget::Expression), } } + +/// Helper function to resolve a module name and create a navigation target. +fn resolve_module_to_navigation_target( + db: &dyn crate::Db, + module_name_str: &str, +) -> Option { + use ty_python_semantic::{ModuleName, resolve_module}; + + if let Some(module_name) = ModuleName::new(module_name_str) { + if let Some(resolved_module) = resolve_module(db, &module_name) { + if let Some(module_file) = resolved_module.file() { + return Some(crate::NavigationTargets::single(crate::NavigationTarget { + file: module_file, + focus_range: TextRange::default(), + full_range: TextRange::default(), + })); + } + } + } + None +} + +/// Helper function to extract module component information from a dotted module name +fn find_module_component( + full_module_name: &str, + module_start: TextSize, + offset: TextSize, +) -> Option<(usize, TextRange)> { + let pos_in_module = offset - module_start; + let pos_in_module = pos_in_module.to_usize(); + + // Split the module name into components and find which one contains the offset + let mut current_pos = 0; + let components: Vec<&str> = full_module_name.split('.').collect(); + + for (i, component) in components.iter().enumerate() { + let component_start = current_pos; + let component_end = current_pos + component.len(); + + // Check if the offset is within this component or at its right boundary + if pos_in_module >= component_start && pos_in_module <= component_end { + let component_range = TextRange::new( + module_start + TextSize::from(u32::try_from(component_start).ok()?), + module_start + TextSize::from(u32::try_from(component_end).ok()?), + ); + return Some((i, component_range)); + } + + // Move past this component and the dot + current_pos = component_end + 1; // +1 for the dot + } + + None +} diff --git a/crates/ty_ide/src/goto_declaration.rs b/crates/ty_ide/src/goto_declaration.rs index e403340eb4..5f5f358153 100644 --- a/crates/ty_ide/src/goto_declaration.rs +++ b/crates/ty_ide/src/goto_declaration.rs @@ -610,6 +610,229 @@ def another_helper(): "#); } + #[test] + fn goto_declaration_import_as_alias_name() { + let test = CursorTest::builder() + .source( + "main.py", + " +import mymodule.submodule as sub +print(sub.helper()) +", + ) + .source( + "mymodule/__init__.py", + " +# Main module init +", + ) + .source( + "mymodule/submodule.py", + r#" +FOO = 0 +"#, + ) + .build(); + + assert_snapshot!(test.goto_declaration(), @r" + info[goto-declaration]: Declaration + --> mymodule/submodule.py:1:1 + | + 1 | + | ^ + 2 | FOO = 0 + | + info: Source + --> main.py:2:30 + | + 2 | import mymodule.submodule as sub + | ^^^ + 3 | print(sub.helper()) + | + "); + } + + #[test] + fn goto_declaration_import_as_alias_name_on_module() { + let test = CursorTest::builder() + .source( + "main.py", + " +import mymodule.submodule as sub +print(sub.helper()) +", + ) + .source( + "mymodule/__init__.py", + " +# Main module init +", + ) + .source( + "mymodule/submodule.py", + r#" +FOO = 0 +"#, + ) + .build(); + + assert_snapshot!(test.goto_declaration(), @r" + info[goto-declaration]: Declaration + --> mymodule/submodule.py:1:1 + | + 1 | + | ^ + 2 | FOO = 0 + | + info: Source + --> main.py:2:17 + | + 2 | import mymodule.submodule as sub + | ^^^^^^^^^ + 3 | print(sub.helper()) + | + "); + } + + #[test] + fn goto_declaration_from_import_symbol_original() { + let test = CursorTest::builder() + .source( + "main.py", + r#" +from mypackage.utils import helper as h +result = h("/a", "/b") +"#, + ) + .source( + "mypackage/__init__.py", + r#" +# Package init +"#, + ) + .source( + "mypackage/utils.py", + r#" +def helper(a, b): + return a + "/" + b + +def another_helper(path): + return "processed" +"#, + ) + .build(); + + assert_snapshot!(test.goto_declaration(), @r#" + info[goto-declaration]: Declaration + --> mypackage/utils.py:2:5 + | + 2 | def helper(a, b): + | ^^^^^^ + 3 | return a + "/" + b + | + info: Source + --> main.py:2:29 + | + 2 | from mypackage.utils import helper as h + | ^^^^^^ + 3 | result = h("/a", "/b") + | + "#); + } + + #[test] + fn goto_declaration_from_import_symbol_alias() { + let test = CursorTest::builder() + .source( + "main.py", + r#" +from mypackage.utils import helper as h +result = h("/a", "/b") +"#, + ) + .source( + "mypackage/__init__.py", + r#" +# Package init +"#, + ) + .source( + "mypackage/utils.py", + r#" +def helper(a, b): + return a + "/" + b + +def another_helper(path): + return "processed" +"#, + ) + .build(); + + assert_snapshot!(test.goto_declaration(), @r#" + info[goto-declaration]: Declaration + --> mypackage/utils.py:2:5 + | + 2 | def helper(a, b): + | ^^^^^^ + 3 | return a + "/" + b + | + info: Source + --> main.py:2:39 + | + 2 | from mypackage.utils import helper as h + | ^ + 3 | result = h("/a", "/b") + | + "#); + } + + #[test] + fn goto_declaration_from_import_module() { + let test = CursorTest::builder() + .source( + "main.py", + r#" +from mypackage.utils import helper as h +result = h("/a", "/b") +"#, + ) + .source( + "mypackage/__init__.py", + r#" +# Package init +"#, + ) + .source( + "mypackage/utils.py", + r#" +def helper(a, b): + return a + "/" + b + +def another_helper(path): + return "processed" +"#, + ) + .build(); + + assert_snapshot!(test.goto_declaration(), @r#" + info[goto-declaration]: Declaration + --> mypackage/utils.py:1:1 + | + 1 | + | ^ + 2 | def helper(a, b): + 3 | return a + "/" + b + | + info: Source + --> main.py:2:16 + | + 2 | from mypackage.utils import helper as h + | ^^^^^ + 3 | result = h("/a", "/b") + | + "#); + } + #[test] fn goto_declaration_instance_attribute() { let test = cursor_test( diff --git a/crates/ty_python_semantic/src/lib.rs b/crates/ty_python_semantic/src/lib.rs index 06d8d6217c..11a842cb26 100644 --- a/crates/ty_python_semantic/src/lib.rs +++ b/crates/ty_python_semantic/src/lib.rs @@ -17,8 +17,10 @@ pub use program::{ pub use python_platform::PythonPlatform; pub use semantic_model::{Completion, CompletionKind, HasType, NameKind, SemanticModel}; pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin}; -pub use types::ide_support::ResolvedDefinition; -pub use types::{definitions_for_attribute, definitions_for_name}; +pub use types::ide_support::{ + ResolvedDefinition, definitions_for_attribute, definitions_for_imported_symbol, + definitions_for_name, +}; pub use util::diagnostics::add_inferred_python_version_hint_to_diagnostic; pub mod ast_node_ref; diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 8e57c04598..f12bb63ab2 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -49,7 +49,8 @@ use crate::types::generics::{ }; pub use crate::types::ide_support::{ CallSignatureDetails, Member, all_members, call_signature_details, definition_kind_for_name, - definitions_for_attribute, definitions_for_keyword_argument, definitions_for_name, + definitions_for_attribute, definitions_for_imported_symbol, definitions_for_keyword_argument, + definitions_for_name, }; use crate::types::infer::infer_unpack_types; use crate::types::mro::{Mro, MroError, MroIterator}; diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index b34064f603..b0eb2dfb93 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -682,6 +682,25 @@ pub fn definitions_for_keyword_argument<'db>( resolved_definitions } +/// Find the definitions for a symbol imported via `from x import y as z` statement. +/// This function handles the case where the cursor is on the original symbol name `y`. +/// Returns the same definitions as would be found for the alias `z`. +pub fn definitions_for_imported_symbol<'db>( + db: &'db dyn Db, + file: File, + import_node: &ast::StmtImportFrom, + symbol_name: &str, +) -> Vec> { + let mut visited = FxHashSet::default(); + resolve_definition::resolve_from_import_definitions( + db, + file, + import_node, + symbol_name, + &mut visited, + ) +} + /// Details about a callable signature for IDE support. #[derive(Debug, Clone)] pub struct CallSignatureDetails<'db> { @@ -888,7 +907,7 @@ mod resolve_definition { } /// Helper function to resolve import definitions for `ImportFrom` and `StarImport` cases. - fn resolve_from_import_definitions<'db>( + pub(crate) fn resolve_from_import_definitions<'db>( db: &'db dyn Db, file: File, import_node: &ast::StmtImportFrom,