Rewrite DeriveEntity

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"

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<syn::Ident>,
+        pub entity: Option<syn::Ident>,
+        pub model: Option<syn::Ident>,
+        pub primary_key: Option<syn::Ident>,
+        pub relation: Option<syn::Ident>,
+        pub schema_name: Option<syn::Lit>,
+        pub table_name: Option<syn::Lit>,
+    }
+}
+
+pub mod field_attr {
+    use bae::FromAttributes;
+
+    #[derive(Default, FromAttributes)]
+    pub struct Sea {
+        pub auto_increment: Option<syn::Lit>,
+        pub belongs_to: Option<syn::Lit>,
+        pub column_type: Option<syn::Lit>,
+        pub column_type_raw: Option<syn::Lit>,
+        pub from: Option<syn::Lit>,
+        pub indexed: Option<()>,
+        pub null: Option<()>,
+        pub primary_key: Option<()>,
+        pub to: Option<syn::Lit>,
+        pub unique: Option<()>,
+    }
+}

sea-orm-macros/src/derives/active_model.rs

@@ -46,13 +46,13 @@ pub fn expand_derive_active_model(ident: Ident, data: Data) -> syn::Result<Token
             }
         }
 
-        impl Default for ActiveModel {
+        impl std::default::Default for ActiveModel {
             fn default() -> Self {
                 <Self as sea_orm::ActiveModelBehavior>::new()
             }
         }
 
-        impl From<<Entity as EntityTrait>::Model> for ActiveModel {
+        impl std::convert::From<<Entity as EntityTrait>::Model> for ActiveModel {
             fn from(m: <Entity as EntityTrait>::Model) -> Self {
                 Self {
                     #(#field: sea_orm::unchanged_active_value_not_intended_for_public_use(m.#field)),*

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<syn::Lit> {
-    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<syn::Lit>,
+    table_name: Option<syn::Lit>,
+}
+
+impl DeriveEntity {
+    fn new(input: syn::DeriveInput) -> Result<Self, syn::Error> {
+        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 {
+                    <Self as sea_orm::EntityName>::table_name(self)
+                }
+            }
+        )
     }
-    None
 }
 
-pub fn expand_derive_entity(ident: Ident, attrs: Vec<Attribute>) -> syn::Result<TokenStream> {
-    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 {
-                <Self as sea_orm::EntityName>::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<TokenStream> {
+    Ok(DeriveEntity::new(input)?.expand())
 }

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))]