[ty] Fix find-references for import aliases

crates/ty_ide/src/find_references.rs

@@ -1722,4 +1722,152 @@ func<CURSOR>_alias()
           |
         "###);
     }
+
+    #[test]
+    fn import_alias() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                import warnings
+                import warnings as <CURSOR>abc
+
+                x = abc
+                y = warnings
+            "#,
+            )
+            .build();
+
+        assert_snapshot!(test.references(), @r"
+        info[references]: Reference 1
+         --> main.py:3:20
+          |
+        2 | import warnings
+        3 | import warnings as abc
+          |                    ^^^
+        4 |
+        5 | x = abc
+          |
+
+        info[references]: Reference 2
+         --> main.py:5:5
+          |
+        3 | import warnings as abc
+        4 |
+        5 | x = abc
+          |     ^^^
+        6 | y = warnings
+          |
+        ");
+    }
+
+    #[test]
+    fn import_alias_use() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                import warnings
+                import warnings as abc
+
+                x = abc<CURSOR>
+                y = warnings
+            "#,
+            )
+            .build();
+
+        assert_snapshot!(test.references(), @r"
+        info[references]: Reference 1
+         --> main.py:3:20
+          |
+        2 | import warnings
+        3 | import warnings as abc
+          |                    ^^^
+        4 |
+        5 | x = abc
+          |
+
+        info[references]: Reference 2
+         --> main.py:5:5
+          |
+        3 | import warnings as abc
+        4 |
+        5 | x = abc
+          |     ^^^
+        6 | y = warnings
+          |
+        ");
+    }
+
+    #[test]
+    fn import_from_alias() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                from warnings import deprecated as xyz<CURSOR>
+                from warnings import deprecated
+
+                y = xyz
+                z = deprecated
+            "#,
+            )
+            .build();
+
+        assert_snapshot!(test.references(), @r"
+        info[references]: Reference 1
+         --> main.py:2:36
+          |
+        2 | from warnings import deprecated as xyz
+          |                                    ^^^
+        3 | from warnings import deprecated
+          |
+
+        info[references]: Reference 2
+         --> main.py:5:5
+          |
+        3 | from warnings import deprecated
+        4 |
+        5 | y = xyz
+          |     ^^^
+        6 | z = deprecated
+          |
+        ");
+    }
+
+    #[test]
+    fn import_from_alias_use() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                from warnings import deprecated as xyz
+                from warnings import deprecated
+
+                y = xyz<CURSOR>
+                z = deprecated
+            "#,
+            )
+            .build();
+
+        assert_snapshot!(test.references(), @r"
+        info[references]: Reference 1
+         --> main.py:2:36
+          |
+        2 | from warnings import deprecated as xyz
+          |                                    ^^^
+        3 | from warnings import deprecated
+          |
+
+        info[references]: Reference 2
+         --> main.py:5:5
+          |
+        3 | from warnings import deprecated
+        4 |
+        5 | y = xyz
+          |     ^^^
+        6 | z = deprecated
+          |
+        ");
+    }
 }

crates/ty_ide/src/goto.rs

@@ -396,13 +396,19 @@ impl GotoTarget<'_> {
             GotoTarget::ImportSymbolAlias {
                 alias, import_from, ..
             } => {
-                let symbol_name = alias.name.as_str();
-                Some(definitions_for_imported_symbol(
-                    model,
-                    import_from,
-                    symbol_name,
-                    alias_resolution,
-                ))
+                if let Some(asname) = alias.asname.as_ref()
+                    && alias_resolution == ImportAliasResolution::PreserveAliases
+                {
+                    Some(definitions_for_name(model, asname.as_str(), asname.into()))
+                } else {
+                    let symbol_name = alias.name.as_str();
+                    Some(definitions_for_imported_symbol(
+                        model,
+                        import_from,
+                        symbol_name,
+                        alias_resolution,
+                    ))
+                }
             }
 
             GotoTarget::ImportModuleComponent {
@@ -418,12 +424,12 @@ impl GotoTarget<'_> {
 
             // Handle import aliases (offset within 'z' in "import x.y as z")
             GotoTarget::ImportModuleAlias { alias } => {
-                if alias_resolution == ImportAliasResolution::ResolveAliases {
-                    definitions_for_module(model, Some(alias.name.as_str()), 0)
+                if let Some(asname) = alias.asname.as_ref()
+                    && alias_resolution == ImportAliasResolution::PreserveAliases
+                {
+                    Some(definitions_for_name(model, asname.as_str(), asname.into()))
                 } else {
-                    alias.asname.as_ref().map(|name| {
-                        definitions_for_name(model, name.as_str(), AnyNodeRef::Identifier(name))
-                    })
+                    definitions_for_module(model, Some(alias.name.as_str()), 0)
                 }
             }