diff --git a/crates/pyrefly_types/src/display.rs b/crates/pyrefly_types/src/display.rs index ee9614e37c..cdc0973470 100644 --- a/crates/pyrefly_types/src/display.rs +++ b/crates/pyrefly_types/src/display.rs @@ -171,6 +171,7 @@ pub struct TypeDisplayContext<'a> { /// Render a `DataFrame` as its plain underlying class rather than the schema form. strip_library_schemas: bool, render_unknown_as_any: bool, + render_source_annotation: bool, /// Optional stdlib reference for resolving builtin type locations stdlib: Option<&'a Stdlib>, /// Stack of identities of type variables currently bound by enclosing Foralls. @@ -342,6 +343,7 @@ impl<'a> TypeDisplayContext<'a> { fn fmt_quantified(&self, quantified: &Quantified, output: &mut impl TypeOutput) -> fmt::Result { write!(output, "{}", quantified.name)?; if self.always_display_module_name + && !self.render_source_annotation && !self .forall_tparam_uniques .borrow() @@ -920,6 +922,60 @@ impl<'a> TypeDisplayContext<'a> { is_toplevel: bool, output: &mut impl TypeOutput, ) -> fmt::Result { + if self.render_source_annotation + && let Some(metadata) = t.toplevel_func_metadata() + && let FunctionKind::CallbackProtocol(cls) = &metadata.kind + { + output.write_qname(cls.qname())?; + return output.write_targs(cls.targs()); + } + if self.render_source_annotation + && let Some((callable, _)) = t.toplevel_callable_signatures().next() + { + let bound_signature = if matches!(t, Type::BoundMethod(_)) { + callable.strip_first_param() + } else { + None + }; + let callable = bound_signature.as_ref().unwrap_or(callable); + self.maybe_fmt_with_module("typing", "Callable", output)?; + output.write_str("[")?; + match &callable.params { + Params::List(params) + if params + .items() + .iter() + .all(|p| matches!(p, Param::PosOnly(..) | Param::Pos(..))) => + { + output.write_str("[")?; + self.fmt_type_sequence( + params.items().iter().map(|p| p.as_type()), + ", ", + false, + output, + )?; + output.write_str("]")?; + } + Params::ParamSpec(prefix, param_spec) => { + if prefix.is_empty() { + self.fmt_helper_generic(param_spec, false, output)?; + } else { + self.maybe_fmt_with_module("typing", "Concatenate", output)?; + output.write_str("[")?; + for param in prefix { + self.fmt_helper_generic(param.ty(), false, output)?; + output.write_str(", ")?; + } + self.fmt_helper_generic(param_spec, false, output)?; + output.write_str("]")?; + } + } + _ => output.write_str("...")?, + } + output.write_str(", ")?; + self.fmt_helper_generic(&callable.ret, false, output)?; + return output.write_str("]"); + } match t { // Things that have QName's and need qualifying Type::ClassDef(cls) => { @@ -1638,6 +1694,12 @@ impl<'a> TypeDisplayContext<'a> { Type::Quantified(var) => self.fmt_quantified(var, output), Type::QuantifiedValue(var) => write!(output, "{var}"), Type::ElementOfTypeVarTuple(var) => write!(output, "ElementOf[{var}]"), + Type::Args(q) if self.render_source_annotation => { + write!(output, "{}.args", q.name()) + } + Type::Kwargs(q) if self.render_source_annotation => { + write!(output, "{}.kwargs", q.name()) + } Type::Args(q) => { output.write_str("Args[")?; write!(output, "{q}")?; @@ -1857,6 +1919,7 @@ impl Type { pub fn get_annotation_parts(&self, stdlib: Option<&Stdlib>) -> Vec { let mut ctx = annotation_context(self, stdlib); ctx.render_unknown_as_any = true; + ctx.render_source_annotation = true; ctx.always_display_module_name_except_builtins(); let mut output = AnnotationOutput::new(&ctx); ctx.fmt_helper_generic(self, false, &mut output).unwrap(); @@ -3617,6 +3680,52 @@ def overloaded_func[T]( assert!(parts.iter().any(|(_, loc)| loc.is_some())); } + #[test] + fn source_annotation_renders_nested_callables() { + let inner = Type::Callable(Box::new(Callable::list( + ParamList::new(vec![Param::PosOnly(None, Type::None, Required::Required)]), + Type::any_implicit(), + ))); + let ty = Type::Callable(Box::new(Callable::list( + ParamList::new(vec![Param::PosOnly(None, inner, Required::Required)]), + Type::None, + ))); + assert_eq!(ty.to_string(), "((None) -> Unknown) -> None"); + assert_eq!( + ty.get_annotation_parts(None), + vec![ + AnnotationPart::Reference { + module: ModuleName::from_str("typing"), + name: "Callable".to_owned() + }, + AnnotationPart::Text("[[".to_owned()), + AnnotationPart::Reference { + module: ModuleName::from_str("typing"), + name: "Callable".to_owned() + }, + AnnotationPart::Text("[[None], ".to_owned()), + AnnotationPart::Reference { + module: ModuleName::from_str("typing"), + name: "Any".to_owned() + }, + AnnotationPart::Text("]], None]".to_owned()), + ] + ); + } + + #[test] + fn source_annotation_omits_bound_receiver() { + let annotation = fake_bound_method("method", "Class", "module") + .get_annotation_parts(None) + .into_iter() + .map(|part| match part { + AnnotationPart::Text(text) => text, + AnnotationPart::Reference { name, .. } => name, + }) + .collect::(); + assert_eq!(annotation, "Callable[[Any, Any], None]"); + } + #[test] fn source_annotation_surfaces_render_self_as_self() { let cls = fake_class("MyClass", "mymodule", 40); diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 0f29a4100d..f424a38d70 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3691,6 +3691,13 @@ impl<'a> Transaction<'a> { // Actions such as missing-`@override` and context-manager annotations // carry multiple edits that must apply together. let mut multi_actions: Vec<(String, Vec<(Module, TextRange, String)>)> = Vec::new(); + if let Some(actions) = self.implement_abstract_members_code_actions(handle, range) { + multi_actions.extend( + actions + .into_iter() + .map(|action| (action.title, action.edits)), + ); + } // Deduplicates the actions pushed onto `other_actions`. The same fix can // be generated more than once -- e.g. several errors on one line all // produce an identical "add pyrefly ignore" action, or one unused binding @@ -4195,6 +4202,16 @@ impl<'a> Transaction<'a> { quick_fixes::extract_superclass::extract_superclass_code_actions(self, handle, selection) } + pub fn implement_abstract_members_code_actions( + &self, + handle: &Handle, + selection: TextRange, + ) -> Option> { + quick_fixes::implement_abstract_members::implement_abstract_members_code_actions( + self, handle, selection, + ) + } + pub fn pull_members_up_code_actions( &self, handle: &Handle, diff --git a/pyrefly/lib/state/lsp/quick_fixes.rs b/pyrefly/lib/state/lsp/quick_fixes.rs index 5bbb87fc63..af57ed4feb 100644 --- a/pyrefly/lib/state/lsp/quick_fixes.rs +++ b/pyrefly/lib/state/lsp/quick_fixes.rs @@ -18,6 +18,7 @@ mod extract_shared; pub(crate) mod extract_superclass; pub(crate) mod extract_variable; pub(crate) mod generate_code; +pub(crate) mod implement_abstract_members; pub(crate) mod implicit_import; pub(crate) mod inline_method; pub(crate) mod inline_parameter; diff --git a/pyrefly/lib/state/lsp/quick_fixes/implement_abstract_members.rs b/pyrefly/lib/state/lsp/quick_fixes/implement_abstract_members.rs new file mode 100644 index 0000000000..00f17cd949 --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/implement_abstract_members.rs @@ -0,0 +1,587 @@ +/* + * 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 lsp_types::CodeActionKind; +use pyrefly_build::handle::Handle; +use pyrefly_python::docstring::dedent_block_preserving_layout; +use pyrefly_python::module::Module; +use pyrefly_python::module_name::ModuleName; +use pyrefly_python::short_identifier::ShortIdentifier; +use pyrefly_types::callable::Callable; +use pyrefly_types::callable::DefaultValue; +use pyrefly_types::callable::Param; +use pyrefly_types::callable::ParamList; +use pyrefly_types::callable::Params; +use pyrefly_types::callable::Required; +use pyrefly_types::literal::Lit; +use pyrefly_types::stdlib::Stdlib; +use pyrefly_types::type_output::AnnotationPart; +use pyrefly_types::type_var::Restriction; +use pyrefly_types::types::TParams; +use pyrefly_types::types::Type; +use ruff_python_ast::ModModule; +use ruff_python_ast::Stmt; +use ruff_python_ast::StmtClassDef; +use ruff_python_ast::helpers::is_docstring_stmt; +use ruff_python_ast::name::Name; +use ruff_text_size::Ranged; +use ruff_text_size::TextRange; +use ruff_text_size::TextSize; +use starlark_map::Hashed; + +use super::extract_shared::line_indent_and_start; +use super::extract_shared::selection_anchor; +use crate::binding::binding::KeyClass; +use crate::state::ide::import_regular_import_edit; +use crate::state::ide::insert_import_edit; +use crate::state::import_tracker::ImportTracker; +use crate::state::lsp::ImportFormat; +use crate::state::lsp::LocalRefactorCodeAction; +use crate::state::lsp::Transaction; + +const DEFAULT_INDENT: &str = " "; + +#[derive(Clone)] +struct AbstractMemberInfo { + name: Name, + ty: Type, + defining_module: Module, + docstring_range: Option, + is_class_var: bool, +} + +/// Builds a quick fix that implements all inherited abstract members in a class. +pub(crate) fn implement_abstract_members_code_actions( + 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 class_def = find_class_context(ast.as_ref(), selection_point)?; + let (method_indent, insert_range, prefix) = insertion_point(class_def, source)?; + let indent_unit = class_indent_unit(source, class_def, &method_indent); + + let members = transaction + .ad_hoc_solve( + handle, + "implement_abstract_members", + |solver| -> Option> { + let key = KeyClass(ShortIdentifier::new(&class_def.name)); + let class_idx = solver.bindings().key_to_idx_hashed_opt(Hashed::new(&key))?; + let class = solver.get_idx(class_idx).0.clone()?; + let abstract_members = solver.get_abstract_members_for_class(&class); + let mut infos = Vec::new(); + for name in abstract_members.unimplemented_abstract_methods() { + let member = solver.get_class_member_with_defining_class(&class, name)?; + if member.defining_class == class { + continue; + } + let mut ty = member.value.ty(); + if ty + .toplevel_func_metadata() + .is_some_and(|m| m.flags.is_async) + { + ty.transform_toplevel_callable_signatures(|callable, _| { + if let Some((_, _, ret)) = solver.unwrap_coroutine(&callable.ret) { + callable.ret = ret; + } + }); + } + infos.push(AbstractMemberInfo { + name: name.clone(), + ty, + is_class_var: member.value.is_class_var(), + defining_module: member.defining_class.module().dupe(), + docstring_range: solver + .get_class_fields(&member.defining_class) + .and_then(|fields| fields.field_docstring_range(name)), + }); + } + Some(infos) + }, + ) + .flatten()?; + + if members.is_empty() { + return None; + } + + let mut members = members; + members.sort_by(|left, right| left.name.cmp(&right.name)); + let stdlib = transaction.get_stdlib(handle); + let mut renderer = MemberRenderer { + stdlib: &stdlib, + parts: vec![AnnotationPart::Text(prefix)], + }; + for (i, member) in members.iter().enumerate() { + if i > 0 { + renderer.text("\n\n"); + } + renderer.member(member, &method_indent, &indent_unit)?; + } + renderer.text("\n"); + let parts = renderer.parts; + + let tracker = ImportTracker::from_ast( + &ast, + handle.module(), + handle.path().is_init(), + handle.sys_info().dupe(), + ); + let direct = tracker.direct_import_candidates(&parts, handle.module()); + let direct_refs = direct + .iter() + .map(|(module, name)| (*module, name.as_str())) + .collect::>(); + let (combined, missing) = tracker.resolve_annotation(&parts, handle.module(), &direct_refs); + let mut imports = String::new(); + for (module, name) in direct { + let import_handle = transaction.import_handle(handle, module, None).finding()?; + let edit = insert_import_edit( + &ast, + transaction.config_finder(), + handle.dupe(), + import_handle, + &name, + ImportFormat::Absolute, + ); + imports.push_str(&edit.insert_text); + } + let mut missing = missing.into_iter().collect::>(); + missing.sort_by(|a, b| a.as_str().cmp(b.as_str())); + for module in missing { + let import_handle = transaction.import_handle(handle, module, None).finding()?; + let (_, text, _) = import_regular_import_edit(&ast, import_handle, None); + imports.push_str(&text); + } + let mut edits = vec![(module_info.dupe(), insert_range, combined)]; + if !imports.is_empty() { + // Future imports must precede every generated import. + let position = ast.body.iter().find(|stmt| { + !is_docstring_stmt(stmt) + && !matches!(stmt, Stmt::ImportFrom(import) if import.module.as_ref().is_some_and(|m| m.as_str() == "__future__")) + })?.range().start(); + edits.push((module_info.dupe(), TextRange::empty(position), imports)); + } + + Some(vec![LocalRefactorCodeAction { + title: "Implement abstract members".to_owned(), + edits, + kind: CodeActionKind::QuickFix, + }]) +} + +fn find_class_context<'a>(ast: &'a ModModule, selection: TextSize) -> Option<&'a StmtClassDef> { + for stmt in &ast.body { + if let Stmt::ClassDef(class_def) = stmt + && class_def.range().contains(selection) + { + if let Some(inner) = find_class_in_body(&class_def.body, selection) { + return Some(inner); + } + return Some(class_def); + } + } + None +} + +fn find_class_in_body<'a>(body: &'a [Stmt], selection: TextSize) -> Option<&'a StmtClassDef> { + for stmt in body { + if let Stmt::ClassDef(class_def) = stmt + && class_def.range().contains(selection) + { + if let Some(inner) = find_class_in_body(&class_def.body, selection) { + return Some(inner); + } + return Some(class_def); + } + } + None +} + +fn insertion_point(class_def: &StmtClassDef, source: &str) -> Option<(String, TextRange, String)> { + let first = class_def.body.first()?; + let (class_indent, _) = line_indent_and_start(source, class_def.range().start())?; + let (_, body_line_start) = line_indent_and_start(source, first.start())?; + if !source[body_line_start.to_usize()..first.start().to_usize()] + .trim() + .is_empty() + { + let start = source[..first.start().to_usize()].rfind(':')? + 1; + let end = line_end_position(source, class_def.end()); + let indent = format!("{class_indent}{DEFAULT_INDENT}"); + let mut prefix = "\n".to_owned(); + if replaceable_pass_stmt(class_def).is_none() || class_def.body.len() > 1 { + prefix.push_str(&indent); + prefix.push_str( + &source[first.start().to_usize()..end.to_usize()] + .trim_end_matches(['\r', '\n']) + .replace("\r\n", "\n"), + ); + prefix.push('\n'); + } + return Some(( + indent, + TextRange::new(TextSize::try_from(start).ok()?, end), + prefix, + )); + } + if let Some(pass_stmt) = replaceable_pass_stmt(class_def) { + let (indent, line_start) = line_indent_and_start(source, pass_stmt.range().start())?; + if source[line_start.to_usize()..pass_stmt.start().to_usize()] + .trim() + .is_empty() + { + let line_end = line_end_position(source, pass_stmt.range().end()); + return Some((indent, TextRange::new(line_start, line_end), String::new())); + } + } + for stmt in &class_def.body { + if is_docstring_stmt(stmt) { + continue; + } + let anchor = match stmt { + Stmt::FunctionDef(function) => function + .decorator_list + .first() + .map_or(stmt.start(), |d| d.start()), + Stmt::ClassDef(class) => class + .decorator_list + .first() + .map_or(stmt.start(), |d| d.start()), + _ => stmt.start(), + }; + let (indent, line_start) = line_indent_and_start(source, anchor)?; + if !source[line_start.to_usize()..anchor.to_usize()] + .trim() + .is_empty() + { + continue; + } + return Some((indent, TextRange::empty(line_start), String::new())); + } + let (indent, _) = line_indent_and_start(source, first.start())?; + let line_end = line_end_position(source, class_def.end()); + let prefix = if source[..line_end.to_usize()].ends_with('\n') { + "" + } else { + "\n" + }; + Some((indent, TextRange::empty(line_end), prefix.to_owned())) +} + +fn replaceable_pass_stmt(class_def: &StmtClassDef) -> Option<&Stmt> { + let mut non_docstring = class_def + .body + .iter() + .filter(|stmt| !is_docstring_stmt(stmt)); + let only_stmt = non_docstring.next()?; + if non_docstring.next().is_some() { + return None; + } + match only_stmt { + Stmt::Pass(_) => Some(only_stmt), + _ => None, + } +} + +fn class_indent_unit(source: &str, class_def: &StmtClassDef, method_indent: &str) -> String { + let (class_indent, _) = match line_indent_and_start(source, class_def.range().start()) { + Some(value) => value, + None => return DEFAULT_INDENT.to_owned(), + }; + if let Some(unit) = method_indent.strip_prefix(&class_indent) + && !unit.is_empty() + { + return unit.to_owned(); + } + DEFAULT_INDENT.to_owned() +} + +fn line_end_position(source: &str, position: TextSize) -> TextSize { + let mut idx = position.to_usize().min(source.len()); + let bytes = source.as_bytes(); + while idx < bytes.len() && bytes[idx] != b'\n' { + idx += 1; + } + if idx < bytes.len() { + idx += 1; + } + TextSize::try_from(idx).expect("The source length fits in TextSize.") +} + +/// Builds source text while retaining the module references used by annotations. +struct MemberRenderer<'a> { + stdlib: &'a Stdlib, + parts: Vec, +} + +impl MemberRenderer<'_> { + fn text(&mut self, text: &str) { + if let Some(AnnotationPart::Text(previous)) = self.parts.last_mut() { + previous.push_str(text); + } else { + self.parts.push(AnnotationPart::Text(text.to_owned())); + } + } + + fn reference(&mut self, module: &str, name: &str) { + self.parts.push(AnnotationPart::Reference { + module: ModuleName::from_str(module), + name: name.to_owned(), + }); + } + + fn annotation(&mut self, ty: &Type) { + self.parts + .extend(ty.get_annotation_parts(Some(self.stdlib))); + } + + fn member( + &mut self, + member: &AbstractMemberInfo, + indent: &str, + indent_unit: &str, + ) -> Option<()> { + let Some(metadata) = member + .ty + .toplevel_func_metadata() + .filter(|_| !member.is_class_var) + else { + self.text(indent); + self.text(member.name.as_str()); + self.text(": "); + if member.is_class_var { + self.reference("typing", "ClassVar"); + self.text("["); + } + self.annotation(&member.ty); + if member.is_class_var { + self.text("]"); + } + return Some(()); + }; + let property = member.ty.property_metadata(); + let getter = property.map(|p| &p.getter); + let setter = property.and_then(|p| p.setter.as_ref()); + let is_cached = metadata.flags.is_cached_property; + let is_static = metadata.flags.is_staticmethod; + let is_class = metadata.flags.is_classmethod; + let is_async = metadata.flags.is_async; + let overloaded = member.ty.toplevel_callable_signatures().count() > 1; + for (index, (callable, tparams)) in getter + .unwrap_or(&member.ty) + .toplevel_callable_signatures() + .enumerate() + { + if index > 0 { + self.text("\n\n"); + } + if overloaded { + self.text(indent); + self.text("@"); + self.reference("typing", "overload"); + self.text("\n"); + } + if getter.is_some() || is_cached { + self.text(indent); + self.text("@"); + if is_cached { + self.reference("functools", "cached_property"); + } else { + self.text("property"); + } + self.text("\n"); + } else if is_class || is_static { + self.text(indent); + self.text(if is_class { + "@classmethod\n" + } else { + "@staticmethod\n" + }); + } + self.text(indent); + self.text(if is_async { "async def " } else { "def " }); + self.text(member.name.as_str()); + self.signature(callable, tparams.map(|p| p.as_ref()), !is_static)?; + self.text(":\n"); + if !overloaded { + self.body(member, indent, indent_unit); + } else { + self.text(&format!("{indent}{indent_unit}...")); + } + } + if overloaded { + // The overloads define the contract. The implementation accepts every overload. + self.text("\n"); + if is_class || is_static { + self.text(&format!( + "\n\n{indent}@{}", + if is_class { + "classmethod" + } else { + "staticmethod" + } + )); + } + self.text(&format!( + "\n{indent}{}def {}(", + if is_async { "async " } else { "" }, + member.name + )); + if !is_static { + self.text(if is_class { "cls, " } else { "self, " }); + } + self.text("*args, **kwargs) -> "); + self.reference("typing", "Any"); + self.text(":\n"); + self.body(member, indent, indent_unit); + } + if let Some(setter) = setter { + let (callable, tparams) = setter.toplevel_callable_signatures().next()?; + self.text(&format!( + "\n\n{indent}@{}.setter\n{indent}def {}", + member.name, member.name + )); + self.signature(callable, tparams.map(|p| p.as_ref()), true)?; + self.text(":\n"); + self.body(member, indent, indent_unit); + } + Some(()) + } + + fn body(&mut self, member: &AbstractMemberInfo, indent: &str, indent_unit: &str) { + let body_indent = format!("{indent}{indent_unit}"); + if let Some(range) = member.docstring_range + && let Some(docstring) = + dedent_block_preserving_layout(member.defining_module.code_at(range)) + { + for line in docstring.trim_end().lines() { + self.text(&body_indent); + self.text(line); + self.text("\n"); + } + } + self.text(&body_indent); + self.text("raise NotImplementedError"); + } + + fn signature( + &mut self, + callable: &Callable, + tparams: Option<&TParams>, + strip_receiver: bool, + ) -> Option<()> { + if let Some(tparams) = tparams.filter(|p| !p.is_empty()) { + self.text("["); + for (i, param) in tparams.iter().enumerate() { + if i > 0 { + self.text(", "); + } + self.text(¶m.display_name_with_prefix().to_string()); + match param.restriction() { + Restriction::Bound(ty) => { + self.text(": "); + self.annotation(ty); + } + Restriction::Constraints(types) => { + self.text(": ("); + for (i, ty) in types.iter().enumerate() { + if i > 0 { + self.text(", "); + } + self.annotation(ty); + } + self.text(")"); + } + Restriction::Unrestricted => {} + Restriction::ShapeExtension(_) => return None, + } + if let Some(default) = param.default() { + self.text(" = "); + self.annotation(default); + } + } + self.text("]"); + } + self.text("("); + let Params::List(params) = &callable.params else { + return None; + }; + self.params(params, strip_receiver); + self.text(") -> "); + self.annotation(&callable.ret); + Some(()) + } + + fn params(&mut self, params: &ParamList, strip_receiver: bool) { + let mut posonly = false; + let mut kwonly = false; + for (i, param) in params.items().iter().enumerate() { + if i > 0 { + self.text(", "); + } + if matches!(param, Param::PosOnly(..)) { + posonly = true; + } else if posonly { + self.text("/, "); + posonly = false; + } + if matches!(param, Param::Varargs(..)) { + kwonly = true; + self.text("*"); + } else if matches!(param, Param::Kwargs(..)) { + self.text("**"); + } else if matches!(param, Param::KwOnly(..)) && !kwonly { + kwonly = true; + self.text("*, "); + } + if let Some(name) = param.name() { + self.text(name.as_str()); + } else { + self.text(&format!("arg{i}")); + } + if !(strip_receiver && i == 0 && matches!(param, Param::PosOnly(..) | Param::Pos(..))) { + self.text(": "); + self.annotation(param.as_type()); + } + if let Param::PosOnly(_, _, Required::Optional(default)) + | Param::Pos(_, _, Required::Optional(default)) + | Param::KwOnly(_, _, Required::Optional(default)) = param + { + self.text(" = "); + match default.as_ref() { + Some(DefaultValue { + display: Some(text), + .. + }) => self.text(text), + Some(DefaultValue { ty: Type::None, .. }) => self.text("None"), + Some(DefaultValue { + ty: Type::Literal(literal), + .. + }) if !matches!(literal.value, Lit::Enum(_)) => { + self.text(&literal.value.to_string()); + } + _ => { + // Arbitrary default expressions can depend on names in the defining module. + self.reference("typing", "cast"); + self.text("("); + self.reference("typing", "Any"); + self.text(", ...)"); + } + } + } + } + if posonly { + self.text(", /"); + } + } +} diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 60c46eb3cb..d67a85098a 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -654,6 +654,49 @@ fn compute_extract_superclass_actions( }) } +fn compute_implement_abstract_actions( + code: &str, +) -> ( + ModuleInfo, + Vec>, + Vec, +) { + let (handles, state) = + mk_multi_file_state_assert_no_errors(&[("main", code)], Require::Everything); + let handle = handles.get("main").unwrap(); + let transaction = state.transaction(); + let module_info = transaction.get_module_info(handle).unwrap(); + let selection = find_marked_range_with(code, "# ACTION-START", "# ACTION-END"); + let actions = transaction + .implement_abstract_members_code_actions(handle, selection) + .unwrap_or_default(); + let edit_sets: Vec> = + actions.iter().map(|action| action.edits.clone()).collect(); + let titles = actions.iter().map(|action| action.title.clone()).collect(); + (module_info, edit_sets, titles) +} + +fn compute_implement_abstract_actions_allow_errors( + code: &str, +) -> ( + ModuleInfo, + Vec>, + Vec, +) { + let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Everything, false); + let handle = handles.get("main").unwrap(); + let transaction = state.transaction(); + let module_info = transaction.get_module_info(handle).unwrap(); + let selection = find_marked_range_with(code, "# ACTION-START", "# ACTION-END"); + let actions = transaction + .implement_abstract_members_code_actions(handle, selection) + .unwrap_or_default(); + let edit_sets: Vec> = + actions.iter().map(|action| action.edits.clone()).collect(); + let titles = actions.iter().map(|action| action.title.clone()).collect(); + (module_info, edit_sets, titles) +} + #[test] fn basic_test() { let report = get_batched_lsp_operations_report_allow_error( @@ -4363,25 +4406,29 @@ class Sub(Super): #[test] fn push_member_down_basic() { let code = r#" +from abc import ABC + class Base: # MOVE-START def foo(self): pass # MOVE-END -class Child(Base): +class Child(Base, ABC): pass "#; let (module_info, actions, titles) = compute_push_down_actions(code); assert_eq!(vec!["Push `foo` down to `Child`"], titles); let updated = apply_refactor_edits_for_module(&module_info, &actions[0]); let expected = r#" +from abc import ABC + class Base: # MOVE-START pass # MOVE-END -class Child(Base): +class Child(Base, ABC): def foo(self): pass "#; @@ -4975,6 +5022,321 @@ x = foo() assert_eq!(expected_c.trim(), updated_c.trim()); } +#[test] +fn implement_abstract_members_preserves_class_body() { + for (body, expected) in [ + ("pass", ""), + ("pass;", ""), + ("\"\"\"Child doc.\"\"\"", " \"\"\"Child doc.\"\"\"\n"), + ( + "\n \"\"\"Child doc.\"\"\"", + " \"\"\"Child doc.\"\"\"\n", + ), + ("pass; value = 1", " pass; value = 1\n"), + ( + "\n \"\"\"Child doc.\"\"\"; pass", + " \"\"\"Child doc.\"\"\"; pass\n", + ), + ] { + let separator = if body.starts_with('\n') { "" } else { " " }; + let code = format!( + "from abc import ABC, abstractmethod\nclass Base(ABC):\n @abstractmethod\n def foo(self) -> int: ...\n# ACTION-START\nclass Child(Base):{separator}{body}" + ); + let (handles, state) = mk_multi_file_state(&[("main", &code)], Require::Everything, false); + let handle = &handles["main"]; + let transaction = state.transaction(); + let info = transaction.get_module_info(handle).unwrap(); + let start = TextSize::try_from(code.find("class Child").unwrap()).unwrap(); + let actions = transaction + .implement_abstract_members_code_actions(handle, TextRange::empty(start)) + .unwrap(); + let after = apply_refactor_edits_for_module(&info, &actions[0].edits); + assert!(after.contains(&format!( + "class Child(Base):\n{expected} def foo(self) -> int:\n raise NotImplementedError\n" + )), "{after}"); + mk_multi_file_state_assert_no_errors(&[("main", &after)], Require::Everything); + } +} + +#[test] +fn implement_abstract_members_indentation_and_decorators() { + let code = r#" +from abc import ABC, abstractmethod +class Base(ABC): + @classmethod + @abstractmethod + def create(cls, /, value: int = 0, *, flag: bool = False) -> int: ... + @staticmethod + @abstractmethod + def parse(value: str, scale: float = 1.5) -> int: ... +class Outer: + # ACTION-START + class Child(Base): + @staticmethod + def existing() -> None: pass + # ACTION-END +"# + .replace(" ", "\t") + .replace('\n', "\r\n"); + let (info, actions, _) = compute_implement_abstract_actions_allow_errors(&code); + let after = apply_refactor_edits_for_module(&info, &actions[0]); + let normalized = after.replace("\r\n", "\n"); + for expected in [ + "\t\t@classmethod\n\t\tdef create(cls, /, value: int = 0, *, flag: bool = False) -> int:", + "\t\t@staticmethod\n\t\tdef parse(value: str, scale: float = 1.5) -> int:", + "\t\t@staticmethod\n\t\tdef existing() -> None: pass", + ] { + assert!(normalized.contains(expected), "{after}"); + } + mk_multi_file_state_assert_no_errors(&[("main", &after)], Require::Everything); +} + +#[test] +fn implement_abstract_members_preserves_contracts_and_imports() { + let base = r#" +from abc import ABC, abstractmethod +from typing import Callable, Protocol, Self, overload + +class Payload: pass + +class Callback(Protocol): + def __call__(self, value: int, *, flag: bool) -> str: ... + +class Base[T](ABC): + @abstractmethod + def callback(self, cb: Callable[[int], str], *args: int, flag: bool) -> Self: ... + + @abstractmethod + def callback_protocol(self, cb: Callback) -> None: ... + + @abstractmethod + def unknown(self, value = object()): ... + + @abstractmethod + async def fetch(self, value: Payload) -> T: ... + + @abstractmethod + def identity[U: Payload](self, value: U) -> U: ... + + @overload + @abstractmethod + def convert(self, value: int) -> str: ... + @overload + @abstractmethod + def convert(self, value: str) -> int: ... + + @property + @abstractmethod + def value(self) -> int: ... + @value.setter + @abstractmethod + def value(self, value: int) -> None: ... +"#; + let code = "from __future__ import annotations\nfrom base import Base\n# ACTION-START\nclass Child(Base[int]):\n pass\n# ACTION-END\n"; + let (handles, state) = mk_multi_file_state( + &[("base", base), ("main", code)], + Require::Everything, + false, + ); + let transaction = state.transaction(); + let handle = &handles["main"]; + let info = transaction.get_module_info(handle).unwrap(); + let actions = transaction + .implement_abstract_members_code_actions( + handle, + find_marked_range_with(code, "# ACTION-START", "# ACTION-END"), + ) + .unwrap(); + let after = apply_refactor_edits_for_module(&info, &actions[0].edits); + for expected in [ + "cb: Callable[[int], str], *args: int, flag: bool) -> Self", + "def callback_protocol(self, cb: Callback) -> None:", + "async def fetch(self, value: Payload) -> int:", + "def identity[U: Payload](self, value: U) -> U:", + "def convert(self, value: int) -> str:", + "def convert(self, value: str) -> int:", + "@value.setter\n def value(self, value: int) -> None:", + "from base import Payload", + ] { + assert!(after.contains(expected), "Missing {expected:?}:\n{after}"); + } + assert!(!after.contains("Unknown"), "{after}"); + assert!(!after.contains("Self@"), "{after}"); + mk_multi_file_state_assert_no_errors(&[("base", base), ("main", &after)], Require::Everything); +} + +#[test] +fn implement_abstract_members_resolves_import_aliases_and_collisions() { + let base = "from abc import ABC, abstractmethod\nclass Payload: pass\nclass Base(ABC):\n @abstractmethod\n def foo(self, value: Payload) -> Payload: ...\n"; + for (imports, annotation) in [ + ("from base import Payload as P\n", "P"), + ("import base as b\n", "b.Payload"), + ("class Payload: pass\n", "base.Payload"), + ] { + let code = format!( + "from base import Base\n{imports}# ACTION-START\nclass Child(Base):\n pass\n# ACTION-END\n" + ); + let (handles, state) = mk_multi_file_state( + &[("base", base), ("main", &code)], + Require::Everything, + false, + ); + let transaction = state.transaction(); + let handle = &handles["main"]; + let info = transaction.get_module_info(handle).unwrap(); + let actions = transaction + .implement_abstract_members_code_actions( + handle, + find_marked_range_with(&code, "# ACTION-START", "# ACTION-END"), + ) + .unwrap(); + let after = apply_refactor_edits_for_module(&info, &actions[0].edits); + assert!( + after.contains(&format!( + "def foo(self, value: {annotation}) -> {annotation}:" + )), + "{after}" + ); + mk_multi_file_state_assert_no_errors( + &[("base", base), ("main", &after)], + Require::Everything, + ); + } +} + +#[test] +fn implement_abstract_members_preserves_protocol_attributes() { + let code = r#" +from typing import ClassVar, Protocol +class Proto(Protocol): + count: ClassVar[int] + name: str +# ACTION-START +class Child(Proto): + pass +# ACTION-END +"#; + let (info, actions, _) = compute_implement_abstract_actions_allow_errors(code); + let after = apply_refactor_edits_for_module(&info, &actions[0]); + assert!( + after.contains("class Child(Proto):\n count: ClassVar[int]\n"), + "{after}" + ); + mk_multi_file_state_assert_no_errors(&[("main", &after)], Require::Everything); +} + +#[test] +fn implement_abstract_members_inserts_stubs() { + let code = r#" +from abc import ABC, abstractmethod + +class Base(ABC): + @abstractmethod + def foo(self, x: int) -> str: + """Foo doc.""" + ... + + @property + @abstractmethod + def bar(self) -> int: + """Bar doc.""" + ... + +class Child(Base): + # ACTION-START + pass + # ACTION-END +"#; + let (module_info, actions, titles) = compute_implement_abstract_actions(code); + assert_eq!(vec!["Implement abstract members"], titles); + let edits = actions.first().expect("missing abstract member edits"); + let after = apply_refactor_edits_for_module(&module_info, edits); + assert_eq!( + r#" +from abc import ABC, abstractmethod + +class Base(ABC): + @abstractmethod + def foo(self, x: int) -> str: + """Foo doc.""" + ... + + @property + @abstractmethod + def bar(self) -> int: + """Bar doc.""" + ... + +class Child(Base): + # ACTION-START + @property + def bar(self) -> int: + """Bar doc.""" + raise NotImplementedError + + def foo(self, x: int) -> str: + """Foo doc.""" + raise NotImplementedError + # ACTION-END +"#, + after + ); +} + +#[test] +fn implement_protocol_members_inserts_stubs() { + let code = r#" +from typing import Protocol + +class Proto(Protocol): + def foo(self, x: int) -> str: + """Foo doc.""" + ... + + @property + def bar(self) -> int: + """Bar doc.""" + ... + +class Impl(Proto): + # ACTION-START + pass + # ACTION-END +"#; + let (module_info, actions, titles) = compute_implement_abstract_actions_allow_errors(code); + assert_eq!(vec!["Implement abstract members"], titles); + let edits = actions.first().expect("missing protocol member edits"); + let after = apply_refactor_edits_for_module(&module_info, edits); + assert_eq!( + r#" +from typing import Protocol + +class Proto(Protocol): + def foo(self, x: int) -> str: + """Foo doc.""" + ... + + @property + def bar(self) -> int: + """Bar doc.""" + ... + +class Impl(Proto): + # ACTION-START + @property + def bar(self) -> int: + """Bar doc.""" + raise NotImplementedError + + def foo(self, x: int) -> str: + """Foo doc.""" + raise NotImplementedError + # ACTION-END +"#, + after + ); +} + #[test] fn make_local_function_top_level() { let code = r#" diff --git a/pyrefly/lib/test/lsp/lsp_interaction.rs b/pyrefly/lib/test/lsp/lsp_interaction.rs index e2ab90c30b..0ef386e460 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction.rs @@ -27,6 +27,7 @@ mod empty_response_reason; mod file_watcher; mod folding_range; mod hover; +mod implement_abstract_members; mod implementation; mod inlay_hint; mod io; diff --git a/pyrefly/lib/test/lsp/lsp_interaction/implement_abstract_members.rs b/pyrefly/lib/test/lsp/lsp_interaction/implement_abstract_members.rs new file mode 100644 index 0000000000..42e90693eb --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/implement_abstract_members.rs @@ -0,0 +1,71 @@ +/* + * 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::fs; + +use lsp_types::CodeActionKind; +use lsp_types::CodeActionRequest; +use lsp_types::CodeActionResponse; +use lsp_types::Uri; +use pyrefly_lsp_test::object_model::InitializeSettings; +use pyrefly_lsp_test::object_model::LspInteraction; +use serde_json::json; + +use crate::test::lsp::lsp_interaction::util::get_test_files_root; + +#[test] +fn test_implement_abstract_members_action_filter() { + let root = get_test_files_root(); + let path = root.path().join("abstract_members.py"); + fs::write(&path, "from abc import ABC, abstractmethod\nclass Base(ABC):\n @abstractmethod\n def foo(self) -> int: ...\nclass Child(Base): pass\n").unwrap(); + let uri = Uri::from_file_path(&path).unwrap(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction + .initialize(InitializeSettings { + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + interaction.client.did_open("abstract_members.py"); + for (only, trigger_kind, expected) in [ + (json!(["quickfix"]), 1, true), + (json!(["quickfix"]), 2, true), + (json!(["refactor"]), 1, false), + (json!(null), 1, true), + ] { + interaction + .client + .send_request::(json!({ + "textDocument": { "uri": uri }, + "range": { + "start": { "line": 4, "character": 6 }, + "end": { "line": 4, "character": 11 } + }, + "context": { "diagnostics": [], "only": only, "triggerKind": trigger_kind } + })) + .expect_response_with(|response| { + let action = response.as_ref().and_then(|actions| { + actions.iter().find_map(|action| match action { + CodeActionResponse::CodeAction(action) + if action.title == "Implement abstract members" => + { + Some(action) + } + _ => None, + }) + }); + assert_eq!(action.is_some(), expected); + if let Some(action) = action { + assert_eq!(action.kind, Some(CodeActionKind::QuickFix)); + } + true + }) + .unwrap(); + } + interaction.shutdown().unwrap(); +}