[ty] Automatically update import statements and references when a module is renamed

crates/ty_server/src/capabilities.rs

@@ -501,6 +501,29 @@ pub(crate) fn server_capabilities(
                 supported: Some(true),
                 change_notifications: Some(true.into()),
             }),
+            file_operations: Some(lsp_types::FileOperationOptions {
+                will_rename: Some(lsp_types::FileOperationRegistrationOptions {
+                    filters: vec![
+                        lsp_types::FileOperationFilter {
+                            scheme: Some("file".to_string()),
+                            pattern: lsp_types::FileOperationPattern {
+                                glob: "**/*.{py,pyi}".to_string(),
+                                matches: Some(lsp_types::FileOperationPatternKind::File),
+                                options: None,
+                            },
+                        },
+                        lsp_types::FileOperationFilter {
+                            scheme: Some("file".to_string()),
+                            pattern: lsp_types::FileOperationPattern {
+                                glob: "**".to_string(),
+                                matches: Some(lsp_types::FileOperationPatternKind::Folder),
+                                options: None,
+                            },
+                        },
+                    ],
+                }),
+                ..Default::default()
+            }),
             ..Default::default()
         }),
         type_hierarchy_provider: Some(true.into()),

crates/ty_server/src/server/api.rs

@@ -108,6 +108,9 @@ pub(super) fn request(req: server::Request) -> Task {
         >(
             req, BackgroundSchedule::Worker
         ),
+        requests::WillRenameFilesHandler::METHOD => background_request_task::<
+            requests::WillRenameFilesHandler,
+        >(req, BackgroundSchedule::Worker),
         requests::PrepareTypeHierarchyRequestHandler::METHOD => background_document_request_task::<
             requests::PrepareTypeHierarchyRequestHandler,
         >(

crates/ty_server/src/server/api/requests.rs

@@ -38,6 +38,7 @@ mod shutdown;
 mod signature_help;
 mod type_hierarchy_subtypes;
 mod type_hierarchy_supertypes;
+mod will_rename_files;
 mod workspace_diagnostic;
 mod workspace_symbols;
 
@@ -67,5 +68,6 @@ pub(super) use shutdown::ShutdownHandler;
 pub(super) use signature_help::SignatureHelpRequestHandler;
 pub(super) use type_hierarchy_subtypes::TypeHierarchySubtypesRequestHandler;
 pub(super) use type_hierarchy_supertypes::TypeHierarchySupertypesRequestHandler;
+pub(super) use will_rename_files::WillRenameFilesHandler;
 pub(super) use workspace_diagnostic::WorkspaceDiagnosticRequestHandler;
 pub(super) use workspace_symbols::WorkspaceSymbolRequestHandler;

crates/ty_server/src/server/api/requests/will_rename_files.rs

@@ -0,0 +1,182 @@
+use std::collections::HashMap;
+
+use lsp_types::{RenameFilesParams, TextEdit, Uri, WillRenameFilesRequest, WorkspaceEdit};
+use ruff_db::system::SystemPathBuf;
+use ruff_db::{Db as _, files::File};
+use ty_ide::{PathRename, will_rename_paths_in_files};
+use ty_project::{Db as _, ProjectDatabase};
+
+use crate::document::ToRangeExt;
+use crate::server::api::traits::{
+    BackgroundRequestHandler, RequestHandler, RetriableRequestHandler,
+};
+use crate::session::SessionSnapshot;
+use crate::session::client::Client;
+
+/// Handles `workspace/willRenameFiles` requests for Python modules and regular packages.
+pub(crate) struct WillRenameFilesHandler;
+
+impl RequestHandler for WillRenameFilesHandler {
+    type RequestType = WillRenameFilesRequest;
+}
+
+impl BackgroundRequestHandler for WillRenameFilesHandler {
+    fn run(
+        snapshot: &SessionSnapshot,
+        client: &Client,
+        params: RenameFilesParams,
+    ) -> crate::server::Result<Option<WorkspaceEdit>> {
+        let encoding = snapshot.position_encoding();
+        let mut all_changes: HashMap<Uri, Vec<TextEdit>> = HashMap::new();
+        let mut paths = Vec::new();
+        let mut project_index = None;
+
+        for file_rename in &params.files {
+            let (Some(old_path), Some(new_path)) = (
+                file_uri_to_path(&file_rename.old_uri),
+                file_uri_to_path(&file_rename.new_uri),
+            ) else {
+                return Ok(None);
+            };
+
+            let Some(rename_project_index) = snapshot.project_index_for_path(&old_path) else {
+                return Ok(None);
+            };
+            if snapshot
+                .enclosing_project_index_for_path(&new_path)
+                .is_some_and(|index| index != rename_project_index)
+            {
+                return Ok(unsupported_move(client));
+            }
+            if project_index
+                .replace(rename_project_index)
+                .is_some_and(|project_index| project_index != rename_project_index)
+            {
+                return Ok(unsupported_move(client));
+            }
+            paths.push((old_path, new_path));
+        }
+
+        // Use the project that owns the old path. Files in other workspaces stay unchanged even if
+        // that project can import them through an extra search path.
+        let Some(project_index) = project_index else {
+            return Ok(None);
+        };
+        let Some(db) = snapshot.projects().get(project_index) else {
+            return Ok(None);
+        };
+        let Some(workspace_settings) = snapshot.workspace_settings(project_index) else {
+            return Ok(None);
+        };
+        if workspace_settings.is_language_services_disabled() {
+            return Ok(None);
+        }
+
+        // The filesystem tells file moves from folder moves. A folder without an initializer is a
+        // namespace package, which the IDE layer deliberately leaves unchanged.
+        let mut renames = Vec::with_capacity(paths.len());
+        for (old_path, new_path) in paths {
+            if db.system().is_directory(&old_path) {
+                if !["__init__.py", "__init__.pyi"]
+                    .into_iter()
+                    .any(|name| db.system().is_file(&old_path.join(name)))
+                {
+                    return Ok(None);
+                }
+                renames.push(PathRename::directory(old_path, new_path));
+            } else {
+                if !matches!(old_path.extension(), Some("py" | "pyi")) {
+                    return Ok(None);
+                }
+                if old_path.extension() != new_path.extension() {
+                    return Ok(unsupported_move(client));
+                }
+                renames.push(PathRename::file(old_path, new_path));
+            }
+        }
+
+        let project = db.project();
+        let indexed_files = project.files(db);
+        let open_files = project.open_files(db);
+        let files = (&indexed_files)
+            .into_iter()
+            .chain(open_files.iter().copied());
+
+        let Ok(edits) = will_rename_paths_in_files(db, &renames, files, |file| {
+            file_is_in_rename_scope(snapshot, db, project_index, file)
+        }) else {
+            return Ok(unsupported_move(client));
+        };
+        // Convert ty's file and byte ranges to the URI and position ranges required by LSP.
+        for edit in edits {
+            let (file, range, new_text) = edit.into_parts();
+            if !file_is_in_rename_scope(snapshot, db, project_index, file) {
+                continue;
+            }
+            let Some(range) = range.to_lsp_range(db, file, encoding) else {
+                return Ok(unsupported_move(client));
+            };
+            let Some(location) = range.into_location() else {
+                return Ok(unsupported_move(client));
+            };
+
+            all_changes.entry(location.uri).or_default().push(TextEdit {
+                range: location.range,
+                new_text,
+            });
+        }
+
+        if all_changes.values_mut().any(|edits| {
+            edits.sort_by_key(|edit| (edit.range.start, edit.range.end));
+            edits.dedup();
+            edits.windows(2).any(|edits| {
+                edits[0].range.start == edits[1].range.start
+                    || edits[0].range.end > edits[1].range.start
+            })
+        }) {
+            return Ok(unsupported_move(client));
+        }
+
+        if all_changes.is_empty() {
+            Ok(None)
+        } else {
+            Ok(Some(WorkspaceEdit {
+                changes: Some(all_changes),
+                ..Default::default()
+            }))
+        }
+    }
+}
+
+impl RetriableRequestHandler for WillRenameFilesHandler {
+    const RETRY_ON_CANCELLATION: bool = true;
+}
+
+/// Converts an LSP file URI to a system path.
+fn file_uri_to_path(uri: &str) -> Option<SystemPathBuf> {
+    let uri = Uri::parse(uri).ok()?;
+    SystemPathBuf::from_path_buf(uri.to_file_path().ok()?).ok()
+}
+
+/// Returns whether the selected project owns a file or the file has no on-disk path.
+fn file_is_in_rename_scope(
+    snapshot: &SessionSnapshot,
+    db: &ProjectDatabase,
+    project_index: usize,
+    file: File,
+) -> bool {
+    file.path(db).as_system_path().is_none_or(|path| {
+        snapshot
+            .enclosing_project_index_for_path(path)
+            .is_none_or(|index| index == project_index)
+    })
+}
+
+/// Warns the user that ty deliberately left an unsafe move unchanged.
+fn unsupported_move(client: &Client) -> Option<WorkspaceEdit> {
+    client.show_warning_message(
+        "ty could not safely update imports and references for this move. \
+         No automated code changes were applied.",
+    );
+    None
+}

crates/ty_server/src/session.rs

@@ -1145,20 +1145,37 @@ impl Session {
 
     /// Creates a snapshot of the current state of the [`Session`].
     pub(crate) fn snapshot_session(&self) -> SessionSnapshot {
+        // Keep roots and settings in the same order as the cloned databases. The databases remain
+        // in their own final field so they are dropped after the snapshot's other shared state.
+        let (project_contexts, projects) = self
+            .projects
+            .iter()
+            .map(|(workspace_root, project)| {
+                let settings = self
+                    .workspaces
+                    .settings_for_path(workspace_root)
+                    .or_else(|| self.workspaces.settings_virtual_fallback())
+                    .unwrap_or_default();
+                (
+                    ProjectSnapshotContext {
+                        workspace_root: workspace_root.clone(),
+                        settings,
+                    },
+                    project.db.clone(),
+                )
+            })
+            .unzip();
+
         SessionSnapshot {
-            projects: self
-                .projects
-                .values()
-                .map(|project| &project.db)
-                .cloned()
-                .collect(),
             index: self.index.clone().unwrap(),
             global_settings: self.global_settings.clone(),
             position_encoding: self.position_encoding,
             in_test: self.in_test,
             resolved_client_capabilities: self.resolved_client_capabilities,
             revision: self.revision,
             client_name: self.client_name,
+            project_contexts,
+            projects,
         }
     }
 
@@ -1437,6 +1454,8 @@ pub(crate) struct SessionSnapshot {
     in_test: bool,
     revision: u64,
     client_name: ClientName,
+    /// Workspace roots and settings, in the same order as `projects`.
+    project_contexts: Vec<ProjectSnapshotContext>,
 
     /// IMPORTANT: It's important that the databases come last, or at least,
     /// after any `Arc` that we try to extract or mutate in-place using `Arc::into_inner`
@@ -1451,10 +1470,34 @@ pub(crate) struct SessionSnapshot {
 }
 
 impl SessionSnapshot {
+    /// Returns the project databases captured by this snapshot.
     pub(crate) fn projects(&self) -> &[ProjectDatabase] {
         &self.projects
     }
 
+    /// Returns the project owned by the closest enclosing workspace, falling back to the first.
+    pub(crate) fn project_index_for_path(&self, path: &SystemPath) -> Option<usize> {
+        self.enclosing_project_index_for_path(path)
+            .or_else(|| (!self.projects.is_empty()).then_some(0))
+    }
+
+    /// Returns the project owned by the closest enclosing workspace, without a fallback.
+    pub(crate) fn enclosing_project_index_for_path(&self, path: &SystemPath) -> Option<usize> {
+        self.project_contexts
+            .iter()
+            .enumerate()
+            .filter(|(_, context)| path.starts_with(&context.workspace_root))
+            .max_by_key(|(_, context)| context.workspace_root.as_str().len())
+            .map(|(index, _)| index)
+    }
+
+    /// Returns the workspace settings associated with `project_index`.
+    pub(crate) fn workspace_settings(&self, project_index: usize) -> Option<&WorkspaceSettings> {
+        self.project_contexts
+            .get(project_index)
+            .map(|context| context.settings.as_ref())
+    }
+
     pub(crate) fn index(&self) -> &Index {
         &self.index
     }
@@ -1487,7 +1530,9 @@ impl SessionSnapshot {
 /// Represents the client (editor) that's connected to the language server.
 #[derive(Debug, Clone, Copy)]
 pub(crate) enum ClientName {
+    /// The Zed editor.
     Zed,
+    /// Any other editor or an unidentified client.
     Other,
 }
 
@@ -1516,6 +1561,14 @@ impl ClientName {
     }
 }
 
+/// Workspace information kept alongside each project database in a session snapshot.
+struct ProjectSnapshotContext {
+    /// The workspace directory used to decide which project owns a path.
+    workspace_root: SystemPathBuf,
+    /// The workspace settings in effect when the snapshot was created.
+    settings: Arc<WorkspaceSettings>,
+}
+
 #[derive(Debug, Default)]
 pub(crate) struct Workspaces {
     workspaces: BTreeMap<SystemPathBuf, Workspace>,

crates/ty_server/tests/e2e/main.rs

@@ -43,6 +43,7 @@ mod rename;
 mod semantic_tokens;
 mod signature_help;
 mod type_hierarchy;
+mod will_rename_files;
 mod workspace_folders;
 
 use std::collections::{BTreeMap, HashMap, VecDeque};
@@ -73,9 +74,10 @@ use lsp_types::{
     SemanticTokens, ShutdownRequest, SignatureHelp, SignatureHelpParams, SignatureHelpRequest,
     SignatureHelpTriggerKind, TextDocumentClientCapabilities, TextDocumentContentChangeEvent,
     TextDocumentIdentifier, TextDocumentItem, TextDocumentPositionParams, Uri,
-    VersionedTextDocumentIdentifier, WorkDoneProgressParams, WorkspaceClientCapabilities,
-    WorkspaceDiagnosticParams, WorkspaceDiagnosticReport, WorkspaceDiagnosticRequest,
-    WorkspaceEdit, WorkspaceFolder, WorkspaceFoldersChangeEvent, WorkspaceFoldersInitializeParams,
+    VersionedTextDocumentIdentifier, WillRenameFilesRequest, WorkDoneProgressParams,
+    WorkspaceClientCapabilities, WorkspaceDiagnosticParams, WorkspaceDiagnosticReport,
+    WorkspaceDiagnosticRequest, WorkspaceEdit, WorkspaceFolder, WorkspaceFoldersChangeEvent,
+    WorkspaceFoldersInitializeParams,
 };
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf, TestSystem};
 use rustc_hash::FxHashMap;
@@ -932,6 +934,15 @@ impl TestServer {
         )
     }
 
+    pub(crate) fn will_rename_files(
+        &mut self,
+        renames: Vec<lsp_types::FileRename>,
+    ) -> Option<WorkspaceEdit> {
+        self.send_request_await::<WillRenameFilesRequest>(lsp_types::RenameFilesParams {
+            files: renames,
+        })
+    }
+
     /// Send a `textDocument/diagnostic` request for the document at the given path.
     pub(crate) fn document_diagnostic_request(
         &mut self,