crates/ide_assists/src/handlers/generate_enum_is_method.rs

   1 use stdx::to_lower_snake_case;
   2 use syntax::ast::VisibilityOwner;
   3 use syntax::ast::{self, AstNode, NameOwner};
   4
   5 use crate::{
   6     utils::{add_method_to_adt, find_struct_impl},
   7     AssistContext, AssistId, AssistKind, Assists,
   8 };
   9
  10 // Assist: generate_enum_is_method
  11 //
  12 // Generate an `is_` method for an enum variant.
  13 //
  14 // ```
  15 // enum Version {
  16 //  Undefined,
  17 //  Minor$0,
  18 //  Major,
  19 // }
  20 // ```
  21 // ->
  22 // ```
  23 // enum Version {
  24 //  Undefined,
  25 //  Minor,
  26 //  Major,
  27 // }
  28 //
  29 // impl Version {
  30 //     /// Returns `true` if the version is [`Minor`].
  31 //     fn is_minor(&self) -> bool {
  32 //         matches!(self, Self::Minor)
  33 //     }
  34 // }
  35 // ```
  36 pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  37     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  38     let variant_name = variant.name()?;
  39     let parent_enum = ast::Adt::Enum(variant.parent_enum());
  40     let pattern_suffix = match variant.kind() {
  41         ast::StructKind::Record(_) => " { .. }",
  42         ast::StructKind::Tuple(_) => "(..)",
  43         ast::StructKind::Unit => "",
  44     };
  45
  46     let enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string());
  47     let fn_name = format!("is_{}", &to_lower_snake_case(&variant_name.text()));
  48
  49     // Return early if we've found an existing new fn
  50     let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?;
  51
  52     let target = variant.syntax().text_range();
  53     acc.add(
  54         AssistId("generate_enum_is_method", AssistKind::Generate),
  55         "Generate an `is_` method for an enum variant",
  56         target,
  57         |builder| {
  58             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
  59             let method = format!(
  60                 "    /// Returns `true` if the {} is [`{}`].
  61     {}fn {}(&self) -> bool {{
  62         matches!(self, Self::{}{})
  63     }}",
  64                 enum_lowercase_name, variant_name, vis, fn_name, variant_name, pattern_suffix,
  65             );
  66
  67             add_method_to_adt(builder, &parent_enum, impl_def, &method);
  68         },
  69     )
  70 }
  71
  72 #[cfg(test)]
  73 mod tests {
  74     use crate::tests::{check_assist, check_assist_not_applicable};
  75
  76     use super::*;
  77
  78     #[test]
  79     fn test_generate_enum_is_from_variant() {
  80         check_assist(
  81             generate_enum_is_method,
  82             r#"
  83 enum Variant {
  84     Undefined,
  85     Minor$0,
  86     Major,
  87 }"#,
  88             r#"enum Variant {
  89     Undefined,
  90     Minor,
  91     Major,
  92 }
  93
  94 impl Variant {
  95     /// Returns `true` if the variant is [`Minor`].
  96     fn is_minor(&self) -> bool {
  97         matches!(self, Self::Minor)
  98     }
  99 }"#,
 100         );
 101     }
 102
 103     #[test]
 104     fn test_generate_enum_is_already_implemented() {
 105         check_assist_not_applicable(
 106             generate_enum_is_method,
 107             r#"
 108 enum Variant {
 109     Undefined,
 110     Minor$0,
 111     Major,
 112 }
 113
 114 impl Variant {
 115     fn is_minor(&self) -> bool {
 116         matches!(self, Self::Minor)
 117     }
 118 }"#,
 119         );
 120     }
 121
 122     #[test]
 123     fn test_generate_enum_is_from_tuple_variant() {
 124         check_assist(
 125             generate_enum_is_method,
 126             r#"
 127 enum Variant {
 128     Undefined,
 129     Minor(u32)$0,
 130     Major,
 131 }"#,
 132             r#"enum Variant {
 133     Undefined,
 134     Minor(u32),
 135     Major,
 136 }
 137
 138 impl Variant {
 139     /// Returns `true` if the variant is [`Minor`].
 140     fn is_minor(&self) -> bool {
 141         matches!(self, Self::Minor(..))
 142     }
 143 }"#,
 144         );
 145     }
 146
 147     #[test]
 148     fn test_generate_enum_is_from_record_variant() {
 149         check_assist(
 150             generate_enum_is_method,
 151             r#"
 152 enum Variant {
 153     Undefined,
 154     Minor { foo: i32 }$0,
 155     Major,
 156 }"#,
 157             r#"enum Variant {
 158     Undefined,
 159     Minor { foo: i32 },
 160     Major,
 161 }
 162
 163 impl Variant {
 164     /// Returns `true` if the variant is [`Minor`].
 165     fn is_minor(&self) -> bool {
 166         matches!(self, Self::Minor { .. })
 167     }
 168 }"#,
 169         );
 170     }
 171
 172     #[test]
 173     fn test_generate_enum_is_from_variant_with_one_variant() {
 174         check_assist(
 175             generate_enum_is_method,
 176             r#"enum Variant { Undefi$0ned }"#,
 177             r#"
 178 enum Variant { Undefined }
 179
 180 impl Variant {
 181     /// Returns `true` if the variant is [`Undefined`].
 182     fn is_undefined(&self) -> bool {
 183         matches!(self, Self::Undefined)
 184     }
 185 }"#,
 186         );
 187     }
 188
 189     #[test]
 190     fn test_generate_enum_is_from_variant_with_visibility_marker() {
 191         check_assist(
 192             generate_enum_is_method,
 193             r#"
 194 pub(crate) enum Variant {
 195     Undefined,
 196     Minor$0,
 197     Major,
 198 }"#,
 199             r#"pub(crate) enum Variant {
 200     Undefined,
 201     Minor,
 202     Major,
 203 }
 204
 205 impl Variant {
 206     /// Returns `true` if the variant is [`Minor`].
 207     pub(crate) fn is_minor(&self) -> bool {
 208         matches!(self, Self::Minor)
 209     }
 210 }"#,
 211         );
 212     }
 213
 214     #[test]
 215     fn test_multiple_generate_enum_is_from_variant() {
 216         check_assist(
 217             generate_enum_is_method,
 218             r#"
 219 enum Variant {
 220     Undefined,
 221     Minor,
 222     Major$0,
 223 }
 224
 225 impl Variant {
 226     /// Returns `true` if the variant is [`Minor`].
 227     fn is_minor(&self) -> bool {
 228         matches!(self, Self::Minor)
 229     }
 230 }"#,
 231             r#"enum Variant {
 232     Undefined,
 233     Minor,
 234     Major,
 235 }
 236
 237 impl Variant {
 238     /// Returns `true` if the variant is [`Minor`].
 239     fn is_minor(&self) -> bool {
 240         matches!(self, Self::Minor)
 241     }
 242
 243     /// Returns `true` if the variant is [`Major`].
 244     fn is_major(&self) -> bool {
 245         matches!(self, Self::Major)
 246     }
 247 }"#,
 248         );
 249     }
 250 }