sea-orm-try-into-active-model-macro 0.1.0

use bae::FromAttributes;
use proc_macro2::{Span, TokenStream};
use quote::{quote, quote_spanned};
use syn::{punctuated::Punctuated, token::Comma, Lit, Meta, Type};

/// Attributes to derive an ActiveModel
#[derive(Default, FromAttributes)]
pub struct SeaOrm {
    pub active_model: Option<syn::Path>,
    pub error: Option<syn::Path>,
}

#[allow(clippy::enum_variant_names)]
enum Error {
    InputNotStruct,
    ErrorTypeNotProvided,
    Syn(syn::Error),
}

struct TryIntoActiveModel {
    attrs: SeaOrm,
    fields: syn::punctuated::Punctuated<syn::Field, syn::token::Comma>,
    field_idents: Vec<syn::Ident>,
    ident: syn::Ident,
}

impl TryIntoActiveModel {
    fn new(input: syn::DeriveInput) -> Result<Self, Error> {
        let fields = match input.data {
            syn::Data::Struct(syn::DataStruct {
                fields: syn::Fields::Named(syn::FieldsNamed { named, .. }), ..
            }) => named,
            _ => return Err(Error::InputNotStruct),
        };

        let attrs = SeaOrm::try_from_attributes(&input.attrs).map_err(Error::Syn)?.unwrap_or_default();
        if attrs.error.is_none() {
            return Err(Error::ErrorTypeNotProvided);
        }

        let ident = input.ident;

        let field_idents = fields.iter().map(|field| field.ident.as_ref().unwrap().clone()).collect();

        Ok(TryIntoActiveModel { attrs, fields, field_idents, ident })
    }

    fn expand(&self) -> syn::Result<TokenStream> {
        let expanded_impl_into_active_model = self.impl_into_active_model();

        Ok(expanded_impl_into_active_model)
    }

    fn impl_into_active_model(&self) -> TokenStream {
        let Self { attrs, ident, field_idents, fields } = self;

        let active_model_ident =
            attrs.active_model.clone().unwrap_or_else(|| syn::Ident::new("ActiveModel", Span::call_site()).into());

        let error_ident = attrs.error.clone().unwrap();

        let expanded_fields_into_active_model = fields.iter().map(|field| {
            let mut try_into = None;
            let mut direct = false;
            for attr in &field.attrs {
                if attr.path.get_ident().map(|i| i == "sea_orm") == Some(true) {
                    if let Ok(list) = attr.parse_args_with(Punctuated::<Meta, Comma>::parse_terminated) {
                        for meta in list.iter() {
                            if let Meta::NameValue(value) = meta {
                                if value.path.get_ident().is_some() && value.path.get_ident().unwrap() == "try_into" {
                                    if let Lit::Str(lit) = &value.lit {
                                        try_into = Some(lit.clone());
                                    }
                                }
                            }
                            if let Meta::Path(value) = meta {
                                if value.get_ident().unwrap() == "direct" {
                                    direct = true;
                                }
                            }
                        }
                    }
                }
            }
            //panic!("field attrs: {:?}", field.attrs);

            let field_ident = field.ident.as_ref().unwrap();
            let mut optional = false;
            if let Type::Path(path) = &field.ty {
                path.path.segments.iter().for_each(|segment| {
                    if segment.ident == "Option" {
                        optional = true;
                    }
                });
            }

            if let Some(try_into) = try_into {
                let try_into: proc_macro2::TokenStream = try_into.parse().unwrap();
                if direct {
                    if optional {
                        quote!(
                            if self.#field_ident.is_none() {
                                sea_orm::entity::ActiveValue::Set(None)
                            } else {
                                #try_into(self.#field_ident)?
                            }
                        )
                    } else {
                        quote!(
                            #try_into(self.#field_ident)?
                        )
                    }
                } else if optional {
                    quote!(
                        if self.#field_ident.is_none() {
                            sea_orm::entity::ActiveValue::Set(None)
                        } else {
                            sea_orm::IntoActiveValue::<_>::into_active_value(#try_into(self.#field_ident)?).into()
                        }
                    )
                } else {
                    quote!(
                        sea_orm::IntoActiveValue::<_>::into_active_value(#try_into(self.#field_ident)?).into()
                    )
                }
            } else if direct {
                if optional {
                    quote!(
                        if self.#field_ident.is_none() {
                            sea_orm::entity::ActiveValue::Set(None)
                        } else {
                        self.#field_ident
                        }
                    )
                } else {
                    quote!(
                        self.#field_ident
                    )
                }
            } else if optional {
                quote!(
                    if self.#field_ident.is_none() {
                        sea_orm::entity::ActiveValue::Set(None)
                    } else {
                        sea_orm::IntoActiveValue::<_>::into_active_value(self.#field_ident).into()
                    }
                )
            } else {
                quote!(
                    sea_orm::IntoActiveValue::<_>::into_active_value(self.#field_ident).into()
                )
            }
        });

        quote!(
            #[automatically_derived]
            impl sea_orm_try_into_active_model::TryIntoActiveModel<#active_model_ident> for #ident {
                type Error = #error_ident;
                fn try_into_active_model(self) -> Result<#active_model_ident, Self::Error> {
                    Ok(#active_model_ident {
                        #( #field_idents: #expanded_fields_into_active_model, )*
                        ..::std::default::Default::default()
                    })
                }
            }
        )
    }
}

/// Method to derive the ActiveModel from the [ActiveModelTrait](sea_orm::ActiveModelTrait)
pub fn expand_try_into_active_model(input: syn::DeriveInput) -> syn::Result<TokenStream> {
    let ident_span = input.ident.span();

    match TryIntoActiveModel::new(input) {
        Ok(model) => model.expand(),
        Err(Error::InputNotStruct) => Ok(quote_spanned! {
            ident_span => compile_error!("you can only derive TryIntoActiveModel on structs");
        }),
        Err(Error::ErrorTypeNotProvided) => Ok(quote_spanned! {
            ident_span => compile_error!("TryIntoActiveModel needs to have provide error type as error parameter");
        }),
        Err(Error::Syn(err)) => Err(err),
    }
}