sea_query_derive/iden/
mod.rs

1mod attr;
2mod error;
3pub mod iden_static;
4mod path;
5mod write_arm;
6
7use darling::FromMeta;
8use heck::ToSnakeCase;
9use proc_macro::TokenStream;
10use quote::{quote, quote_spanned};
11use std::convert::{TryFrom, TryInto};
12use syn::{Attribute, DataEnum, DataStruct, DeriveInput, Fields, Ident, Variant};
13
14use attr::IdenAttr;
15use error::ErrorMsg;
16use path::IdenPath;
17use write_arm::IdenVariant;
18
19struct DeriveIden;
20struct DeriveIdenStatic;
21
22pub fn expand(input: DeriveInput) -> TokenStream {
23    let DeriveInput {
24        ident, data, attrs, ..
25    } = input;
26
27    let table_name = match get_table_name(&ident, attrs) {
28        Ok(v) => v,
29        Err(e) => return e.to_compile_error().into(),
30    };
31
32    // Currently we only support enums and unit structs
33    let variants =
34        match data {
35            syn::Data::Enum(DataEnum { variants, .. }) => variants,
36            syn::Data::Struct(DataStruct {
37                fields: Fields::Unit,
38                ..
39            }) => return impl_iden_for_unit_struct(&ident, &table_name).into(),
40            _ => return quote_spanned! {
41                ident.span() => compile_error!("you can only derive Iden on enums or unit structs");
42            }
43            .into(),
44        };
45
46    if variants.is_empty() {
47        return TokenStream::new();
48    }
49
50    let output = impl_iden_for_enum(&ident, &table_name, variants.iter());
51
52    output.into()
53}
54
55fn impl_iden_for_unit_struct(
56    ident: &proc_macro2::Ident,
57    table_name: &str,
58) -> proc_macro2::TokenStream {
59    let sea_query_path = sea_query_path();
60
61    if is_static_iden(table_name) {
62        quote! {
63            impl #sea_query_path::Iden for #ident {
64                fn quoted(&self) -> std::borrow::Cow<'static, str> {
65                    std::borrow::Cow::Borrowed(#table_name)
66                }
67
68                fn unquoted(&self) -> &str {
69                    #table_name
70                }
71            }
72        }
73    } else {
74        quote! {
75            impl #sea_query_path::Iden for #ident {
76                fn unquoted(&self) -> &str {
77                    #table_name
78                }
79            }
80        }
81    }
82}
83
84fn impl_iden_for_enum<'a, T>(
85    ident: &proc_macro2::Ident,
86    table_name: &str,
87    variants: T,
88) -> proc_macro2::TokenStream
89where
90    T: Iterator<Item = &'a Variant>,
91{
92    let sea_query_path = sea_query_path();
93
94    let mut is_all_static_iden = true;
95
96    let match_arms = match variants
97        .map(|v| {
98            let v = IdenVariant::<DeriveIden>::try_from((table_name, v))?;
99            is_all_static_iden &= v.is_static_iden();
100            Ok(v)
101        })
102        .collect::<syn::Result<Vec<_>>>()
103    {
104        Ok(v) => v,
105        Err(e) => return e.to_compile_error(),
106    };
107
108    if is_all_static_iden {
109        quote! {
110            impl #sea_query_path::Iden for #ident {
111                fn quoted(&self) -> std::borrow::Cow<'static, str> {
112                    std::borrow::Cow::Borrowed(match self {
113                        #(#match_arms),*
114                    })
115                }
116
117                fn unquoted(&self) -> &str {
118                    match self {
119                        #(#match_arms),*
120                    }
121                }
122            }
123        }
124    } else {
125        quote! {
126            impl #sea_query_path::Iden for #ident {
127                fn unquoted(&self) -> &str {
128                    match self {
129                        #(#match_arms),*
130                    }
131                }
132            }
133        }
134    }
135}
136
137fn sea_query_path() -> proc_macro2::TokenStream {
138    if cfg!(feature = "sea-orm") {
139        quote!(sea_orm::sea_query)
140    } else {
141        quote!(sea_query)
142    }
143}
144
145pub struct NamingHolder {
146    pub default: Ident,
147    pub pascal: Ident,
148}
149
150#[derive(Debug, FromMeta)]
151pub struct GenEnumArgs {
152    #[darling(default)]
153    pub prefix: Option<String>,
154    #[darling(default)]
155    pub suffix: Option<String>,
156    #[darling(default)]
157    pub crate_name: Option<String>,
158    #[darling(default)]
159    pub table_name: Option<String>,
160}
161
162pub const DEFAULT_PREFIX: &str = "";
163pub const DEFAULT_SUFFIX: &str = "Iden";
164pub const DEFAULT_CRATE_NAME: &str = "sea_query";
165
166impl Default for GenEnumArgs {
167    fn default() -> Self {
168        Self {
169            prefix: Some(DEFAULT_PREFIX.to_string()),
170            suffix: Some(DEFAULT_SUFFIX.to_string()),
171            crate_name: Some(DEFAULT_CRATE_NAME.to_string()),
172            table_name: None,
173        }
174    }
175}
176
177fn get_table_name(ident: &proc_macro2::Ident, attrs: Vec<Attribute>) -> Result<String, syn::Error> {
178    let table_name = match find_attr(&attrs) {
179        Some(att) => match att.try_into()? {
180            IdenAttr::Rename(lit) => lit,
181            _ => return Err(syn::Error::new_spanned(att, ErrorMsg::ContainerAttr)),
182        },
183        None => ident.to_string().to_snake_case(),
184    };
185    Ok(table_name)
186}
187
188fn find_attr(attrs: &[Attribute]) -> Option<&Attribute> {
189    attrs.iter().find(|attr| {
190        attr.path().is_ident(&IdenPath::Iden) || attr.path().is_ident(&IdenPath::Method)
191    })
192}
193
194pub fn is_static_iden(name: &str) -> bool {
195    // can only begin with [a-z_]
196    name.chars()
197        .take(1)
198        .all(|c| c == '_' || c.is_ascii_alphabetic())
199        && name.chars().all(|c| c == '_' || c.is_ascii_alphanumeric())
200}
sea_query_derive/iden/mod.rs

sea_query_derive/iden/
mod.rs
