From 9d4474cfa6ac388f6d97cc2713766f55f6c4ef91 Mon Sep 17 00:00:00 2001 From: Jnyandeep Date: Sat, 7 Nov 2026 19:59:50 +0530 Subject: [PATCH 1/3] fix: quickfix insert param types --- pyrefly/lib/lsp/non_wasm/server.rs | 127 ++++---- pyrefly/lib/lsp/wasm/inlay_hints.rs | 82 +++-- pyrefly/lib/state/lsp.rs | 16 + pyrefly/lib/state/lsp/quick_fixes.rs | 1 + .../quick_fixes/insert_inferred_annotation.rs | 77 +++++ pyrefly/lib/test/lsp/code_actions.rs | 299 ++++++++++++++++++ pyrefly/lib/test/lsp/inlay_hint.rs | 36 +++ .../lsp_interaction/notebook_code_action.rs | 48 +++ 8 files changed, 605 insertions(+), 81 deletions(-) create mode 100644 pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 8c290c0b38..0db5994024 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5332,71 +5332,82 @@ impl Server { // index so that import quick-fixes can be redirected to the current cell // instead of always targeting cell 1 (position 0 of the combined AST). let triggered_cell_index = self.maybe_get_code_cell_index(uri); - if let Some(quickfixes) = transaction.local_quickfix_code_actions_sorted( + let mut quickfixes = transaction + .local_quickfix_code_actions_sorted( + &handle, + range, + import_format, + Some(&self.lsp_thread_pool), + ) + .unwrap_or_default(); + let annotation_start = Instant::now(); + let inlay_hint_config = lsp_config.and_then(|c| c.inlay_hints).unwrap_or_default(); + if let Some(annotation_fixes) = transaction.insert_inferred_annotation_code_actions( &handle, range, import_format, - Some(&self.lsp_thread_pool), + inlay_hint_config, ) { - actions.extend(quickfixes.into_iter().filter_map(|(title, edits)| { - // Fixes for missing `@override` and context-manager annotations - // can include several edits. Group every edit by its target document - // into a single workspace edit - // so the whole fix applies in one action. - let mut changes: HashMap> = HashMap::new(); - for (info, range, insert_text) in edits { - let lsp_location = self.to_lsp_location(&TextRangeWithModule { - module: info.clone(), - range, - })?; - let mut edit_uri = lsp_location.uri; - let mut edit_range = lsp_location.range; - // Import insertions in notebooks go into the active cell. - // Replacements must retain their target cell and range. - if range.is_empty() - && let Some(current_cell_idx) = triggered_cell_index - { - let edit_cell_idx = info.to_cell_for_lsp(range.start()); - if edit_cell_idx != Some(current_cell_idx) { - // Redirect to the current cell, inserting at line 0. - let open_files = self.open_files.read(); - let notebook_path = - self.open_notebook_cells.read().get(uri).cloned(); - let cell_url = notebook_path.and_then(|path| { - if let Some(LspFile::Notebook(notebook)) = - open_files.get(&path).map(|f| &**f) - { - notebook.get_code_cell_url(current_cell_idx).cloned() - } else { - None - } - }); - if let Some(cell_url) = cell_url { - let top_of_cell = lsp_types::Range { - start: lsp_types::Position::new(0, 0), - end: lsp_types::Position::new(0, 0), - }; - edit_uri = cell_url; - edit_range = top_of_cell; + quickfixes.extend(annotation_fixes); + } + record_code_action_telemetry("insert_inferred_annotation", annotation_start); + actions.extend(quickfixes.into_iter().filter_map(|(title, edits)| { + // Fixes for missing `@override` and context-manager annotations + // can include several edits. Group every edit by its target document + // into a single workspace edit + // so the whole fix applies in one action. + let mut changes: HashMap> = HashMap::new(); + for (info, range, insert_text) in edits { + let lsp_location = self.to_lsp_location(&TextRangeWithModule { + module: info.clone(), + range, + })?; + let mut edit_uri = lsp_location.uri; + let mut edit_range = lsp_location.range; + // Import insertions in notebooks go into the active cell. + // Replacements must retain their target cell and range. + if range.is_empty() + && let Some(current_cell_idx) = triggered_cell_index + { + let edit_cell_idx = info.to_cell_for_lsp(range.start()); + if edit_cell_idx != Some(current_cell_idx) { + // Redirect to the current cell, inserting at line 0. + let open_files = self.open_files.read(); + let notebook_path = self.open_notebook_cells.read().get(uri).cloned(); + let cell_url = notebook_path.and_then(|path| { + if let Some(LspFile::Notebook(notebook)) = + open_files.get(&path).map(|f| &**f) + { + notebook.get_code_cell_url(current_cell_idx).cloned() + } else { + None } + }); + if let Some(cell_url) = cell_url { + let top_of_cell = lsp_types::Range { + start: lsp_types::Position::new(0, 0), + end: lsp_types::Position::new(0, 0), + }; + edit_uri = cell_url; + edit_range = top_of_cell; } - }; - changes.entry(edit_uri).or_default().push(TextEdit { - range: edit_range, - new_text: insert_text, - }); - } - Some(CodeActionResponse::CodeAction(CodeAction { - title, - kind: Some(CodeActionKind::QuickFix), - edit: Some(WorkspaceEdit { - changes: Some(changes), - ..Default::default() - }), + } + }; + changes.entry(edit_uri).or_default().push(TextEdit { + range: edit_range, + new_text: insert_text, + }); + } + Some(CodeActionResponse::CodeAction(CodeAction { + title, + kind: Some(CodeActionKind::QuickFix), + edit: Some(WorkspaceEdit { + changes: Some(changes), ..Default::default() - })) - })); - } + }), + ..Default::default() + })) + })); record_code_action_telemetry("quickfix", start); } if allow_fix_all { diff --git a/pyrefly/lib/lsp/wasm/inlay_hints.rs b/pyrefly/lib/lsp/wasm/inlay_hints.rs index 52ad9568c0..4b3f8f5812 100644 --- a/pyrefly/lib/lsp/wasm/inlay_hints.rs +++ b/pyrefly/lib/lsp/wasm/inlay_hints.rs @@ -118,6 +118,30 @@ impl<'a> Transaction<'a> { handle: &Handle, inlay_hint_config: InlayHintConfig, import_format: ImportFormat, + ) -> Option> { + self.inlay_hints_impl(handle, inlay_hint_config, import_format, None) + } + + /// Inlay hints whose annotated declaration overlaps `range`. + /// + /// Types outside `range` are not rendered. Code actions use this so a cursor + /// request does not format every hint in the file. + pub fn inlay_hints_in_range( + &self, + handle: &Handle, + inlay_hint_config: InlayHintConfig, + import_format: ImportFormat, + range: TextRange, + ) -> Option> { + self.inlay_hints_impl(handle, inlay_hint_config, import_format, Some(range)) + } + + fn inlay_hints_impl( + &self, + handle: &Handle, + inlay_hint_config: InlayHintConfig, + import_format: ImportFormat, + limit: Option, ) -> Option> { let is_interesting = |e: &Expr, ty: &Type, class_name: Option<&Name>| { !ty.is_any() @@ -186,40 +210,48 @@ impl<'a> Transaction<'a> { edits, } }; + // `None` keeps every hint. A selection includes a hint when it overlaps the + // declaration that hint annotates, including a cursor on the insertion point + // at the end of that declaration. + let in_selection = |target: TextRange| match limit { + None => true, + Some(selection) => { + selection.start() <= target.end() && target.start() <= selection.end() + } + }; let mut res = Vec::new(); for idx in bindings.keys::() { match bindings.idx_to_key(idx) { Key::ReturnType(id) if inlay_hint_config.function_return_types => { - match bindings.get(bindings.key_to_idx(&Key::Definition(*id))) { - Binding::Function { decorated_idx, .. } => { - if matches!(&bindings.get(idx), Binding::ReturnType(ret) if !ret.kind.has_return_annotation()) - && let Some(mut ty) = answers.get_type_at(idx) - && !ty.is_any() + if let Binding::Function { decorated_idx, .. } = + bindings.get(bindings.key_to_idx(&Key::Definition(*id))) + { + let fun = bindings.get(bindings.get(*decorated_idx).undecorated_idx); + let is_async = fun.def.is_async; + let position = fun.def.parameters.range.end(); + // The header runs from the name through the parameter list, so a + // cursor in the body does not offer the return annotation. + let header = TextRange::new(fun.def.name.range.start(), position); + if in_selection(header) + && matches!(&bindings.get(idx), Binding::ReturnType(ret) if !ret.kind.has_return_annotation()) + && let Some(mut ty) = answers.get_type_at(idx) + && !ty.is_any() + { + if is_async + && let Some(Some((_, _, return_ty))) = + self.ad_hoc_solve(handle, "inlay_hint_coroutine", |solver| { + solver.unwrap_coroutine(&ty) + }) { - let fun = - bindings.get(bindings.get(*decorated_idx).undecorated_idx); - if fun.def.is_async - && let Some(Some((_, _, return_ty))) = self.ad_hoc_solve( - handle, - "inlay_hint_coroutine", - |solver| solver.unwrap_coroutine(&ty), - ) - { - ty = return_ty; - } - res.push(make_type_hint( - " -> ", - fun.def.parameters.range.end(), - &ty, - true, - )); + ty = return_ty; } + res.push(make_type_hint(" -> ", position, &ty, true)); } - _ => {} } } key @ Key::Definition(_) if inlay_hint_config.variable_types + && in_selection(key.range()) && let Some(ty) = answers.get_type_at(idx) => { let declaration_type = answers.get_class_field_type_at_definition(idx); @@ -299,6 +331,9 @@ impl<'a> Transaction<'a> { if inlay_hint_config.variable_types { for field_idx in bindings.keys::() { let field = bindings.get(field_idx); + if !in_selection(field.range) { + continue; + } if let ClassFieldDefinition::DefinedInMethod { values, annotation: None, @@ -350,6 +385,7 @@ impl<'a> Transaction<'a> { res.extend( self.add_inlay_hints_for_positional_function_args(handle) .into_iter() + .filter(|(pos, _)| in_selection(TextRange::empty(*pos))) .map(|(pos, text)| InlayHintData { position: pos, label_parts: vec![(text.clone(), None)], diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index da73d3e1e8..e5457719a2 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -4211,6 +4211,22 @@ impl<'a> Transaction<'a> { } } + pub fn insert_inferred_annotation_code_actions( + &self, + handle: &Handle, + selection: TextRange, + import_format: ImportFormat, + inlay_hint_config: InlayHintConfig, + ) -> Option)>> { + quick_fixes::insert_inferred_annotation::insert_inferred_annotation_code_actions( + self, + handle, + selection, + import_format, + inlay_hint_config, + ) + } + pub fn pytest_fixture_type_annotation_code_actions( &self, handle: &Handle, diff --git a/pyrefly/lib/state/lsp/quick_fixes.rs b/pyrefly/lib/state/lsp/quick_fixes.rs index dc07441fc6..5dee52c181 100644 --- a/pyrefly/lib/state/lsp/quick_fixes.rs +++ b/pyrefly/lib/state/lsp/quick_fixes.rs @@ -23,6 +23,7 @@ pub(crate) mod implicit_import; pub(crate) mod inline_method; pub(crate) mod inline_parameter; pub(crate) mod inline_variable; +pub(crate) mod insert_inferred_annotation; pub(crate) mod introduce_parameter; pub(crate) mod invert_boolean; pub(crate) mod move_members; diff --git a/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs b/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs new file mode 100644 index 0000000000..2fdfa640db --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs @@ -0,0 +1,77 @@ +/* + * 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 dupe::Dupe; +use pyrefly_build::handle::Handle; +use pyrefly_python::module::Module; +use ruff_text_size::TextRange; + +use crate::error::suppress::detect_line_ending; +use crate::state::lsp::AllOffPartial; +use crate::state::lsp::ImportFormat; +use crate::state::lsp::InlayHintConfig; +use crate::state::state::Transaction; + +/// Offer the insertable inlay hints that overlap `selection` as quick fixes. +/// +/// Each action applies that hint's existing edits: the annotation at the hint +/// position, plus any imports the annotation needs. Call-argument name hints +/// are not type annotations, so they are not requested. +pub(crate) fn insert_inferred_annotation_code_actions( + transaction: &Transaction<'_>, + handle: &Handle, + selection: TextRange, + import_format: ImportFormat, + mut inlay_hint_config: InlayHintConfig, +) -> Option)>> { + if !inlay_hint_config.variable_types && !inlay_hint_config.function_return_types { + return None; + } + inlay_hint_config.call_argument_names = AllOffPartial::Off; + let hints = + transaction.inlay_hints_in_range(handle, inlay_hint_config, import_format, selection)?; + let module = transaction.get_module_info(handle)?; + let line_ending = detect_line_ending(module.contents().as_str()); + let normalize = |text: String| -> String { + if line_ending == "\n" { + text + } else { + text.replace('\n', line_ending) + } + }; + let mut actions = Vec::new(); + for hint in hints { + let Some(edits) = hint.edits else { + continue; + }; + let title = format!("Insert inferred type `{}`", edits.annotation.trim()); + let mut action_edits = Vec::with_capacity(1 + edits.imports.len()); + action_edits.push(( + module.dupe(), + TextRange::empty(hint.position), + normalize(edits.annotation), + )); + for (offset, import_text) in edits.imports { + action_edits.push(( + module.dupe(), + TextRange::empty(offset), + normalize(import_text), + )); + } + actions.push((hint.position, title, action_edits)); + } + if actions.is_empty() { + return None; + } + actions.sort_by_key(|(position, _, _)| *position); + Some( + actions + .into_iter() + .map(|(_, title, edits)| (title, edits)) + .collect(), + ) +} diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index cb23cb1146..40f3d8a9a3 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -23,7 +23,9 @@ use crate::module::module_info::ModuleInfo; use crate::state::ide::import_regular_import_edit; use crate::state::ide::insert_from_import_edit; use crate::state::ide::insert_import_edit; +use crate::state::lsp::AllOffPartial; use crate::state::lsp::ImportFormat; +use crate::state::lsp::InlayHintConfig; use crate::state::lsp::LocalRefactorCodeAction; use crate::state::require::Require; use crate::state::state::State; @@ -8255,3 +8257,300 @@ class Derived(Base): "; assert_eq!(expected, after); } + +fn inferred_annotation_actions( + files: &[(&'static str, &str)], + range: TextRange, + config: InlayHintConfig, +) -> Vec<(String, String)> { + let (handles, state) = mk_multi_file_state_assert_no_errors(files, Require::Exports); + let handle = handles.get("main").unwrap(); + let transaction = state.transaction(); + transaction + .insert_inferred_annotation_code_actions(handle, range, ImportFormat::Absolute, config) + .unwrap_or_default() + .into_iter() + .map(|(title, edits)| { + let info = edits[0].0.clone(); + let after = apply_refactor_edits_for_module(&info, &edits); + (title, after) + }) + .collect() +} + +fn range_of(source: &str, needle: &str) -> TextRange { + let start = source.find(needle).expect("needle"); + TextRange::new( + TextSize::try_from(start).unwrap(), + TextSize::try_from(start + needle.len()).unwrap(), + ) +} + +const ANNOTATED_RETURN: &str = "\ +def grow(x: int): + return x + +yyy = grow(1) +"; + +#[test] +fn insert_inferred_type_on_variable() { + let actions = inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "yyy"), + Default::default(), + ); + assert_eq!(actions.len(), 1); + assert_eq!(actions[0].0, "Insert inferred type `: int`"); + assert_eq!( + actions[0].1, + "\ +def grow(x: int): + return x + +yyy: int = grow(1) +" + ); +} + +#[test] +fn insert_inferred_type_on_return() { + let actions = inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "grow"), + Default::default(), + ); + assert_eq!(actions.len(), 1); + assert_eq!(actions[0].0, "Insert inferred type `-> int`"); + assert_eq!( + actions[0].1, + "\ +def grow(x: int) -> int: + return x + +yyy = grow(1) +" + ); +} + +#[test] +fn insert_inferred_type_at_insertion_point() { + let start = ANNOTATED_RETURN.find("yyy").unwrap() + "yyy".len(); + let cursor = TextRange::empty(TextSize::try_from(start).unwrap()); + let actions = + inferred_annotation_actions(&[("main", ANNOTATED_RETURN)], cursor, Default::default()); + assert_eq!(actions.len(), 1); + assert_eq!(actions[0].0, "Insert inferred type `: int`"); +} + +#[test] +fn insert_inferred_type_on_class_attribute() { + let code = "\ +def value() -> int: + return 1 + +class Box: + def __init__(self) -> None: + self.item = value() +"; + let actions = inferred_annotation_actions( + &[("main", code)], + range_of(code, "item"), + Default::default(), + ); + assert_eq!(actions.len(), 1); + assert_eq!(actions[0].0, "Insert inferred type `: int`"); + assert!(actions[0].1.contains("self.item: int = value()")); +} + +#[test] +fn insert_inferred_type_adds_import() { + let lib = "\ +class Value: + pass + +def make() -> Value: + return Value() +"; + let main = "\ +from lib import make + +value = make() +"; + let actions = inferred_annotation_actions( + &[("lib", lib), ("main", main)], + range_of(main, "value"), + Default::default(), + ); + assert_eq!(actions.len(), 1); + assert_eq!(actions[0].0, "Insert inferred type `: Value`"); + assert_eq!( + actions[0].1, + "\ +from lib import Value +from lib import make + +value: Value = make() +" + ); +} + +#[test] +fn insert_inferred_type_preserves_crlf_imports() { + let lib = "class Value:\r\n pass\r\n\r\ndef make() -> Value:\r\n return Value()\r\n"; + let main = "from lib import make\r\n\r\nvalue = make()\r\n"; + let actions = inferred_annotation_actions( + &[("lib", lib), ("main", main)], + range_of(main, "value"), + Default::default(), + ); + assert_eq!(actions.len(), 1); + assert_eq!( + actions[0].1, + "from lib import Value\r\nfrom lib import make\r\n\r\nvalue: Value = make()\r\n" + ); +} + +#[test] +fn insert_inferred_type_skips_non_insertable_and_unrelated_cursors() { + let unpacked = "\ +def take(pair: tuple[int, int]) -> None: + first, second = pair +"; + assert!( + inferred_annotation_actions( + &[("main", unpacked)], + range_of(unpacked, "first"), + Default::default(), + ) + .is_empty() + ); + + let capture = "\ +def take(xs: list[int]) -> None: + match xs: + case [head]: + print(head) +"; + assert!( + inferred_annotation_actions( + &[("main", capture)], + range_of(capture, "head"), + Default::default(), + ) + .is_empty() + ); + + let annotated = "yyy: int = 1\n"; + assert!( + inferred_annotation_actions( + &[("main", annotated)], + range_of(annotated, "yyy"), + Default::default(), + ) + .is_empty() + ); + + // A cursor in the body is outside the function header the return hint annotates. + assert!( + inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "return"), + Default::default(), + ) + .is_empty() + ); +} + +#[test] +fn insert_inferred_type_respects_inlay_hint_config() { + let hidden = InlayHintConfig { + call_argument_names: AllOffPartial::All, + function_return_types: false, + pytest_parameters: false, + variable_types: false, + }; + assert!( + inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "yyy"), + hidden + ) + .is_empty() + ); + assert!( + inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "grow"), + hidden + ) + .is_empty() + ); + + let mut returns_only = hidden; + returns_only.function_return_types = true; + assert!( + inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "yyy"), + returns_only, + ) + .is_empty() + ); + assert_eq!( + inferred_annotation_actions( + &[("main", ANNOTATED_RETURN)], + range_of(ANNOTATED_RETURN, "grow"), + returns_only, + ) + .len(), + 1 + ); + + // Argument-name hints stay out of this quick fix even when that inlay hint is enabled. + let calls = "\ +def f(value: int) -> None: + return None + +f(1) +"; + let names_on = InlayHintConfig { + call_argument_names: AllOffPartial::All, + ..InlayHintConfig::default() + }; + assert!( + inferred_annotation_actions(&[("main", calls)], range_of(calls, "1"), names_on).is_empty() + ); +} + +#[test] +fn insert_inferred_type_for_each_selected_hint() { + let code = "\ +def grow(x: int): + return x + +yyy = grow(1) +zzz = grow(2) +"; + let start = code.find("yyy").unwrap(); + let end = code.find("zzz").unwrap() + "zzz".len(); + let selection = TextRange::new( + TextSize::try_from(start).unwrap(), + TextSize::try_from(end).unwrap(), + ); + let actions = inferred_annotation_actions(&[("main", code)], selection, Default::default()); + assert_eq!( + actions + .iter() + .map(|(title, _)| title.as_str()) + .collect::>(), + vec![ + "Insert inferred type `: int`", + "Insert inferred type `: int`" + ] + ); + assert!(actions[0].1.contains("yyy: int = grow(1)")); + assert!(actions[1].1.contains("zzz: int = grow(2)")); + assert!(!actions[0].1.contains("zzz: int")); + assert!(!actions[1].1.contains("yyy: int")); +} diff --git a/pyrefly/lib/test/lsp/inlay_hint.rs b/pyrefly/lib/test/lsp/inlay_hint.rs index 3e9265c07d..76e385ce01 100644 --- a/pyrefly/lib/test/lsp/inlay_hint.rs +++ b/pyrefly/lib/test/lsp/inlay_hint.rs @@ -6,6 +6,8 @@ */ use pretty_assertions::assert_eq; +use ruff_text_size::TextRange; +use ruff_text_size::TextSize; use crate::state::lsp::AllOffPartial; use crate::state::lsp::InlayHintConfig; @@ -1703,3 +1705,37 @@ class MyClass: // NewType is a callable alias, so `type[N]` is not a valid annotation to insert. assert!(hints[0].edits.is_none()); } + +#[test] +fn test_inlay_hints_in_range_limits_to_the_selection() { + let code = "\ +def f() -> int: + return 1 + +yyy = f() +zzz = f() +"; + let files = [("main", code)]; + let (handles, state) = mk_multi_file_state_assert_no_errors(&files, Require::Exports); + let handle = handles.get("main").unwrap(); + let transaction = state.transaction(); + let all = transaction + .inlay_hints(handle, Default::default(), Default::default()) + .unwrap(); + let start = code.find("zzz").unwrap(); + let zzz = TextRange::new( + TextSize::try_from(start).unwrap(), + TextSize::try_from(start + 3).unwrap(), + ); + let limited = transaction + .inlay_hints_in_range(handle, Default::default(), Default::default(), zzz) + .unwrap(); + assert!(all.len() > limited.len()); + assert_eq!(limited.len(), 1); + let label: String = limited[0] + .label_parts + .iter() + .map(|(text, _)| text.as_str()) + .collect(); + assert_eq!(label, ": int"); +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/notebook_code_action.rs b/pyrefly/lib/test/lsp/lsp_interaction/notebook_code_action.rs index cd9b3b592c..2aec0e26c1 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/notebook_code_action.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/notebook_code_action.rs @@ -144,3 +144,51 @@ fn test_notebook_code_action_implicit_import_cross_cell() { .unwrap(); interaction.shutdown().unwrap(); } + +#[test] +fn test_notebook_code_action_insert_inferred_type() { + let root = get_test_files_root(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction + .initialize(InitializeSettings { + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + + interaction.open_notebook( + "notebook.ipynb", + vec!["x = 1", "def f() -> int:\n return x\nyyy = f()\n"], + ); + + let cell2_uri = interaction.cell_uri("notebook.ipynb", "cell2"); + interaction + .code_action_cell("notebook.ipynb", "cell2", 2, 0, 2, 3) + .expect_response_with(|response| { + let Some(actions) = response else { + return false; + }; + actions.iter().any(|action| { + let CodeActionResponse::CodeAction(code_action) = action else { + return false; + }; + if code_action.title != "Insert inferred type `: int`" { + return false; + } + let Some(text_edits) = code_action + .edit + .as_ref() + .and_then(|edit| edit.changes.as_ref()) + .and_then(|changes| changes.get(&cell2_uri)) + else { + return false; + }; + text_edits + .iter() + .any(|text_edit| text_edit.new_text == ": int") + }) + }) + .unwrap(); + interaction.shutdown().unwrap(); +} From 5911c1ebf31129f7979fc7bd25d708eeff0692c6 Mon Sep 17 00:00:00 2001 From: Jnyandeep Date: Tue, 10 Nov 2026 11:50:50 +0530 Subject: [PATCH 2/3] feat: added maintainers suggestions --- pyrefly/lib/lsp/non_wasm/server.rs | 128 ++++++++---------- pyrefly/lib/state/lsp.rs | 28 ++-- .../quick_fixes/insert_inferred_annotation.rs | 17 +-- pyrefly/lib/test/lsp/code_actions.rs | 15 +- 4 files changed, 87 insertions(+), 101 deletions(-) diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 0db5994024..b26412fae7 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5332,82 +5332,72 @@ impl Server { // index so that import quick-fixes can be redirected to the current cell // instead of always targeting cell 1 (position 0 of the combined AST). let triggered_cell_index = self.maybe_get_code_cell_index(uri); - let mut quickfixes = transaction - .local_quickfix_code_actions_sorted( - &handle, - range, - import_format, - Some(&self.lsp_thread_pool), - ) - .unwrap_or_default(); - let annotation_start = Instant::now(); - let inlay_hint_config = lsp_config.and_then(|c| c.inlay_hints).unwrap_or_default(); - if let Some(annotation_fixes) = transaction.insert_inferred_annotation_code_actions( + if let Some(quickfixes) = transaction.local_quickfix_code_actions_sorted( &handle, range, import_format, - inlay_hint_config, + lsp_config.and_then(|c| c.inlay_hints).unwrap_or_default(), + Some(&self.lsp_thread_pool), ) { - quickfixes.extend(annotation_fixes); - } - record_code_action_telemetry("insert_inferred_annotation", annotation_start); - actions.extend(quickfixes.into_iter().filter_map(|(title, edits)| { - // Fixes for missing `@override` and context-manager annotations - // can include several edits. Group every edit by its target document - // into a single workspace edit - // so the whole fix applies in one action. - let mut changes: HashMap> = HashMap::new(); - for (info, range, insert_text) in edits { - let lsp_location = self.to_lsp_location(&TextRangeWithModule { - module: info.clone(), - range, - })?; - let mut edit_uri = lsp_location.uri; - let mut edit_range = lsp_location.range; - // Import insertions in notebooks go into the active cell. - // Replacements must retain their target cell and range. - if range.is_empty() - && let Some(current_cell_idx) = triggered_cell_index - { - let edit_cell_idx = info.to_cell_for_lsp(range.start()); - if edit_cell_idx != Some(current_cell_idx) { - // Redirect to the current cell, inserting at line 0. - let open_files = self.open_files.read(); - let notebook_path = self.open_notebook_cells.read().get(uri).cloned(); - let cell_url = notebook_path.and_then(|path| { - if let Some(LspFile::Notebook(notebook)) = - open_files.get(&path).map(|f| &**f) - { - notebook.get_code_cell_url(current_cell_idx).cloned() - } else { - None + actions.extend(quickfixes.into_iter().filter_map(|(title, edits)| { + // Fixes for missing `@override` and context-manager annotations + // can include several edits. Group every edit by its target document + // into a single workspace edit + // so the whole fix applies in one action. + let mut changes: HashMap> = HashMap::new(); + for (info, range, insert_text) in edits { + let lsp_location = self.to_lsp_location(&TextRangeWithModule { + module: info.clone(), + range, + })?; + let mut edit_uri = lsp_location.uri; + let mut edit_range = lsp_location.range; + // Import insertions in notebooks go into the active cell. + // Replacements must retain their target cell and range. + if range.is_empty() + && let Some(current_cell_idx) = triggered_cell_index + { + let edit_cell_idx = info.to_cell_for_lsp(range.start()); + if edit_cell_idx != Some(current_cell_idx) { + // Redirect to the current cell, inserting at line 0. + let open_files = self.open_files.read(); + let notebook_path = + self.open_notebook_cells.read().get(uri).cloned(); + let cell_url = notebook_path.and_then(|path| { + if let Some(LspFile::Notebook(notebook)) = + open_files.get(&path).map(|f| &**f) + { + notebook.get_code_cell_url(current_cell_idx).cloned() + } else { + None + } + }); + if let Some(cell_url) = cell_url { + let top_of_cell = lsp_types::Range { + start: lsp_types::Position::new(0, 0), + end: lsp_types::Position::new(0, 0), + }; + edit_uri = cell_url; + edit_range = top_of_cell; } - }); - if let Some(cell_url) = cell_url { - let top_of_cell = lsp_types::Range { - start: lsp_types::Position::new(0, 0), - end: lsp_types::Position::new(0, 0), - }; - edit_uri = cell_url; - edit_range = top_of_cell; } - } - }; - changes.entry(edit_uri).or_default().push(TextEdit { - range: edit_range, - new_text: insert_text, - }); - } - Some(CodeActionResponse::CodeAction(CodeAction { - title, - kind: Some(CodeActionKind::QuickFix), - edit: Some(WorkspaceEdit { - changes: Some(changes), + }; + changes.entry(edit_uri).or_default().push(TextEdit { + range: edit_range, + new_text: insert_text, + }); + } + Some(CodeActionResponse::CodeAction(CodeAction { + title, + kind: Some(CodeActionKind::QuickFix), + edit: Some(WorkspaceEdit { + changes: Some(changes), + ..Default::default() + }), ..Default::default() - }), - ..Default::default() - })) - })); + })) + })); + } record_code_action_telemetry("quickfix", start); } if allow_fix_all { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index e5457719a2..e7ebf15f7a 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3716,6 +3716,7 @@ impl<'a> Transaction<'a> { handle: &Handle, range: TextRange, import_format: ImportFormat, + inlay_hint_config: InlayHintConfig, custom_thread_pool: Option<&ThreadPool>, ) -> Option)>> { let module_info = self.get_module_info(handle)?; @@ -3990,6 +3991,17 @@ impl<'a> Transaction<'a> { _ => {} } } + if let Some(actions) = + quick_fixes::insert_inferred_annotation::insert_inferred_annotation_code_actions( + self, + handle, + range, + import_format, + inlay_hint_config, + ) + { + multi_actions.extend(actions); + } if let Some(answers) = self.get_answers(handle) { let bindings = answers.bindings(); for unused in bindings.unused_variables() { @@ -4211,22 +4223,6 @@ impl<'a> Transaction<'a> { } } - pub fn insert_inferred_annotation_code_actions( - &self, - handle: &Handle, - selection: TextRange, - import_format: ImportFormat, - inlay_hint_config: InlayHintConfig, - ) -> Option)>> { - quick_fixes::insert_inferred_annotation::insert_inferred_annotation_code_actions( - self, - handle, - selection, - import_format, - inlay_hint_config, - ) - } - pub fn pytest_fixture_type_annotation_code_actions( &self, handle: &Handle, diff --git a/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs b/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs index 2fdfa640db..75e3e8dc9e 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs @@ -10,7 +10,6 @@ use pyrefly_build::handle::Handle; use pyrefly_python::module::Module; use ruff_text_size::TextRange; -use crate::error::suppress::detect_line_ending; use crate::state::lsp::AllOffPartial; use crate::state::lsp::ImportFormat; use crate::state::lsp::InlayHintConfig; @@ -35,14 +34,6 @@ pub(crate) fn insert_inferred_annotation_code_actions( let hints = transaction.inlay_hints_in_range(handle, inlay_hint_config, import_format, selection)?; let module = transaction.get_module_info(handle)?; - let line_ending = detect_line_ending(module.contents().as_str()); - let normalize = |text: String| -> String { - if line_ending == "\n" { - text - } else { - text.replace('\n', line_ending) - } - }; let mut actions = Vec::new(); for hint in hints { let Some(edits) = hint.edits else { @@ -53,14 +44,10 @@ pub(crate) fn insert_inferred_annotation_code_actions( action_edits.push(( module.dupe(), TextRange::empty(hint.position), - normalize(edits.annotation), + edits.annotation, )); for (offset, import_text) in edits.imports { - action_edits.push(( - module.dupe(), - TextRange::empty(offset), - normalize(import_text), - )); + action_edits.push((module.dupe(), TextRange::empty(offset), import_text)); } actions.push((hint.position, title, action_edits)); } diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 40f3d8a9a3..81051c3036 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -56,6 +56,7 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default() @@ -975,6 +976,7 @@ fn code_action_imports_with_future_imports() { &handles["main"], TextRange::empty(TextSize::try_from(header.len()).unwrap()), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap(); @@ -1140,6 +1142,7 @@ fn insertion_test_implicit_import() { handle, TextRange::empty(position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1181,6 +1184,7 @@ fn insertion_test_implicit_import() { handle, TextRange::empty(TextSize::new(0)), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1228,6 +1232,7 @@ np handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1260,6 +1265,7 @@ fn insert_import_uses_file_line_ending() { handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1621,6 +1627,7 @@ fn apply_contextmanager_fix(code: &str, minor: u32) -> Option { handle, TextRange::empty(position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -2312,6 +2319,7 @@ fn remove_unused_ignore_results( &handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default() @@ -2737,6 +2745,7 @@ fn unused_import_action_after_with_env( handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -3058,6 +3067,7 @@ fn redundant_cast_action_after(code: &str, cursor_offset: usize) -> Option Option<(St handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -3447,6 +3458,7 @@ fn test_import_for_unimported_directives() { handle, TextRange::new(TextSize::new(0), TextSize::new(0)), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -8105,6 +8117,7 @@ fn add_override_quickfix_edits( &handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default() @@ -8267,7 +8280,7 @@ fn inferred_annotation_actions( let handle = handles.get("main").unwrap(); let transaction = state.transaction(); transaction - .insert_inferred_annotation_code_actions(handle, range, ImportFormat::Absolute, config) + .local_quickfix_code_actions_sorted(handle, range, ImportFormat::Absolute, config, None) .unwrap_or_default() .into_iter() .map(|(title, edits)| { From dd3128d905849b55734dcfda1bcce6904becf5fa Mon Sep 17 00:00:00 2001 From: Jnyandeep Date: Sat, 14 Nov 2026 15:12:34 +0530 Subject: [PATCH 3/3] fix: ci fixed --- crates/pyrefly_config/src/config.rs | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/crates/pyrefly_config/src/config.rs b/crates/pyrefly_config/src/config.rs index 64e04a57e7..a60bfdee74 100644 --- a/crates/pyrefly_config/src/config.rs +++ b/crates/pyrefly_config/src/config.rs @@ -4673,7 +4673,7 @@ output-format = "omit-errors" // File contents should still be relative to the location of the config file, not src/. assert_eq!( config.project_includes, - Globs::new(vec![root.path().join("**/*").to_string_lossy().to_string()]).unwrap(), + Globs::new_with_root(root.path(), vec!["**/*".to_owned()]).unwrap(), ); assert_eq!( config.search_path().cloned().collect::>(), @@ -5494,15 +5494,14 @@ matches = "generated/**" let (config, errors) = ConfigFile::from_file(&child); assert!(errors.is_empty()); - let shared_glob = |pattern: &str| { - Globs::new(vec![shared.join(pattern).to_string_lossy().into_owned()]).unwrap() - }; + let shared_glob = + |pattern: &str| Globs::new_with_root(&shared, vec![pattern.to_owned()]).unwrap(); assert_eq!(config.project_includes, shared_glob("base/**/*.py")); assert_eq!(config.search_path_from_file, vec![shared.join("stubs")]); assert_eq!(config.coverage.includes, Some(shared_glob("covered"))); assert_eq!( config.sub_configs[0].matches, - Glob::new(shared.join("generated/**").to_string_lossy().into_owned()).unwrap() + Glob::new_with_root(&shared, "generated/**".to_owned()).unwrap() ); assert_eq!( config.python_environment.site_package_path, @@ -5538,13 +5537,7 @@ search-path = [] assert!(errors.is_empty()); assert_eq!( config.project_includes, - Globs::new(vec![ - root.path() - .join("child/**/*.py") - .to_string_lossy() - .to_string() - ]) - .unwrap() + Globs::new_with_root(root.path(), vec!["child/**/*.py".to_owned()]).unwrap() ); assert!(config.search_path_from_file.is_empty()); }