From 6156d4761757f6f67d84834c872ed8c606e604cf Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 20 Apr 2026 18:42:48 +0900 Subject: [PATCH] fix --- pyrefly/lib/lsp/non_wasm/mod.rs | 1 + .../lib/lsp/non_wasm/move_symbol_new_file.rs | 199 ++++++++++++++++ pyrefly/lib/lsp/non_wasm/server.rs | 14 ++ pyrefly/lib/state/lsp.rs | 43 ++++ .../lib/state/lsp/quick_fixes/move_module.rs | 212 +++++++++++++++--- pyrefly/lib/test/lsp/lsp_interaction/mod.rs | 1 + .../lsp_interaction/move_symbol_new_file.rs | 114 ++++++++++ .../move_symbol_to_new_file/consumer.py | 3 + .../move_symbol_to_new_file/source.py | 2 + 9 files changed, 564 insertions(+), 25 deletions(-) create mode 100644 pyrefly/lib/lsp/non_wasm/move_symbol_new_file.rs create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/move_symbol_new_file.rs create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/consumer.py create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/source.py diff --git a/pyrefly/lib/lsp/non_wasm/mod.rs b/pyrefly/lib/lsp/non_wasm/mod.rs index f8a20a5803..919a140db0 100644 --- a/pyrefly/lib/lsp/non_wasm/mod.rs +++ b/pyrefly/lib/lsp/non_wasm/mod.rs @@ -14,6 +14,7 @@ pub mod external_provider; pub mod folding_ranges; pub mod lsp; pub mod module_helpers; +pub mod move_symbol_new_file; mod mru; pub mod protocol; pub mod queue; diff --git a/pyrefly/lib/lsp/non_wasm/move_symbol_new_file.rs b/pyrefly/lib/lsp/non_wasm/move_symbol_new_file.rs new file mode 100644 index 0000000000..8b039f97c0 --- /dev/null +++ b/pyrefly/lib/lsp/non_wasm/move_symbol_new_file.rs @@ -0,0 +1,199 @@ +/* + * 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 dupe::Dupe; +use lsp_types::ClientCapabilities; +use lsp_types::CodeAction; +use lsp_types::CodeActionKind; +use lsp_types::CodeActionOrCommand; +use lsp_types::CreateFile; +use lsp_types::DocumentChangeOperation; +use lsp_types::DocumentChanges; +use lsp_types::OneOf; +use lsp_types::OptionalVersionedTextDocumentIdentifier; +use lsp_types::Position; +use lsp_types::Range; +use lsp_types::ResourceOp; +use lsp_types::ResourceOperationKind; +use lsp_types::TextDocumentEdit; +use lsp_types::TextEdit; +use lsp_types::Url; +use lsp_types::WorkspaceEdit; +use pyrefly_build::handle::Handle; +use pyrefly_python::PYTHON_EXTENSIONS; +use pyrefly_python::module_name::ModuleName; +use pyrefly_python::module_path::ModulePath; +use pyrefly_util::absolutize::Absolutize as _; +use ruff_text_size::TextRange; + +use crate::lsp::non_wasm::module_helpers::PathRemapper; +use crate::lsp::non_wasm::module_helpers::module_info_to_uri; +use crate::state::lsp::ImportFormat; +use crate::state::state::Transaction; + +fn supports_workspace_edit_document_changes(capabilities: &ClientCapabilities) -> bool { + capabilities + .workspace + .as_ref() + .and_then(|workspace| workspace.workspace_edit.as_ref()) + .and_then(|workspace_edit| workspace_edit.document_changes) + .unwrap_or(false) +} + +fn supports_workspace_edit_resource_ops( + capabilities: &ClientCapabilities, + required: &[ResourceOperationKind], +) -> bool { + let supported = capabilities + .workspace + .as_ref() + .and_then(|workspace| workspace.workspace_edit.as_ref()) + .and_then(|workspace_edit| workspace_edit.resource_operations.as_ref()); + required + .iter() + .all(|kind| supported.is_some_and(|ops| ops.contains(kind))) +} + +fn path_to_uri(path: &std::path::Path, remapper: Option<&PathRemapper>) -> Option { + let final_path = remapper + .map(|remap| remap(path).into_owned()) + .unwrap_or_else(|| path.to_path_buf()); + let abs_path = final_path.absolutize(); + Url::from_file_path(abs_path).ok() +} + +pub(crate) fn move_symbol_to_new_file_code_action( + capabilities: &ClientCapabilities, + transaction: &Transaction<'_>, + handle: &Handle, + uri: &Url, + selection: TextRange, + import_format: ImportFormat, + path_remapper: Option<&PathRemapper>, +) -> Option { + if !supports_workspace_edit_document_changes(capabilities) { + return None; + } + if !supports_workspace_edit_resource_ops(capabilities, &[ResourceOperationKind::Create]) { + return None; + } + + let path = uri.to_file_path().ok()?; + let extension = path.extension().and_then(|ext| ext.to_str())?; + if !PYTHON_EXTENSIONS.contains(&extension) { + return None; + } + + let context = transaction.module_member_move_context(handle, selection)?; + let new_path = path + .parent()? + .join(format!("{}.{}", context.member_name, extension)); + if new_path == path || new_path.exists() { + return None; + } + + let config = transaction + .config_finder() + .python_file(handle.module_kind(), context.module_info.path()); + let new_module_name = ModuleName::from_path( + &new_path, + config.search_path().chain( + config + .fallback_search_path + .for_directory(new_path.parent()) + .iter(), + ), + &config.extra_file_extensions, + )?; + let target_handle = Handle::new( + new_module_name, + ModulePath::filesystem(new_path.clone()), + handle.sys_info().dupe(), + ); + + let mut edits = transaction.module_member_source_move_edits( + handle, + &context, + &target_handle, + import_format, + )?; + edits.extend(transaction.module_member_consumer_import_updates( + handle, + &context.module_info, + &context.member_name, + &target_handle, + import_format, + )); + + let new_uri = path_to_uri(&new_path, path_remapper)?; + let mut changes: HashMap> = HashMap::new(); + for (module, range, new_text) in edits { + let Some(edit_uri) = module_info_to_uri(&module, path_remapper) else { + continue; + }; + changes.entry(edit_uri).or_default().push(TextEdit { + range: module.to_lsp_range(range), + new_text, + }); + } + + let mut operations = vec![ + DocumentChangeOperation::Op(ResourceOp::Create(CreateFile { + uri: new_uri.clone(), + options: None, + annotation_id: None, + })), + DocumentChangeOperation::Edit(TextDocumentEdit { + text_document: OptionalVersionedTextDocumentIdentifier { + uri: new_uri, + version: None, + }, + edits: vec![OneOf::Left(TextEdit { + range: Range { + start: Position::new(0, 0), + end: Position::new(0, 0), + }, + new_text: context.member_text, + })], + }), + ]; + + let mut sorted_changes: Vec<(Url, Vec)> = changes.into_iter().collect(); + sorted_changes.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str())); + for (uri, mut text_edits) in sorted_changes { + text_edits.sort_by(|a, b| { + ( + a.range.start.line, + a.range.start.character, + a.range.end.line, + a.range.end.character, + ) + .cmp(&( + b.range.start.line, + b.range.start.character, + b.range.end.line, + b.range.end.character, + )) + }); + operations.push(DocumentChangeOperation::Edit(TextDocumentEdit { + text_document: OptionalVersionedTextDocumentIdentifier { uri, version: None }, + edits: text_edits.into_iter().map(OneOf::Left).collect(), + })); + } + + Some(CodeActionOrCommand::CodeAction(CodeAction { + title: format!("Move `{}` to new file", context.member_name), + kind: Some(CodeActionKind::new("refactor.move")), + edit: Some(WorkspaceEdit { + document_changes: Some(DocumentChanges::Operations(operations)), + ..Default::default() + }), + ..Default::default() + })) +} diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index ffefbbc391..01fedfa21d 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -288,6 +288,7 @@ use crate::lsp::non_wasm::module_helpers::ThriftRemapper; use crate::lsp::non_wasm::module_helpers::handle_from_module_path; use crate::lsp::non_wasm::module_helpers::make_open_handle; use crate::lsp::non_wasm::module_helpers::module_info_to_uri; +use crate::lsp::non_wasm::move_symbol_new_file::move_symbol_to_new_file_code_action; use crate::lsp::non_wasm::mru::CompletionMru; use crate::lsp::non_wasm::protocol::Message; use crate::lsp::non_wasm::protocol::Notification; @@ -4608,6 +4609,19 @@ impl Server { actions.push(action); } record_code_action_telemetry("convert_module_package", start); + let start = Instant::now(); + if let Some(action) = move_symbol_to_new_file_code_action( + &self.initialize_params.capabilities, + transaction, + &handle, + uri, + range, + import_format, + self.path_remapper.as_ref(), + ) { + actions.push(action); + } + record_code_action_telemetry("move_symbol_new_file", start); } let start = Instant::now(); if let Some(action) = safe_delete_file_code_action( diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 2859671591..621a741e4c 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -93,6 +93,7 @@ use crate::types::types::Type; mod dict_completions; mod quick_fixes; +pub(crate) use self::quick_fixes::move_module::MoveModuleMemberContext; pub(crate) use self::quick_fixes::types::LocalRefactorCodeAction; #[derive(Debug)] @@ -2667,6 +2668,48 @@ impl<'a> Transaction<'a> { ) } + pub(crate) fn module_member_move_context( + &self, + handle: &Handle, + selection: TextRange, + ) -> Option { + quick_fixes::move_module::module_member_move_context(self, handle, selection) + } + + pub(crate) fn module_member_source_move_edits( + &self, + handle: &Handle, + context: &MoveModuleMemberContext, + target_handle: &Handle, + import_format: ImportFormat, + ) -> Option> { + quick_fixes::move_module::build_module_member_source_move_edits( + self, + handle, + context, + target_handle, + import_format, + ) + } + + pub(crate) fn module_member_consumer_import_updates( + &self, + source_handle: &Handle, + source_module_info: &ModuleInfo, + member_name: &str, + target_handle: &Handle, + import_format: ImportFormat, + ) -> Vec<(ModuleInfo, TextRange, String)> { + quick_fixes::move_module::build_module_member_consumer_import_updates( + self, + source_handle, + source_module_info, + member_name, + target_handle, + import_format, + ) + } + pub fn make_local_function_top_level_code_actions( &self, handle: &Handle, diff --git a/pyrefly/lib/state/lsp/quick_fixes/move_module.rs b/pyrefly/lib/state/lsp/quick_fixes/move_module.rs index 61c89fe50a..7161e0c340 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/move_module.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/move_module.rs @@ -5,11 +5,14 @@ * LICENSE file in the root directory of this source tree. */ +use std::sync::Arc; + use dupe::Dupe; use lsp_types::CodeActionKind; use pyrefly_build::handle::Handle; use pyrefly_python::module::Module; use pyrefly_python::module_name::ModuleName; +use pyrefly_python::module_path::ModulePath; use pyrefly_python::module_path::ModulePathDetails; use ruff_python_ast::Decorator; use ruff_python_ast::ModModule; @@ -40,6 +43,15 @@ fn move_kind() -> CodeActionKind { CodeActionKind::new("refactor.move") } +#[derive(Clone)] +pub(crate) struct MoveModuleMemberContext { + pub module_info: Module, + pub ast: Arc, + pub member_name: String, + pub member_text: String, + pub removal_range: TextRange, +} + #[derive(Clone, Copy, Debug)] enum ParentKind<'a> { Module, @@ -68,42 +80,40 @@ pub(crate) fn move_module_member_code_actions( selection: TextRange, import_format: ImportFormat, ) -> Option> { - let module_info = transaction.get_module_info(handle)?; - let ast = transaction.get_ast(handle)?; - let source = module_info.contents(); - let selection_point = selection_anchor(source, selection); - let member_stmt = find_module_member(ast.as_ref(), selection_point)?; - let member_name = member_name_from_stmt(member_stmt)?; - let (from_indent, _) = line_indent_and_start(source, member_stmt.range().start())?; - let member_text = reindent_statement(source, member_stmt.range(), &from_indent, ""); - - let removal_range = statement_removal_range(source, member_stmt)?; + let context = module_member_move_context(transaction, handle, selection)?; let mut actions = Vec::new(); for (target_handle, target_info, target_ast) in - sibling_module_targets(transaction, handle, &module_info)? + sibling_module_targets(transaction, handle, &context.module_info)? { - if target_info.path() == module_info.path() { + if target_info.path() == context.module_info.path() { continue; } - let insert_edit = - build_module_insertion_edit(&target_info, target_ast.as_ref(), &member_text, None)?; - let (removal_edit, import_edit) = build_removal_and_import_edits( + let insert_edit = build_module_insertion_edit( + &target_info, + target_ast.as_ref(), + &context.member_text, + None, + )?; + let mut edits = vec![insert_edit]; + edits.extend(build_module_member_source_move_edits( transaction, handle, - &module_info, - ast.as_ref(), - &member_name, + &context, &target_handle, import_format, - removal_range, - )?; - let mut edits = vec![insert_edit, removal_edit]; - if let Some(import_edit) = import_edit { - edits.push(import_edit); - } + )?); + edits.extend(build_module_member_consumer_import_updates( + transaction, + handle, + &context.module_info, + &context.member_name, + &target_handle, + import_format, + )); actions.push(LocalRefactorCodeAction { title: format!( - "Move `{member_name}` to `{}`", + "Move `{}` to `{}`", + context.member_name, target_handle.module().as_str() ), edits, @@ -118,6 +128,151 @@ pub(crate) fn move_module_member_code_actions( } } +pub(crate) fn module_member_move_context( + transaction: &Transaction<'_>, + handle: &Handle, + selection: TextRange, +) -> Option { + let module_info = transaction.get_module_info(handle)?; + let ast = transaction.get_ast(handle)?; + let source = module_info.contents(); + let selection_point = selection_anchor(source, selection); + let member_stmt = find_module_member(ast.as_ref(), selection_point)?; + let member_name = member_name_from_stmt(member_stmt)?; + let (from_indent, _) = line_indent_and_start(source, member_stmt.range().start())?; + let member_text = reindent_statement(source, member_stmt.range(), &from_indent, ""); + let removal_range = statement_removal_range(source, member_stmt)?; + Some(MoveModuleMemberContext { + module_info, + ast, + member_name, + member_text, + removal_range, + }) +} + +pub(crate) fn build_module_member_source_move_edits( + transaction: &Transaction<'_>, + handle: &Handle, + context: &MoveModuleMemberContext, + target_handle: &Handle, + import_format: ImportFormat, +) -> Option> { + let (removal_edit, import_edit) = build_removal_and_import_edits( + transaction, + handle, + &context.module_info, + context.ast.as_ref(), + &context.member_name, + target_handle, + import_format, + context.removal_range, + )?; + let mut edits = vec![removal_edit]; + if let Some(import_edit) = import_edit { + edits.push(import_edit); + } + Some(edits) +} + +pub(crate) fn build_module_member_consumer_import_updates( + transaction: &Transaction<'_>, + source_handle: &Handle, + source_module_info: &Module, + member_name: &str, + target_handle: &Handle, + import_format: ImportFormat, +) -> Vec<(Module, TextRange, String)> { + let old_module_name = source_handle.module(); + let rdep_root = match source_handle.path().details() { + ModulePathDetails::Memory(path) => Handle::new( + source_handle.module(), + ModulePath::filesystem((**path).clone()), + source_handle.sys_info().dupe(), + ), + _ => source_handle.dupe(), + }; + let rdeps = transaction.get_transitive_rdeps(rdep_root); + let unique_rdeps: Vec<_> = { + let mut seen = std::collections::HashSet::new(); + rdeps + .into_iter() + .filter(|handle| seen.insert(handle.path().as_path().to_owned())) + .collect() + }; + + let mut edits = Vec::new(); + for rdep_handle in unique_rdeps { + if rdep_handle.path() == source_module_info.path() { + continue; + } + let Some(module_info) = transaction.get_module_info(&rdep_handle) else { + continue; + }; + let Some(ast) = transaction.get_ast(&rdep_handle) else { + continue; + }; + let has_target_import = + has_existing_import(ast.as_ref(), target_handle.module(), member_name); + let (_, _, new_module_text) = insert_import_edit( + ast.as_ref(), + transaction.config_finder(), + rdep_handle.dupe(), + target_handle.dupe(), + member_name, + import_format, + ); + for stmt in &ast.body { + let Stmt::ImportFrom(import_from) = stmt else { + continue; + }; + let Some(module) = &import_from.module else { + continue; + }; + if import_from.level != 0 || ModuleName::from_name(&module.id) != old_module_name { + continue; + } + + let mut moved_aliases = Vec::new(); + let mut remaining_aliases = Vec::new(); + for alias in &import_from.names { + let rendered = render_import_alias(alias); + if alias.name.id.as_str() == member_name { + moved_aliases.push(rendered); + } else { + remaining_aliases.push(rendered); + } + } + if moved_aliases.is_empty() { + continue; + } + + let Some((indent, _)) = + line_indent_and_start(module_info.contents(), import_from.range().start()) + else { + continue; + }; + + let mut replacement = String::new(); + if !remaining_aliases.is_empty() { + replacement.push_str(&format!( + "{indent}from {} import {}\n", + module.id.as_str(), + remaining_aliases.join(", ") + )); + } + if !has_target_import { + replacement.push_str(&format!( + "{indent}from {new_module_text} import {}\n", + moved_aliases.join(", ") + )); + } + edits.push((module_info.dupe(), import_from.range(), replacement)); + } + } + edits +} + /// Builds make-local-function/method-top-level code actions. pub(crate) fn make_local_function_top_level_code_actions( transaction: &Transaction<'_>, @@ -417,6 +572,13 @@ fn has_existing_import(ast: &ModModule, module_name: ModuleName, name: &str) -> }) } +fn render_import_alias(alias: &ruff_python_ast::Alias) -> String { + match &alias.asname { + Some(asname) => format!("{} as {}", alias.name.id, asname.id), + None => alias.name.id.to_string(), + } +} + fn build_module_insertion_edit( module_info: &Module, ast: &ModModule, diff --git a/pyrefly/lib/test/lsp/lsp_interaction/mod.rs b/pyrefly/lib/test/lsp/lsp_interaction/mod.rs index a3e782df08..0f521aa407 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/mod.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/mod.rs @@ -28,6 +28,7 @@ mod hover; mod implementation; mod inlay_hint; mod io; +mod move_symbol_new_file; mod no_config_warnings; mod notebook_code_action; mod notebook_completion; diff --git a/pyrefly/lib/test/lsp/lsp_interaction/move_symbol_new_file.rs b/pyrefly/lib/test/lsp/lsp_interaction/move_symbol_new_file.rs new file mode 100644 index 0000000000..1e1255c2f1 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/move_symbol_new_file.rs @@ -0,0 +1,114 @@ +/* + * 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 lsp_types::CodeActionOrCommand; +use lsp_types::DocumentChangeOperation; +use lsp_types::DocumentChanges; +use lsp_types::ResourceOp; +use lsp_types::TextEdit; +use lsp_types::Url; +use lsp_types::request::CodeActionRequest; +use serde_json::json; + +use crate::object_model::InitializeSettings; +use crate::object_model::LspInteraction; +use crate::util::get_test_files_root; + +fn init_with_create_support(root_path: &std::path::Path) -> (LspInteraction, Url) { + let scope_uri = Url::from_file_path(root_path).unwrap(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root_path.to_path_buf()); + interaction + .initialize(InitializeSettings { + workspace_folders: Some(vec![("test".to_owned(), scope_uri.clone())]), + capabilities: Some(json!({ + "workspace": { + "workspaceEdit": { + "documentChanges": true, + "resourceOperations": ["create"] + } + } + })), + ..Default::default() + }) + .unwrap(); + (interaction, scope_uri) +} + +fn has_edit(ops: &[DocumentChangeOperation], uri: &Url, expected_text: &str) -> bool { + ops.iter().any(|op| { + let DocumentChangeOperation::Edit(edit) = op else { + return false; + }; + edit.text_document.uri == *uri + && edit.edits.iter().any(|edit| match edit { + lsp_types::OneOf::Left(TextEdit { new_text, .. }) => new_text == expected_text, + lsp_types::OneOf::Right(_) => false, + }) + }) +} + +#[test] +fn test_move_symbol_to_new_file_code_action() { + let root = get_test_files_root(); + let root_path = root.path().join("move_symbol_to_new_file"); + let (interaction, _scope_uri) = init_with_create_support(&root_path); + + let source_path = root_path.join("source.py"); + let source_uri = Url::from_file_path(&source_path).unwrap(); + let consumer_uri = Url::from_file_path(root_path.join("consumer.py")).unwrap(); + let new_uri = Url::from_file_path(root_path.join("test.py")).unwrap(); + + interaction.client.did_open("source.py"); + interaction.client.did_open("consumer.py"); + + interaction + .client + .send_request::(json!({ + "textDocument": { "uri": source_uri }, + "range": { + "start": { "line": 0, "character": 4 }, + "end": { "line": 0, "character": 4 } + }, + "context": { "diagnostics": [] } + })) + .expect_response_with(|response: Option>| { + let Some(actions) = response else { + return false; + }; + actions.iter().any(|action| { + let CodeActionOrCommand::CodeAction(code_action) = action else { + return false; + }; + if code_action.title != "Move `test` to new file" { + return false; + } + let Some(edit) = &code_action.edit else { + return false; + }; + let Some(DocumentChanges::Operations(ops)) = &edit.document_changes else { + return false; + }; + if ops.len() != 4 { + return false; + } + let has_create = ops.iter().any(|op| match op { + DocumentChangeOperation::Op(ResourceOp::Create(create)) => { + create.uri == new_uri + } + _ => false, + }); + has_create + && has_edit(ops, &new_uri, "def test(x, y):\n return x + y\n") + && has_edit(ops, &source_uri, "from test import test\n") + && has_edit(ops, &consumer_uri, "from test import test\n") + }) + }) + .unwrap(); + + interaction.shutdown().unwrap(); +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/consumer.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/consumer.py new file mode 100644 index 0000000000..9d864c3e03 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/consumer.py @@ -0,0 +1,3 @@ +from source import test + +result = test(1, 2) diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/source.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/source.py new file mode 100644 index 0000000000..eb7aadb40b --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/move_symbol_to_new_file/source.py @@ -0,0 +1,2 @@ +def test(x, y): + return x + y