crates/ide_assists/src/handlers/generate_enum_is_method.rs

   1 use stdx::to_lower_snake_case;
   2 use syntax::ast::HasVisibility;
   3 use syntax::ast::{self, AstNode, HasName};
   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 //     ///
  32 //     /// [`Minor`]: Version::Minor
  33 //     fn is_minor(&self) -> bool {
  34 //         matches!(self, Self::Minor)
  35 //     }
  36 // }
  37 // ```
  38 pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  39     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  40     let variant_name = variant.name()?;
  41     let parent_enum = ast::Adt::Enum(variant.parent_enum());
  42     let pattern_suffix = match variant.kind() {
  43         ast::StructKind::Record(_) => " { .. }",
  44         ast::StructKind::Tuple(_) => "(..)",
  45         ast::StructKind::Unit => "",
  46     };
  47
  48     let enum_name = parent_enum.name()?;
  49     let enum_lowercase_name = to_lower_snake_case(&enum_name.to_string()).replace('_', " ");
  50     let fn_name = format!("is_{}", &to_lower_snake_case(&variant_name.text()));
  51
  52     // Return early if we've found an existing new fn
  53     let impl_def = find_struct_impl(ctx, &parent_enum, &fn_name)?;
  54
  55     let target = variant.syntax().text_range();
  56     acc.add(
  57         AssistId("generate_enum_is_method", AssistKind::Generate),
  58         "Generate an `is_` method for an enum variant",
  59         target,
  60         |builder| {
  61             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
  62             let method = format!(
  63                 "    /// Returns `true` if the {} is [`{variant}`].
  64     ///
  65     /// [`{variant}`]: {}::{variant}
  66     {}fn {}(&self) -> bool {{
  67         matches!(self, Self::{variant}{})
  68     }}",
  69                 enum_lowercase_name,
  70                 enum_name,
  71                 vis,
  72                 fn_name,
  73                 pattern_suffix,
  74                 variant = variant_name
  75             );
  76
  77             add_method_to_adt(builder, &parent_enum, impl_def, &method);
  78         },
  79     )
  80 }
  81
  82 #[cfg(test)]
  83 mod tests {
  84     use crate::tests::{check_assist, check_assist_not_applicable};
  85
  86     use super::*;
  87
  88     #[test]
  89     fn test_generate_enum_is_from_variant() {
  90         check_assist(
  91             generate_enum_is_method,
  92             r#"
  93 enum Variant {
  94     Undefined,
  95     Minor$0,
  96     Major,
  97 }"#,
  98             r#"enum Variant {
  99     Undefined,
 100     Minor,
 101     Major,
 102 }
 103
 104 impl Variant {
 105     /// Returns `true` if the variant is [`Minor`].
 106     ///
 107     /// [`Minor`]: Variant::Minor
 108     fn is_minor(&self) -> bool {
 109         matches!(self, Self::Minor)
 110     }
 111 }"#,
 112         );
 113     }
 114
 115     #[test]
 116     fn test_generate_enum_is_already_implemented() {
 117         check_assist_not_applicable(
 118             generate_enum_is_method,
 119             r#"
 120 enum Variant {
 121     Undefined,
 122     Minor$0,
 123     Major,
 124 }
 125
 126 impl Variant {
 127     fn is_minor(&self) -> bool {
 128         matches!(self, Self::Minor)
 129     }
 130 }"#,
 131         );
 132     }
 133
 134     #[test]
 135     fn test_generate_enum_is_from_tuple_variant() {
 136         check_assist(
 137             generate_enum_is_method,
 138             r#"
 139 enum Variant {
 140     Undefined,
 141     Minor(u32)$0,
 142     Major,
 143 }"#,
 144             r#"enum Variant {
 145     Undefined,
 146     Minor(u32),
 147     Major,
 148 }
 149
 150 impl Variant {
 151     /// Returns `true` if the variant is [`Minor`].
 152     ///
 153     /// [`Minor`]: Variant::Minor
 154     fn is_minor(&self) -> bool {
 155         matches!(self, Self::Minor(..))
 156     }
 157 }"#,
 158         );
 159     }
 160
 161     #[test]
 162     fn test_generate_enum_is_from_record_variant() {
 163         check_assist(
 164             generate_enum_is_method,
 165             r#"
 166 enum Variant {
 167     Undefined,
 168     Minor { foo: i32 }$0,
 169     Major,
 170 }"#,
 171             r#"enum Variant {
 172     Undefined,
 173     Minor { foo: i32 },
 174     Major,
 175 }
 176
 177 impl Variant {
 178     /// Returns `true` if the variant is [`Minor`].
 179     ///
 180     /// [`Minor`]: Variant::Minor
 181     fn is_minor(&self) -> bool {
 182         matches!(self, Self::Minor { .. })
 183     }
 184 }"#,
 185         );
 186     }
 187
 188     #[test]
 189     fn test_generate_enum_is_from_variant_with_one_variant() {
 190         check_assist(
 191             generate_enum_is_method,
 192             r#"enum Variant { Undefi$0ned }"#,
 193             r#"
 194 enum Variant { Undefined }
 195
 196 impl Variant {
 197     /// Returns `true` if the variant is [`Undefined`].
 198     ///
 199     /// [`Undefined`]: Variant::Undefined
 200     fn is_undefined(&self) -> bool {
 201         matches!(self, Self::Undefined)
 202     }
 203 }"#,
 204         );
 205     }
 206
 207     #[test]
 208     fn test_generate_enum_is_from_variant_with_visibility_marker() {
 209         check_assist(
 210             generate_enum_is_method,
 211             r#"
 212 pub(crate) enum Variant {
 213     Undefined,
 214     Minor$0,
 215     Major,
 216 }"#,
 217             r#"pub(crate) enum Variant {
 218     Undefined,
 219     Minor,
 220     Major,
 221 }
 222
 223 impl Variant {
 224     /// Returns `true` if the variant is [`Minor`].
 225     ///
 226     /// [`Minor`]: Variant::Minor
 227     pub(crate) fn is_minor(&self) -> bool {
 228         matches!(self, Self::Minor)
 229     }
 230 }"#,
 231         );
 232     }
 233
 234     #[test]
 235     fn test_multiple_generate_enum_is_from_variant() {
 236         check_assist(
 237             generate_enum_is_method,
 238             r#"
 239 enum Variant {
 240     Undefined,
 241     Minor,
 242     Major$0,
 243 }
 244
 245 impl Variant {
 246     /// Returns `true` if the variant is [`Minor`].
 247     ///
 248     /// [`Minor`]: Variant::Minor
 249     fn is_minor(&self) -> bool {
 250         matches!(self, Self::Minor)
 251     }
 252 }"#,
 253             r#"enum Variant {
 254     Undefined,
 255     Minor,
 256     Major,
 257 }
 258
 259 impl Variant {
 260     /// Returns `true` if the variant is [`Minor`].
 261     ///
 262     /// [`Minor`]: Variant::Minor
 263     fn is_minor(&self) -> bool {
 264         matches!(self, Self::Minor)
 265     }
 266
 267     /// Returns `true` if the variant is [`Major`].
 268     ///
 269     /// [`Major`]: Variant::Major
 270     fn is_major(&self) -> bool {
 271         matches!(self, Self::Major)
 272     }
 273 }"#,
 274         );
 275     }
 276
 277     #[test]
 278     fn test_generate_enum_is_variant_names() {
 279         check_assist(
 280             generate_enum_is_method,
 281             r#"
 282 enum GeneratorState {
 283     Yielded,
 284     Complete$0,
 285     Major,
 286 }"#,
 287             r#"enum GeneratorState {
 288     Yielded,
 289     Complete,
 290     Major,
 291 }
 292
 293 impl GeneratorState {
 294     /// Returns `true` if the generator state is [`Complete`].
 295     ///
 296     /// [`Complete`]: GeneratorState::Complete
 297     fn is_complete(&self) -> bool {
 298         matches!(self, Self::Complete)
 299     }
 300 }"#,
 301         );
 302     }
 303 }