implement basic will_rename_files (#1284)

Summary:

basic support for willRenameFiles.

I found document_changes to be the only way that worked reliably on windows so I implemented both document_changes and changes (depending on the client capability).

Differential Revision: D84537597

Author: kinto0

pyrefly/lib/lsp/features/mod.rs

@@ -7,3 +7,4 @@
 
 pub mod hover;
 pub mod provide_type;
+pub mod will_rename_files;

pyrefly/lib/lsp/features/will_rename_files.rs

@@ -0,0 +1,311 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use lsp_types::DocumentChangeOperation;
+use lsp_types::DocumentChanges;
+use lsp_types::OneOf;
+use lsp_types::OptionalVersionedTextDocumentIdentifier;
+use lsp_types::RenameFilesParams;
+use lsp_types::TextDocumentEdit;
+use lsp_types::TextEdit;
+use lsp_types::Url;
+use lsp_types::WorkspaceEdit;
+use pyrefly_python::PYTHON_EXTENSIONS;
+use pyrefly_python::ast::Ast;
+use pyrefly_python::module_name::ModuleName;
+use pyrefly_python::module_path::ModulePath;
+use pyrefly_util::lined_buffer::LinedBuffer;
+use pyrefly_util::lock::RwLock;
+use rayon::prelude::*;
+use ruff_python_ast::Stmt;
+use ruff_text_size::Ranged;
+
+use crate::lsp::module_helpers::handle_from_module_path;
+use crate::lsp::module_helpers::module_info_to_uri;
+use crate::state::state::State;
+use crate::state::state::Transaction;
+
+/// Visitor that looks for imports of an old module name and creates TextEdits to update them
+struct RenameUsageVisitor<'a> {
+    edits: Vec<TextEdit>,
+    old_module_name: &'a ModuleName,
+    new_module_name: &'a ModuleName,
+    lined_buffer: &'a LinedBuffer,
+}
+
+impl<'a> RenameUsageVisitor<'a> {
+    fn new(
+        old_module_name: &'a ModuleName,
+        new_module_name: &'a ModuleName,
+        lined_buffer: &'a LinedBuffer,
+    ) -> Self {
+        Self {
+            edits: Vec::new(),
+            old_module_name,
+            new_module_name,
+            lined_buffer,
+        }
+    }
+
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        match stmt {
+            Stmt::Import(import) => {
+                for alias in &import.names {
+                    let imported_module = ModuleName::from_name(&alias.name.id);
+                    if imported_module == *self.old_module_name
+                        || imported_module
+                            .as_str()
+                            .starts_with(&format!("{}.", self.old_module_name.as_str()))
+                    {
+                        // Replace the module name
+                        let new_import_name = if imported_module == *self.old_module_name {
+                            self.new_module_name.as_str().to_owned()
+                        } else {
+                            // Replace the prefix
+                            imported_module.as_str().replace(
+                                self.old_module_name.as_str(),
+                                self.new_module_name.as_str(),
+                            )
+                        };
+
+                        self.edits.push(TextEdit {
+                            range: self.lined_buffer.to_lsp_range(alias.name.range()),
+                            new_text: new_import_name,
+                        });
+                    }
+                }
+            }
+            Stmt::ImportFrom(import_from) => {
+                if let Some(module) = &import_from.module {
+                    let imported_module = ModuleName::from_name(&module.id);
+                    if imported_module == *self.old_module_name
+                        || imported_module
+                            .as_str()
+                            .starts_with(&format!("{}.", self.old_module_name.as_str()))
+                    {
+                        // Replace the module name
+                        let new_import_name = if imported_module == *self.old_module_name {
+                            self.new_module_name.as_str().to_owned()
+                        } else {
+                            // Replace the prefix
+                            imported_module.as_str().replace(
+                                self.old_module_name.as_str(),
+                                self.new_module_name.as_str(),
+                            )
+                        };
+
+                        self.edits.push(TextEdit {
+                            range: self.lined_buffer.to_lsp_range(module.range()),
+                            new_text: new_import_name,
+                        });
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+
+    fn take_edits(self) -> Vec<TextEdit> {
+        self.edits
+    }
+}
+
+/// Handle workspace/willRenameFiles request to update imports when files are renamed.
+///
+/// This function:
+/// 1. Converts file paths to module names
+/// 2. Uses get_transitive_rdeps to find all files that depend on the renamed module
+/// 3. Uses a visitor pattern to find imports of the old module and creates TextEdits
+/// 4. Returns a WorkspaceEdit with all necessary changes
+///
+/// If the client supports `workspace.workspaceEdit.documentChanges`, the response will use
+/// `document_changes` instead of `changes` for better ordering guarantees and version checking.
+pub fn will_rename_files(
+    state: &Arc<State>,
+    transaction: &Transaction<'_>,
+    _open_files: &Arc<RwLock<HashMap<std::path::PathBuf, Arc<String>>>>,
+    params: RenameFilesParams,
+    supports_document_changes: bool,
+) -> Option<WorkspaceEdit> {
+    eprintln!(
+        "will_rename_files called with {} file(s)",
+        params.files.len()
+    );
+
+    let mut all_changes: HashMap<Url, Vec<TextEdit>> = HashMap::new();
+
+    for file_rename in &params.files {
+        eprintln!(
+            "  Processing rename: {} -> {}",
+            file_rename.old_uri, file_rename.new_uri
+        );
+
+        // Convert URLs to paths
+        let old_uri = match Url::parse(&file_rename.old_uri) {
+            Ok(uri) => uri,
+            Err(_) => {
+                eprintln!("    Failed to parse old_uri");
+                continue;
+            }
+        };
+
+        let new_uri = match Url::parse(&file_rename.new_uri) {
+            Ok(uri) => uri,
+            Err(_) => {
+                eprintln!("    Failed to parse new_uri");
+                continue;
+            }
+        };
+
+        let old_path = match old_uri.to_file_path() {
+            Ok(path) => path,
+            Err(_) => {
+                eprintln!("    Failed to convert old_uri to path");
+                continue;
+            }
+        };
+
+        let new_path = match new_uri.to_file_path() {
+            Ok(path) => path,
+            Err(_) => {
+                eprintln!("    Failed to convert new_uri to path");
+                continue;
+            }
+        };
+
+        // Only process Python files
+        if !PYTHON_EXTENSIONS
+            .iter()
+            .any(|ext| old_path.extension().and_then(|e| e.to_str()) == Some(*ext))
+        {
+            eprintln!("    Skipping non-Python file");
+            continue;
+        }
+
+        // Important: only use filesystem handle (never use an in-memory handle)
+        let old_handle = handle_from_module_path(state, ModulePath::filesystem(old_path.clone()));
+        let config = state
+            .config_finder()
+            .python_file(old_handle.module(), old_handle.path());
+
+        // Convert paths to module names
+        let old_module_name =
+            ModuleName::from_path(&old_path, config.search_path()).or_else(|| {
+                // Fallback: try to get module name from the handle
+                Some(old_handle.module())
+            });
+
+        // For the new module name, we can't rely on from_path because the file doesn't exist yet.
+        // Instead, we compute the relative path from the old to new file and adjust the module name.
+        let new_module_name = ModuleName::from_path(&new_path, config.search_path());
+
+        let (old_module_name, new_module_name) = match (old_module_name, new_module_name) {
+            (Some(old), Some(new)) => (old, new),
+            _ => {
+                eprintln!(
+                    "    Could not determine module names for the rename (old={:?}, new={:?})",
+                    old_module_name, new_module_name
+                );
+                continue;
+            }
+        };
+
+        eprintln!(
+            "    Module rename: {} -> {}",
+            old_module_name, new_module_name
+        );
+
+        // If module names are the same, no need to update imports
+        if old_module_name == new_module_name {
+            eprintln!("    Module names are the same, skipping");
+            continue;
+        }
+
+        // Use get_transitive_rdeps to find all files that depend on this module
+        let rdeps = transaction.get_transitive_rdeps(old_handle.clone());
+
+        eprintln!("    Found {} transitive rdeps", rdeps.len());
+
+        // Visit each dependent file to find and update imports (parallelized)
+        let rdeps_changes: Vec<(Url, Vec<TextEdit>)> = rdeps
+            .into_par_iter()
+            .filter_map(|rdep_handle| {
+                let module_info = transaction.get_module_info(&rdep_handle)?;
+
+                let ast = Ast::parse(module_info.contents()).0;
+                let mut visitor = RenameUsageVisitor::new(
+                    &old_module_name,
+                    &new_module_name,
+                    module_info.lined_buffer(),
+                );
+
+                for stmt in &ast.body {
+                    visitor.visit_stmt(stmt);
+                }
+
+                let edits_for_file = visitor.take_edits();
+
+                if !edits_for_file.is_empty() {
+                    let uri = module_info_to_uri(&module_info)?;
+                    eprintln!(
+                        "    Found {} import(s) to update in {}",
+                        edits_for_file.len(),
+                        uri
+                    );
+                    Some((uri, edits_for_file))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        // Merge results into all_changes
+        for (uri, edits) in rdeps_changes {
+            all_changes.entry(uri).or_default().extend(edits);
+        }
+    }
+
+    if all_changes.is_empty() {
+        eprintln!("  No import updates needed");
+        None
+    } else {
+        eprintln!(
+            "  Returning {} file(s) with import updates",
+            all_changes.len()
+        );
+
+        if supports_document_changes {
+            // Use document_changes for better ordering guarantees and version checking
+            let document_changes: Vec<DocumentChangeOperation> = all_changes
+                .into_iter()
+                .map(|(uri, edits)| {
+                    DocumentChangeOperation::Edit(TextDocumentEdit {
+                        text_document: OptionalVersionedTextDocumentIdentifier {
+                            uri,
+                            version: None, // None means "any version"
+                        },
+                        edits: edits.into_iter().map(OneOf::Left).collect(),
+                    })
+                })
+                .collect();
+
+            Some(WorkspaceEdit {
+                document_changes: Some(DocumentChanges::Operations(document_changes)),
+                ..Default::default()
+            })
+        } else {
+            // Fall back to changes for older clients
+            Some(WorkspaceEdit {
+                changes: Some(all_changes),
+                ..Default::default()
+            })
+        }
+    }
+}

pyrefly/lib/lsp/server.rs

@@ -176,6 +176,7 @@ use crate::lsp::features::hover::get_hover;
 use crate::lsp::features::provide_type::ProvideType;
 use crate::lsp::features::provide_type::ProvideTypeResponse;
 use crate::lsp::features::provide_type::provide_type;
+use crate::lsp::features::will_rename_files::will_rename_files;
 use crate::lsp::lsp::apply_change_events;
 use crate::lsp::lsp::as_notification;
 use crate::lsp::lsp::as_request;
@@ -908,9 +909,21 @@ impl Server {
                     {
                         let transaction =
                             ide_transaction_manager.non_committable_transaction(&self.state);
+                        let supports_document_changes = self
+                            .initialize_params
+                            .capabilities
+                            .workspace
+                            .as_ref()
+                            .and_then(|w| w.workspace_edit.as_ref())
+                            .and_then(|we| we.document_changes)
+                            .unwrap_or(false);
                         self.send_response(new_response(
                             x.id,
-                            Ok(self.will_rename_files(&transaction, params)),
+                            Ok(self.will_rename_files(
+                                &transaction,
+                                params,
+                                supports_document_changes,
+                            )),
                         ));
                         ide_transaction_manager.save(transaction);
                     }
@@ -2143,20 +2156,17 @@ impl Server {
 
     fn will_rename_files(
         &self,
-        _transaction: &Transaction<'_>,
+        transaction: &Transaction<'_>,
         params: RenameFilesParams,
+        supports_document_changes: bool,
     ) -> Option<WorkspaceEdit> {
-        // TODO: Implement import updates when files are renamed
-        // For now, return None to indicate no edits are needed
-        // This is similar to how basedpyright initially implemented this feature
-        eprintln!(
-            "will_rename_files called with {} file(s)",
-            params.files.len()
-        );
-        for file in &params.files {
-            eprintln!("  Renaming: {} -> {}", file.old_uri, file.new_uri);
-        }
-        None
+        will_rename_files(
+            &self.state,
+            transaction,
+            &self.open_files,
+            params,
+            supports_document_changes,
+        )
     }
 }