crates/ide_assists/src/handlers/generate_setter.rs

   1 use stdx::{format_to, to_lower_snake_case};
   2 use syntax::ast::{self, AstNode, HasName, HasVisibility};
   3
   4 use crate::{
   5     utils::{find_impl_block_end, find_struct_impl, generate_impl_text},
   6     AssistContext, AssistId, AssistKind, Assists, GroupLabel,
   7 };
   8
   9 // Assist: generate_setter
  10 //
  11 // Generate a setter method.
  12 //
  13 // ```
  14 // struct Person {
  15 //     nam$0e: String,
  16 // }
  17 // ```
  18 // ->
  19 // ```
  20 // struct Person {
  21 //     name: String,
  22 // }
  23 //
  24 // impl Person {
  25 //     /// Set the person's name.
  26 //     fn set_name(&mut self, name: String) {
  27 //         self.name = name;
  28 //     }
  29 // }
  30 // ```
  31 pub(crate) fn generate_setter(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  32     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  33     let field = ctx.find_node_at_offset::<ast::RecordField>()?;
  34
  35     let strukt_name = strukt.name()?;
  36     let field_name = field.name()?;
  37     let field_ty = field.ty()?;
  38
  39     // Return early if we've found an existing fn
  40     let fn_name = to_lower_snake_case(&field_name.to_string());
  41     let impl_def = find_struct_impl(
  42         ctx,
  43         &ast::Adt::Struct(strukt.clone()),
  44         format!("set_{}", fn_name).as_str(),
  45     )?;
  46
  47     let target = field.syntax().text_range();
  48     acc.add_group(
  49         &GroupLabel("Generate getter/setter".to_owned()),
  50         AssistId("generate_setter", AssistKind::Generate),
  51         "Generate a setter method",
  52         target,
  53         |builder| {
  54             let mut buf = String::with_capacity(512);
  55
  56             let fn_name_spaced = fn_name.replace('_', " ");
  57             let strukt_name_spaced =
  58                 to_lower_snake_case(&strukt_name.to_string()).replace('_', " ");
  59
  60             if impl_def.is_some() {
  61                 buf.push('\n');
  62             }
  63
  64             let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
  65             format_to!(
  66                 buf,
  67                 "    /// Set the {}'s {}.
  68     {}fn set_{}(&mut self, {}: {}) {{
  69         self.{} = {};
  70     }}",
  71                 strukt_name_spaced,
  72                 fn_name_spaced,
  73                 vis,
  74                 fn_name,
  75                 fn_name,
  76                 field_ty,
  77                 fn_name,
  78                 fn_name,
  79             );
  80
  81             let start_offset = impl_def
  82                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
  83                 .unwrap_or_else(|| {
  84                     buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
  85                     strukt.syntax().text_range().end()
  86                 });
  87
  88             builder.insert(start_offset, buf);
  89         },
  90     )
  91 }
  92
  93 #[cfg(test)]
  94 mod tests {
  95     use crate::tests::{check_assist, check_assist_not_applicable};
  96
  97     use super::*;
  98
  99     fn check_not_applicable(ra_fixture: &str) {
 100         check_assist_not_applicable(generate_setter, ra_fixture)
 101     }
 102
 103     #[test]
 104     fn test_generate_setter_from_field() {
 105         check_assist(
 106             generate_setter,
 107             r#"
 108 struct Person<T: Clone> {
 109     dat$0a: T,
 110 }"#,
 111             r#"
 112 struct Person<T: Clone> {
 113     data: T,
 114 }
 115
 116 impl<T: Clone> Person<T> {
 117     /// Set the person's data.
 118     fn set_data(&mut self, data: T) {
 119         self.data = data;
 120     }
 121 }"#,
 122         );
 123     }
 124
 125     #[test]
 126     fn test_generate_setter_already_implemented() {
 127         check_not_applicable(
 128             r#"
 129 struct Person<T: Clone> {
 130     dat$0a: T,
 131 }
 132
 133 impl<T: Clone> Person<T> {
 134     fn set_data(&mut self, data: T) {
 135         self.data = data;
 136     }
 137 }"#,
 138         );
 139     }
 140
 141     #[test]
 142     fn test_generate_setter_from_field_with_visibility_marker() {
 143         check_assist(
 144             generate_setter,
 145             r#"
 146 pub(crate) struct Person<T: Clone> {
 147     dat$0a: T,
 148 }"#,
 149             r#"
 150 pub(crate) struct Person<T: Clone> {
 151     data: T,
 152 }
 153
 154 impl<T: Clone> Person<T> {
 155     /// Set the person's data.
 156     pub(crate) fn set_data(&mut self, data: T) {
 157         self.data = data;
 158     }
 159 }"#,
 160         );
 161     }
 162
 163     #[test]
 164     fn test_multiple_generate_setter() {
 165         check_assist(
 166             generate_setter,
 167             r#"
 168 struct Context<T: Clone> {
 169     data: T,
 170     cou$0nt: usize,
 171 }
 172
 173 impl<T: Clone> Context<T> {
 174     /// Set the context's data.
 175     fn set_data(&mut self, data: T) {
 176         self.data = data;
 177     }
 178 }"#,
 179             r#"
 180 struct Context<T: Clone> {
 181     data: T,
 182     count: usize,
 183 }
 184
 185 impl<T: Clone> Context<T> {
 186     /// Set the context's data.
 187     fn set_data(&mut self, data: T) {
 188         self.data = data;
 189     }
 190
 191     /// Set the context's count.
 192     fn set_count(&mut self, count: usize) {
 193         self.count = count;
 194     }
 195 }"#,
 196         );
 197     }
 198 }