[ty] Implemented support for "rename" language server feature (#19551)

This PR adds support for the "rename" language server feature. It builds
upon existing functionality used for "go to references".

The "rename" feature involves two language server requests. The first is
a "prepare rename" request that determines whether renaming should be
possible for the identifier at the current offset. The second is a
"rename" request that returns a list of file ranges where the rename
should be applied.

Care must be taken when attempting to rename symbols that span files,
especially if the symbols are defined in files that are not part of the
project. We don't want to modify code in the user's Python environment
or in the vendored stub files.

I found a few bugs in the "go to references" feature when implementing
"rename", and those bug fixes are included in this PR.

---------

Co-authored-by: UnboundVariable <unbound@gmail.com>

Author: UnboundVariable

crates/ty_ide/src/goto.rs

@@ -10,6 +10,7 @@ use ruff_db::parsed::ParsedModuleRef;
 use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_python_parser::TokenKind;
 use ruff_text_size::{Ranged, TextRange, TextSize};
+use ty_python_semantic::ImportAliasResolution;
 use ty_python_semantic::types::Type;
 use ty_python_semantic::types::definitions_for_keyword_argument;
 use ty_python_semantic::{
@@ -172,12 +173,16 @@ impl GotoTarget<'_> {
 
     /// Gets the navigation ranges for this goto target.
     /// If a stub mapper is provided, definitions from stub files will be mapped to
-    /// their corresponding source file implementations.
+    /// their corresponding source file implementations. The `alias_resolution`
+    /// parameter controls whether import aliases (i.e. "x" in "from a import b as x") are
+    /// resolved or returned as is. We want to resolve them in some cases (like
+    /// "goto declaration") but not in others (like find references or rename).
     pub(crate) fn get_definition_targets(
         &self,
         file: ruff_db::files::File,
         db: &dyn crate::Db,
         stub_mapper: Option<&StubMapper>,
+        alias_resolution: ImportAliasResolution,
     ) -> Option<crate::NavigationTargets> {
         use crate::NavigationTarget;
         use ruff_python_ast as ast;
@@ -229,10 +234,14 @@ impl GotoTarget<'_> {
             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);
+                let definitions = definitions_for_imported_symbol(
+                    db,
+                    file,
+                    import_from,
+                    symbol_name,
+                    alias_resolution,
+                );
 
                 definitions_to_navigation_targets(db, stub_mapper, definitions)
             }
@@ -254,12 +263,18 @@ impl GotoTarget<'_> {
 
             // 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)
+                if alias_resolution == ImportAliasResolution::ResolveAliases {
+                    let full_module_name = alias.name.as_str();
+                    // Try to resolve the module
+                    resolve_module_to_navigation_target(db, full_module_name)
+                } else {
+                    let alias_range = alias.asname.as_ref().unwrap().range;
+                    Some(crate::NavigationTargets::single(NavigationTarget {
+                        file,
+                        focus_range: alias_range,
+                        full_range: alias.range(),
+                    }))
+                }
             }
 
             // Handle keyword arguments in call expressions

crates/ty_ide/src/goto_declaration.rs

@@ -3,6 +3,7 @@ use crate::{Db, NavigationTargets, RangedValue};
 use ruff_db::files::{File, FileRange};
 use ruff_db::parsed::parsed_module;
 use ruff_text_size::{Ranged, TextSize};
+use ty_python_semantic::ImportAliasResolution;
 
 /// Navigate to the declaration of a symbol.
 ///
@@ -17,7 +18,12 @@ pub fn goto_declaration(
     let module = parsed_module(db, file).load(db);
     let goto_target = find_goto_target(&module, offset)?;
 
-    let declaration_targets = goto_target.get_definition_targets(file, db, None)?;
+    let declaration_targets = goto_target.get_definition_targets(
+        file,
+        db,
+        None,
+        ImportAliasResolution::ResolveAliases,
+    )?;
 
     Some(RangedValue {
         range: FileRange::new(file, goto_target.range()),

crates/ty_ide/src/goto_definition.rs

@@ -4,6 +4,7 @@ use crate::{Db, NavigationTargets, RangedValue};
 use ruff_db::files::{File, FileRange};
 use ruff_db::parsed::parsed_module;
 use ruff_text_size::{Ranged, TextSize};
+use ty_python_semantic::ImportAliasResolution;
 
 /// Navigate to the definition of a symbol.
 ///
@@ -22,7 +23,12 @@ pub fn goto_definition(
     // Create a StubMapper to map from stub files to source files
     let stub_mapper = StubMapper::new(db);
 
-    let definition_targets = goto_target.get_definition_targets(file, db, Some(&stub_mapper))?;
+    let definition_targets = goto_target.get_definition_targets(
+        file,
+        db,
+        Some(&stub_mapper),
+        ImportAliasResolution::ResolveAliases,
+    )?;
 
     Some(RangedValue {
         range: FileRange::new(file, goto_target.range()),

crates/ty_ide/src/goto_references.rs

@@ -688,34 +688,30 @@ cls = MyClass
             .source(
                 "utils.py",
                 "
-def helper_fun<CURSOR>ction(x):
+def fun<CURSOR>c(x):
     return x * 2
 ",
             )
             .source(
                 "module.py",
                 "
-from utils import helper_function
+from utils import func
 
 def process_data(data):
-    return helper_function(data)
-
-def double_process(data):
-    result = helper_function(data)
-    return helper_function(result)
+    return func(data)
 ",
             )
             .source(
                 "app.py",
                 "
-from utils import helper_function
+from utils import func
 
 class DataProcessor:
     def __init__(self):
-        self.multiplier = helper_function
+        self.multiplier = func
     
     def process(self, value):
-        return helper_function(value)
+        return func(value)
 ",
             )
             .build();
@@ -724,46 +720,44 @@ class DataProcessor:
         info[references]: Reference 1
          --> utils.py:2:5
           |
-        2 | def helper_function(x):
-          |     ^^^^^^^^^^^^^^^
+        2 | def func(x):
+          |     ^^^^
         3 |     return x * 2
           |
 
         info[references]: Reference 2
-         --> module.py:5:12
+         --> module.py:2:19
           |
+        2 | from utils import func
+          |                   ^^^^
+        3 |
         4 | def process_data(data):
-        5 |     return helper_function(data)
-          |            ^^^^^^^^^^^^^^^
-        6 |
-        7 | def double_process(data):
           |
 
         info[references]: Reference 3
-         --> module.py:8:14
+         --> module.py:5:12
           |
-        7 | def double_process(data):
-        8 |     result = helper_function(data)
-          |              ^^^^^^^^^^^^^^^
-        9 |     return helper_function(result)
+        4 | def process_data(data):
+        5 |     return func(data)
+          |            ^^^^
           |
 
         info[references]: Reference 4
-         --> module.py:9:12
+         --> app.py:2:19
           |
-        7 | def double_process(data):
-        8 |     result = helper_function(data)
-        9 |     return helper_function(result)
-          |            ^^^^^^^^^^^^^^^
+        2 | from utils import func
+          |                   ^^^^
+        3 |
+        4 | class DataProcessor:
           |
 
         info[references]: Reference 5
          --> app.py:6:27
           |
         4 | class DataProcessor:
         5 |     def __init__(self):
-        6 |         self.multiplier = helper_function
-          |                           ^^^^^^^^^^^^^^^
+        6 |         self.multiplier = func
+          |                           ^^^^
         7 |     
         8 |     def process(self, value):
           |
@@ -772,8 +766,8 @@ class DataProcessor:
          --> app.py:9:16
           |
         8 |     def process(self, value):
-        9 |         return helper_function(value)
-          |                ^^^^^^^^^^^^^^^
+        9 |         return func(value)
+          |                ^^^^
           |
         ");
     }
@@ -855,4 +849,49 @@ def process_model():
           |
         ");
     }
+
+    #[test]
+    fn test_import_alias_references_should_not_resolve_to_original() {
+        let test = CursorTest::builder()
+            .source(
+                "original.py",
+                "
+def func():
+    pass
+
+func()
+",
+            )
+            .source(
+                "importer.py",
+                "
+from original import func as func_alias
+
+func<CURSOR>_alias()
+",
+            )
+            .build();
+
+        // When finding references to the alias, we should NOT find references
+        // to the original function in the original module
+        assert_snapshot!(test.references(), @r"
+        info[references]: Reference 1
+         --> importer.py:2:30
+          |
+        2 | from original import func as func_alias
+          |                              ^^^^^^^^^^
+        3 |
+        4 | func_alias()
+          |
+
+        info[references]: Reference 2
+         --> importer.py:4:1
+          |
+        2 | from original import func as func_alias
+        3 |
+        4 | func_alias()
+          | ^^^^^^^^^^
+          |
+        ");
+    }
 }

crates/ty_ide/src/lib.rs

@@ -12,6 +12,7 @@ mod hover;
 mod inlay_hints;
 mod markup;
 mod references;
+mod rename;
 mod selection_range;
 mod semantic_tokens;
 mod signature_help;
@@ -29,6 +30,7 @@ pub use hover::hover;
 pub use inlay_hints::inlay_hints;
 pub use markup::MarkupKind;
 pub use references::ReferencesMode;
+pub use rename::{can_rename, rename};
 pub use selection_range::selection_range;
 pub use semantic_tokens::{
     SemanticToken, SemanticTokenModifier, SemanticTokenType, SemanticTokens, semantic_tokens,