diff --git a/sea-orm-macros/Cargo.toml b/sea-orm-macros/Cargo.toml index 9de45ae3..27ea856e 100644 --- a/sea-orm-macros/Cargo.toml +++ b/sea-orm-macros/Cargo.toml @@ -16,6 +16,7 @@ path = "src/lib.rs" proc-macro = true [dependencies] +bae = "^0.1" syn = { version = "^1", default-features = false, features = [ "full", "derive", "clone-impls", "parsing", "proc-macro", "printing", "extra-traits" ] } quote = "^1" heck = "^0.3" diff --git a/sea-orm-macros/src/attributes.rs b/sea-orm-macros/src/attributes.rs new file mode 100644 index 00000000..d195a98f --- /dev/null +++ b/sea-orm-macros/src/attributes.rs @@ -0,0 +1,32 @@ +pub mod derive_attr { + use bae::FromAttributes; + + #[derive(Default, FromAttributes)] + pub struct Sea { + pub column: Option, + pub entity: Option, + pub model: Option, + pub primary_key: Option, + pub relation: Option, + pub schema_name: Option, + pub table_name: Option, + } +} + +pub mod field_attr { + use bae::FromAttributes; + + #[derive(Default, FromAttributes)] + pub struct Sea { + pub auto_increment: Option, + pub belongs_to: Option, + pub column_type: Option, + pub column_type_raw: Option, + pub from: Option, + pub indexed: Option<()>, + pub null: Option<()>, + pub primary_key: Option<()>, + pub to: Option, + pub unique: Option<()>, + } +} diff --git a/sea-orm-macros/src/derives/active_model.rs b/sea-orm-macros/src/derives/active_model.rs index 4d18a169..98bc0aec 100644 --- a/sea-orm-macros/src/derives/active_model.rs +++ b/sea-orm-macros/src/derives/active_model.rs @@ -46,13 +46,13 @@ pub fn expand_derive_active_model(ident: Ident, data: Data) -> syn::Result Self { ::new() } } - impl From<::Model> for ActiveModel { + impl std::convert::From<::Model> for ActiveModel { fn from(m: ::Model) -> Self { Self { #(#field: sea_orm::unchanged_active_value_not_intended_for_public_use(m.#field)),* diff --git a/sea-orm-macros/src/derives/entity.rs b/sea-orm-macros/src/derives/entity.rs index de6b97fb..06696d8f 100644 --- a/sea-orm-macros/src/derives/entity.rs +++ b/sea-orm-macros/src/derives/entity.rs @@ -1,51 +1,135 @@ -use heck::SnakeCase; -use proc_macro2::{Ident, TokenStream}; -use quote::quote; -use syn::{Attribute, Meta}; +use std::iter::FromIterator; -fn get_entity_attr(attrs: &[Attribute]) -> Option { - for attr in attrs { - let name_value = match attr.parse_meta() { - Ok(Meta::NameValue(nv)) => nv, - _ => continue, - }; - if name_value.path.is_ident("table") { - return Some(name_value.lit); - } +use proc_macro2::TokenStream; +use quote::{format_ident, quote}; + +use crate::attributes::derive_attr; + +struct DeriveEntity { + column_ident: syn::Ident, + ident: syn::Ident, + model_ident: syn::Ident, + primary_key_ident: syn::Ident, + relation_ident: syn::Ident, + schema_name: Option, + table_name: Option, +} + +impl DeriveEntity { + fn new(input: syn::DeriveInput) -> Result { + let sea_attr = derive_attr::Sea::try_from_attributes(&input.attrs)?.unwrap_or_default(); + + let ident = input.ident; + let column_ident = sea_attr.column.unwrap_or_else(|| format_ident!("Column")); + let model_ident = sea_attr.model.unwrap_or_else(|| format_ident!("Model")); + let primary_key_ident = sea_attr + .primary_key + .unwrap_or_else(|| format_ident!("PrimaryKey")); + let relation_ident = sea_attr + .relation + .unwrap_or_else(|| format_ident!("Relation")); + + let table_name = sea_attr.table_name; + let schema_name = sea_attr.schema_name; + + Ok(DeriveEntity { + column_ident, + ident, + model_ident, + primary_key_ident, + relation_ident, + schema_name, + table_name, + }) + } + + fn expand(&self) -> TokenStream { + let expanded_impl_entity_name = self.impl_entity_name(); + let expanded_impl_entity_trait = self.impl_entity_trait(); + let expanded_impl_iden = self.impl_iden(); + let expanded_impl_iden_static = self.impl_iden_static(); + + TokenStream::from_iter([ + expanded_impl_entity_name, + expanded_impl_entity_trait, + expanded_impl_iden, + expanded_impl_iden_static, + ]) + } + + fn impl_entity_name(&self) -> TokenStream { + let ident = &self.ident; + let table_name = match &self.table_name { + Some(table_name) => table_name, + None => return TokenStream::new(), // No table name, do not derive EntityName + }; + let expanded_schema_name = self + .schema_name + .as_ref() + .map(|schema| quote!(Some(#schema))) + .unwrap_or_else(|| quote!(None)); + + quote!( + impl sea_orm::entity::EntityName for #ident { + fn schema_name(&self) -> Option<&str> { + #expanded_schema_name + } + + fn table_name(&self) -> &str { + #table_name + } + } + ) + } + + fn impl_entity_trait(&self) -> TokenStream { + let Self { + ident, + model_ident, + column_ident, + primary_key_ident, + relation_ident, + .. + } = self; + + quote!( + impl sea_orm::entity::EntityTrait for #ident { + type Model = #model_ident; + + type Column = #column_ident; + + type PrimaryKey = #primary_key_ident; + + type Relation = #relation_ident; + } + ) + } + + fn impl_iden(&self) -> TokenStream { + let ident = &self.ident; + + quote!( + impl sea_orm::Iden for #ident { + fn unquoted(&self, s: &mut dyn std::fmt::Write) { + write!(s, "{}", self.as_str()).unwrap(); + } + } + ) + } + + fn impl_iden_static(&self) -> TokenStream { + let ident = &self.ident; + + quote!( + impl sea_orm::IdenStatic for #ident { + fn as_str(&self) -> &str { + ::table_name(self) + } + } + ) } - None } -pub fn expand_derive_entity(ident: Ident, attrs: Vec) -> syn::Result { - let _entity_name = match get_entity_attr(&attrs) { - Some(lit) => quote! { #lit }, - None => { - let normalized = ident.to_string().to_snake_case(); - quote! { #normalized } - } - }; - - Ok(quote!( - impl sea_orm::Iden for #ident { - fn unquoted(&self, s: &mut dyn std::fmt::Write) { - write!(s, "{}", self.as_str()).unwrap(); - } - } - - impl sea_orm::IdenStatic for #ident { - fn as_str(&self) -> &str { - ::table_name(self) - } - } - - impl EntityTrait for #ident { - type Model = Model; - - type Column = Column; - - type PrimaryKey = PrimaryKey; - - type Relation = Relation; - } - )) +pub fn expand_derive_entity(input: syn::DeriveInput) -> syn::Result { + Ok(DeriveEntity::new(input)?.expand()) } diff --git a/sea-orm-macros/src/lib.rs b/sea-orm-macros/src/lib.rs index b85a84ad..fe132e5f 100644 --- a/sea-orm-macros/src/lib.rs +++ b/sea-orm-macros/src/lib.rs @@ -1,18 +1,17 @@ extern crate proc_macro; use proc_macro::TokenStream; -use syn::{parse_macro_input, DeriveInput}; +use syn::{parse_macro_input, DeriveInput, Error}; +mod attributes; mod derives; -#[proc_macro_derive(DeriveEntity, attributes(table))] +#[proc_macro_derive(DeriveEntity, attributes(sea_orm))] pub fn derive_entity(input: TokenStream) -> TokenStream { - let DeriveInput { ident, attrs, .. } = parse_macro_input!(input); - - match derives::expand_derive_entity(ident, attrs) { - Ok(ts) => ts.into(), - Err(e) => e.to_compile_error().into(), - } + let input = parse_macro_input!(input as DeriveInput); + derives::expand_derive_entity(input) + .unwrap_or_else(Error::into_compile_error) + .into() } #[proc_macro_derive(DeriveEntityModel, attributes(sea_orm))]