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()); } diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index ff1c941539..90c4fc2aab 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5312,6 +5312,7 @@ impl Server { &handle, range, import_format, + lsp_config.and_then(|c| c.inlay_hints).unwrap_or_default(), Some(&self.lsp_thread_pool), ) { actions.extend(quickfixes.into_iter().filter_map(|(title, edits)| { diff --git a/pyrefly/lib/lsp/wasm/inlay_hints.rs b/pyrefly/lib/lsp/wasm/inlay_hints.rs index 26dd61ffb8..b1f000b36b 100644 --- a/pyrefly/lib/lsp/wasm/inlay_hints.rs +++ b/pyrefly/lib/lsp/wasm/inlay_hints.rs @@ -21,6 +21,7 @@ use ruff_python_ast::ExprDict; use ruff_python_ast::ExprList; use ruff_python_ast::ModModule; use ruff_python_ast::ParameterWithDefault; +use ruff_python_ast::Stmt; use ruff_python_ast::name::Name; use ruff_text_size::Ranged; use ruff_text_size::TextRange; @@ -118,6 +119,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 +211,49 @@ impl<'a> Transaction<'a> { edits, } }; + // `None` keeps every hint. An empty cursor includes the insertion point at + // the end of the declaration. A nonempty selection uses half-open overlap, + // so a declaration that only touches the selection's exclusive end is left out. + let in_selection = |target: TextRange| match limit { + None => true, + Some(selection) if selection.is_empty() => { + selection.start() >= target.start() && selection.start() <= target.end() + } + 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); @@ -247,7 +281,14 @@ impl<'a> Transaction<'a> { // receiver case, could mislead users into // annotating with the RHS-derived type. The same // applies to receiver-bearing unpacked rebinds. - Binding::NameAssign(x) if !x.is_pinned() => (Some(&*x.expr), false), + Binding::NameAssign(x) if !x.is_pinned() => { + // `global` and `nonlocal` names cannot take an annotated assignment. + if self.capture_forbids_annotation(handle, key.range().start(), &x.name) + { + insertable = false; + } + (Some(&*x.expr), false) + } Binding::Expr(None, e) => { // Expression-bound names, such as walrus targets, cannot be annotated inline. insertable = false; @@ -303,6 +344,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, @@ -354,6 +398,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)], @@ -368,6 +413,33 @@ impl<'a> Transaction<'a> { Some(res) } + /// `global` and `nonlocal` names cannot take an annotated assignment in that scope. + fn capture_forbids_annotation(&self, handle: &Handle, position: TextSize, name: &Name) -> bool { + fn declared(stmts: &[Stmt], position: TextSize, name: &Name) -> Option { + for stmt in stmts { + let body = match stmt { + Stmt::FunctionDef(func) if func.range().contains(position) => { + func.body.as_slice() + } + Stmt::ClassDef(class) if class.range().contains(position) => { + class.body.as_slice() + } + _ => continue, + }; + return Some(declared(body, position, name).unwrap_or_else(|| { + body.iter().any(|stmt| match stmt { + Stmt::Global(global) => global.names.iter().any(|id| id.id == *name), + Stmt::Nonlocal(nonlocal) => nonlocal.names.iter().any(|id| id.id == *name), + _ => false, + }) + })); + } + None + } + self.get_ast(handle) + .is_some_and(|ast| declared(&ast.body, position, name).unwrap_or(false)) + } + /// Helper to extract the element expression from an unpacked source. /// Returns the expression at the given position if the source is a tuple or list literal. /// For nested unpacking or function calls, returns None (caller should fall back to diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index ca5ff3723a..8434606381 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3805,6 +3805,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)?; @@ -4079,6 +4080,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() { 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..75e3e8dc9e --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/insert_inferred_annotation.rs @@ -0,0 +1,64 @@ +/* + * 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::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 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), + edits.annotation, + )); + for (offset, import_text) in edits.imports { + action_edits.push((module.dupe(), TextRange::empty(offset), 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 d0bfad0d0d..26fca6944d 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; @@ -54,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() @@ -869,6 +872,7 @@ fn insert_import_skips_module_name_with_non_identifier() { &handles["main"], TextRange::empty(TextSize::new(0)), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default() @@ -1004,6 +1008,7 @@ fn code_action_imports_with_future_imports() { &handles["main"], TextRange::empty(TextSize::try_from(header.len()).unwrap()), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap(); @@ -1169,6 +1174,7 @@ fn insertion_test_implicit_import() { handle, TextRange::empty(position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1210,6 +1216,7 @@ fn insertion_test_implicit_import() { handle, TextRange::empty(TextSize::new(0)), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1257,6 +1264,7 @@ np handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1289,6 +1297,7 @@ fn insert_import_uses_file_line_ending() { handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -1650,6 +1659,7 @@ fn apply_contextmanager_fix(code: &str, minor: u32) -> Option { handle, TextRange::empty(position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -2341,6 +2351,7 @@ fn remove_unused_ignore_results( &handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default() @@ -2766,6 +2777,7 @@ fn unused_import_action_after_with_env( handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -3087,6 +3099,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(); @@ -3476,6 +3490,7 @@ fn test_import_for_unimported_directives() { handle, TextRange::new(TextSize::new(0), TextSize::new(0)), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default(); @@ -8134,6 +8149,7 @@ fn add_override_quickfix_edits( &handle, TextRange::new(position, position), ImportFormat::Absolute, + InlayHintConfig::default(), None, ) .unwrap_or_default() @@ -8286,3 +8302,376 @@ 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 + .local_quickfix_code_actions_sorted(handle, range, ImportFormat::Absolute, config, None) + .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")); +} + +#[test] +fn insert_inferred_type_ignores_declaration_at_selection_end() { + 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(); + 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.len(), 1); + assert!(actions[0].1.contains("yyy: int = grow(1)")); + assert!(!actions[0].1.contains("zzz: int")); +} + +#[test] +fn insert_inferred_type_skips_walrus_and_captured_writes() { + let walrus = "\ +def make() -> int: + return 1 + +if (assigned := make()): + print(assigned) +"; + assert!( + inferred_annotation_actions( + &[("main", walrus)], + range_of(walrus, "assigned :="), + Default::default(), + ) + .is_empty() + ); + + let global = "\ +def make() -> int: + return 1 + +def write() -> None: + global item + item = make() +"; + assert!( + inferred_annotation_actions( + &[("main", global)], + range_of(global, "item = make()"), + Default::default(), + ) + .is_empty() + ); + + let nonlocal_write = "\ +def make() -> int: + return 1 + +def outer() -> None: + item = 1 + def write() -> None: + nonlocal item + item = make() +"; + assert!( + inferred_annotation_actions( + &[("main", nonlocal_write)], + range_of(nonlocal_write, "item = make()"), + Default::default(), + ) + .is_empty() + ); +} diff --git a/pyrefly/lib/test/lsp/inlay_hint.rs b/pyrefly/lib/test/lsp/inlay_hint.rs index e1fbf5e051..d6c7a4f018 100644 --- a/pyrefly/lib/test/lsp/inlay_hint.rs +++ b/pyrefly/lib/test/lsp/inlay_hint.rs @@ -6,6 +6,7 @@ */ use pretty_assertions::assert_eq; +use ruff_text_size::TextRange; use ruff_text_size::TextSize; use crate::state::lsp::AllOffPartial; @@ -1746,3 +1747,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 cfb641206a..8416a0ac6a 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/notebook_code_action.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/notebook_code_action.rs @@ -145,6 +145,54 @@ fn test_notebook_code_action_implicit_import_cross_cell() { 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(); +} + #[test] fn test_notebook_code_action_import_cell_header() { let root = get_test_files_root();