[ty] Implemented "go to definition" support for import statements (#19428)

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 <unbound@gmail.com>

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
+}

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 su<CURSOR>b
+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.submod<CURSOR>ule 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 hel<CURSOR>per 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<CURSOR>
+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.ut<CURSOR>ils 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(

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;

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};

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<ResolvedDefinition<'db>> {
+    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,