codama_syn_helpers/extensions/
derive_input.rs

1use super::ToTokensExtension;
2use codama_errors::CodamaResult;
3use syn::DeriveInput;
4
5pub trait DeriveInputExtension {
6    fn get_self(&self) -> &DeriveInput;
7
8    /// Ensure the derive input is a struct and return the data.
9    fn as_struct(&self) -> CodamaResult<&syn::DataStruct> {
10        let this = self.get_self();
11        match &this.data {
12            syn::Data::Struct(data) => Ok(data),
13            _ => Err(this.error("expected a struct").into()),
14        }
15    }
16
17    /// Ensure the derive input is an enum and return the data.
18    fn as_enum(&self) -> CodamaResult<&syn::DataEnum> {
19        let this = self.get_self();
20        match &this.data {
21            syn::Data::Enum(data) => Ok(data),
22            _ => Err(this.error("expected an enum").into()),
23        }
24    }
25}
26
27impl DeriveInputExtension for DeriveInput {
28    fn get_self(&self) -> &DeriveInput {
29        self
30    }
31}
32
33#[cfg(test)]
34mod tests {
35    use super::*;
36
37    #[test]
38    fn as_struct_ok() {
39        let derive_input: DeriveInput = syn::parse_quote! { struct Foo(u32); };
40        assert!(derive_input.as_struct().is_ok());
41    }
42
43    #[test]
44    fn as_struct_err() {
45        let derive_input: DeriveInput = syn::parse_quote! { enum Foo { Bar } };
46        assert!(derive_input.as_struct().is_err());
47    }
48
49    #[test]
50    fn as_enum_ok() {
51        let derive_input: DeriveInput = syn::parse_quote! { enum Foo { Bar } };
52        assert!(derive_input.as_enum().is_ok());
53    }
54
55    #[test]
56    fn as_enum_err() {
57        let derive_input: DeriveInput = syn::parse_quote! { struct Foo(u32); };
58        assert!(derive_input.as_enum().is_err());
59    }
60}