[ty] Implemented "go to definition" support for import statements

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<Type<'db>> {
+    pub(crate) fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Option<Type<'db>> {
         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<crate::NavigationTargets> {
+    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
+}