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::HasVisibility;
   4 use syntax::ast::{self, AstNode, HasName};
   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 =
 136         format!("{}_{}", props.fn_name_prefix, &to_lower_snake_case(&variant_name.text()));
 137
 138     // Return early if we've found an existing new fn
 139     let impl_def = find_struct_impl(ctx, &parent_enum, &fn_name)?;
 140
 141     let target = variant.syntax().text_range();
 142     acc.add(AssistId(assist_id, AssistKind::Generate), assist_description, target, |builder| {
 143         let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
 144         let method = format!(
 145             "    {0}fn {1}({2}) -> {3}{4}{5} {{
 146         if let Self::{6}{7} = self {{
 147             {8}({9})
 148         }} else {{
 149             {10}
 150         }}
 151     }}",
 152             vis,
 153             fn_name,
 154             props.self_param,
 155             props.return_prefix,
 156             field_type.syntax(),
 157             props.return_suffix,
 158             variant_name,
 159             pattern_suffix,
 160             props.happy_case,
 161             bound_name,
 162             props.sad_case,
 163         );
 164
 165         add_method_to_adt(builder, &parent_enum, impl_def, &method);
 166     })
 167 }
 168
 169 #[cfg(test)]
 170 mod tests {
 171     use crate::tests::{check_assist, check_assist_not_applicable};
 172
 173     use super::*;
 174
 175     #[test]
 176     fn test_generate_enum_try_into_tuple_variant() {
 177         check_assist(
 178             generate_enum_try_into_method,
 179             r#"
 180 enum Value {
 181     Number(i32),
 182     Text(String)$0,
 183 }"#,
 184             r#"enum Value {
 185     Number(i32),
 186     Text(String),
 187 }
 188
 189 impl Value {
 190     fn try_into_text(self) -> Result<String, Self> {
 191         if let Self::Text(v) = self {
 192             Ok(v)
 193         } else {
 194             Err(self)
 195         }
 196     }
 197 }"#,
 198         );
 199     }
 200
 201     #[test]
 202     fn test_generate_enum_try_into_already_implemented() {
 203         check_assist_not_applicable(
 204             generate_enum_try_into_method,
 205             r#"enum Value {
 206     Number(i32),
 207     Text(String)$0,
 208 }
 209
 210 impl Value {
 211     fn try_into_text(self) -> Result<String, Self> {
 212         if let Self::Text(v) = self {
 213             Ok(v)
 214         } else {
 215             Err(self)
 216         }
 217     }
 218 }"#,
 219         );
 220     }
 221
 222     #[test]
 223     fn test_generate_enum_try_into_unit_variant() {
 224         check_assist_not_applicable(
 225             generate_enum_try_into_method,
 226             r#"enum Value {
 227     Number(i32),
 228     Text(String),
 229     Unit$0,
 230 }"#,
 231         );
 232     }
 233
 234     #[test]
 235     fn test_generate_enum_try_into_record_with_multiple_fields() {
 236         check_assist_not_applicable(
 237             generate_enum_try_into_method,
 238             r#"enum Value {
 239     Number(i32),
 240     Text(String),
 241     Both { first: i32, second: String }$0,
 242 }"#,
 243         );
 244     }
 245
 246     #[test]
 247     fn test_generate_enum_try_into_tuple_with_multiple_fields() {
 248         check_assist_not_applicable(
 249             generate_enum_try_into_method,
 250             r#"enum Value {
 251     Number(i32),
 252     Text(String, String)$0,
 253 }"#,
 254         );
 255     }
 256
 257     #[test]
 258     fn test_generate_enum_try_into_record_variant() {
 259         check_assist(
 260             generate_enum_try_into_method,
 261             r#"enum Value {
 262     Number(i32),
 263     Text { text: String }$0,
 264 }"#,
 265             r#"enum Value {
 266     Number(i32),
 267     Text { text: String },
 268 }
 269
 270 impl Value {
 271     fn try_into_text(self) -> Result<String, Self> {
 272         if let Self::Text { text } = self {
 273             Ok(text)
 274         } else {
 275             Err(self)
 276         }
 277     }
 278 }"#,
 279         );
 280     }
 281
 282     #[test]
 283     fn test_generate_enum_as_tuple_variant() {
 284         check_assist(
 285             generate_enum_as_method,
 286             r#"
 287 enum Value {
 288     Number(i32),
 289     Text(String)$0,
 290 }"#,
 291             r#"enum Value {
 292     Number(i32),
 293     Text(String),
 294 }
 295
 296 impl Value {
 297     fn as_text(&self) -> Option<&String> {
 298         if let Self::Text(v) = self {
 299             Some(v)
 300         } else {
 301             None
 302         }
 303     }
 304 }"#,
 305         );
 306     }
 307
 308     #[test]
 309     fn test_generate_enum_as_record_variant() {
 310         check_assist(
 311             generate_enum_as_method,
 312             r#"enum Value {
 313     Number(i32),
 314     Text { text: String }$0,
 315 }"#,
 316             r#"enum Value {
 317     Number(i32),
 318     Text { text: String },
 319 }
 320
 321 impl Value {
 322     fn as_text(&self) -> Option<&String> {
 323         if let Self::Text { text } = self {
 324             Some(text)
 325         } else {
 326             None
 327         }
 328     }
 329 }"#,
 330         );
 331     }
 332 }