crates/ide_assists/src/handlers/generate_enum_projection_method.rs

   1 use itertools::Itertools;
   2 use stdx::to_lower_snake_case;
   3 use syntax::ast::VisibilityOwner;
   4 use syntax::ast::{self, AstNode, NameOwner};
   5
   6 use crate::{
   7     utils::{add_method_to_adt, find_struct_impl},
   8     AssistContext, AssistId, AssistKind, Assists,
   9 };
  10
  11 // Assist: generate_enum_try_into_method
  12 //
  13 // Generate an `try_into_` method for an enum variant.
  14 //
  15 // ```
  16 // enum Value {
  17 //  Number(i32),
  18 //  Text(String)$0,
  19 // }
  20 // ```
  21 // ->
  22 // ```
  23 // enum Value {
  24 //  Number(i32),
  25 //  Text(String),
  26 // }
  27 //
  28 // impl Value {
  29 //     fn try_into_text(self) -> Result<String, Self> {
  30 //         if let Self::Text(v) = self {
  31 //             Ok(v)
  32 //         } else {
  33 //             Err(self)
  34 //         }
  35 //     }
  36 // }
  37 // ```
  38 pub(crate) fn generate_enum_try_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  39     generate_enum_projection_method(
  40         acc,
  41         ctx,
  42         "generate_enum_try_into_method",
  43         "Generate an `try_into_` method for an enum variant",
  44         ProjectionProps {
  45             fn_name_prefix: "try_into",
  46             self_param: "self",
  47             return_prefix: "Result<",
  48             return_suffix: ", Self>",
  49             happy_case: "Ok",
  50             sad_case: "Err(self)",
  51         },
  52     )
  53 }
  54
  55 // Assist: generate_enum_as_method
  56 //
  57 // Generate an `as_` method for an enum variant.
  58 //
  59 // ```
  60 // enum Value {
  61 //  Number(i32),
  62 //  Text(String)$0,
  63 // }
  64 // ```
  65 // ->
  66 // ```
  67 // enum Value {
  68 //  Number(i32),
  69 //  Text(String),
  70 // }
  71 //
  72 // impl Value {
  73 //     fn as_text(&self) -> Option<&String> {
  74 //         if let Self::Text(v) = self {
  75 //             Some(v)
  76 //         } else {
  77 //             None
  78 //         }
  79 //     }
  80 // }
  81 // ```
  82 pub(crate) fn generate_enum_as_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  83     generate_enum_projection_method(
  84         acc,
  85         ctx,
  86         "generate_enum_as_method",
  87         "Generate an `as_` method for an enum variant",
  88         ProjectionProps {
  89             fn_name_prefix: "as",
  90             self_param: "&self",
  91             return_prefix: "Option<&",
  92             return_suffix: ">",
  93             happy_case: "Some",
  94             sad_case: "None",
  95         },
  96     )
  97 }
  98
  99 struct ProjectionProps {
 100     fn_name_prefix: &'static str,
 101     self_param: &'static str,
 102     return_prefix: &'static str,
 103     return_suffix: &'static str,
 104     happy_case: &'static str,
 105     sad_case: &'static str,
 106 }
 107
 108 fn generate_enum_projection_method(
 109     acc: &mut Assists,
 110     ctx: &AssistContext,
 111     assist_id: &'static str,
 112     assist_description: &str,
 113     props: ProjectionProps,
 114 ) -> Option<()> {
 115     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
 116     let variant_name = variant.name()?;
 117     let parent_enum = ast::Adt::Enum(variant.parent_enum());
 118
 119     let (pattern_suffix, field_type, bound_name) = match variant.kind() {
 120         ast::StructKind::Record(record) => {
 121             let (field,) = record.fields().collect_tuple()?;
 122             let name = field.name()?.to_string();
 123             let ty = field.ty()?;
 124             let pattern_suffix = format!(" {{ {} }}", name);
 125             (pattern_suffix, ty, name)
 126         }
 127         ast::StructKind::Tuple(tuple) => {
 128             let (field,) = tuple.fields().collect_tuple()?;
 129             let ty = field.ty()?;
 130             ("(v)".to_owned(), ty, "v".to_owned())
 131         }
 132         ast::StructKind::Unit => return None,
 133     };
 134
 135     let fn_name = format!("{}_{}", props.fn_name_prefix, &to_lower_snake_case(variant_name.text()));
 136
 137     // Return early if we've found an existing new fn
 138     let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?;
 139
 140     let target = variant.syntax().text_range();
 141     acc.add(AssistId(assist_id, AssistKind::Generate), assist_description, target, |builder| {
 142         let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
 143         let method = format!(
 144             "    {0}fn {1}({2}) -> {3}{4}{5} {{
 145         if let Self::{6}{7} = self {{
 146             {8}({9})
 147         }} else {{
 148             {10}
 149         }}
 150     }}",
 151             vis,
 152             fn_name,
 153             props.self_param,
 154             props.return_prefix,
 155             field_type.syntax(),
 156             props.return_suffix,
 157             variant_name,
 158             pattern_suffix,
 159             props.happy_case,
 160             bound_name,
 161             props.sad_case,
 162         );
 163
 164         add_method_to_adt(builder, &parent_enum, impl_def, &method);
 165     })
 166 }
 167
 168 #[cfg(test)]
 169 mod tests {
 170     use crate::tests::{check_assist, check_assist_not_applicable};
 171
 172     use super::*;
 173
 174     #[test]
 175     fn test_generate_enum_try_into_tuple_variant() {
 176         check_assist(
 177             generate_enum_try_into_method,
 178             r#"
 179 enum Value {
 180     Number(i32),
 181     Text(String)$0,
 182 }"#,
 183             r#"enum Value {
 184     Number(i32),
 185     Text(String),
 186 }
 187
 188 impl Value {
 189     fn try_into_text(self) -> Result<String, Self> {
 190         if let Self::Text(v) = self {
 191             Ok(v)
 192         } else {
 193             Err(self)
 194         }
 195     }
 196 }"#,
 197         );
 198     }
 199
 200     #[test]
 201     fn test_generate_enum_try_into_already_implemented() {
 202         check_assist_not_applicable(
 203             generate_enum_try_into_method,
 204             r#"enum Value {
 205     Number(i32),
 206     Text(String)$0,
 207 }
 208
 209 impl Value {
 210     fn try_into_text(self) -> Result<String, Self> {
 211         if let Self::Text(v) = self {
 212             Ok(v)
 213         } else {
 214             Err(self)
 215         }
 216     }
 217 }"#,
 218         );
 219     }
 220
 221     #[test]
 222     fn test_generate_enum_try_into_unit_variant() {
 223         check_assist_not_applicable(
 224             generate_enum_try_into_method,
 225             r#"enum Value {
 226     Number(i32),
 227     Text(String),
 228     Unit$0,
 229 }"#,
 230         );
 231     }
 232
 233     #[test]
 234     fn test_generate_enum_try_into_record_with_multiple_fields() {
 235         check_assist_not_applicable(
 236             generate_enum_try_into_method,
 237             r#"enum Value {
 238     Number(i32),
 239     Text(String),
 240     Both { first: i32, second: String }$0,
 241 }"#,
 242         );
 243     }
 244
 245     #[test]
 246     fn test_generate_enum_try_into_tuple_with_multiple_fields() {
 247         check_assist_not_applicable(
 248             generate_enum_try_into_method,
 249             r#"enum Value {
 250     Number(i32),
 251     Text(String, String)$0,
 252 }"#,
 253         );
 254     }
 255
 256     #[test]
 257     fn test_generate_enum_try_into_record_variant() {
 258         check_assist(
 259             generate_enum_try_into_method,
 260             r#"enum Value {
 261     Number(i32),
 262     Text { text: String }$0,
 263 }"#,
 264             r#"enum Value {
 265     Number(i32),
 266     Text { text: String },
 267 }
 268
 269 impl Value {
 270     fn try_into_text(self) -> Result<String, Self> {
 271         if let Self::Text { text } = self {
 272             Ok(text)
 273         } else {
 274             Err(self)
 275         }
 276     }
 277 }"#,
 278         );
 279     }
 280
 281     #[test]
 282     fn test_generate_enum_as_tuple_variant() {
 283         check_assist(
 284             generate_enum_as_method,
 285             r#"
 286 enum Value {
 287     Number(i32),
 288     Text(String)$0,
 289 }"#,
 290             r#"enum Value {
 291     Number(i32),
 292     Text(String),
 293 }
 294
 295 impl Value {
 296     fn as_text(&self) -> Option<&String> {
 297         if let Self::Text(v) = self {
 298             Some(v)
 299         } else {
 300             None
 301         }
 302     }
 303 }"#,
 304         );
 305     }
 306
 307     #[test]
 308     fn test_generate_enum_as_record_variant() {
 309         check_assist(
 310             generate_enum_as_method,
 311             r#"enum Value {
 312     Number(i32),
 313     Text { text: String }$0,
 314 }"#,
 315             r#"enum Value {
 316     Number(i32),
 317     Text { text: String },
 318 }
 319
 320 impl Value {
 321     fn as_text(&self) -> Option<&String> {
 322         if let Self::Text { text } = self {
 323             Some(text)
 324         } else {
 325             None
 326         }
 327     }
 328 }"#,
 329         );
 330     }
 331 }