Derive attributes

This commit is contained in:
Billy Chan 2021-09-13 18:33:07 +08:00
parent f801590b9b
commit 39c5a4d134
No known key found for this signature in database
GPG Key ID: A2D690CAC7DF3CC7
2 changed files with 39 additions and 5 deletions

View File

@ -2,6 +2,7 @@ use heck::{CamelCase, SnakeCase};
use proc_macro2::{Ident, TokenStream}; use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote}; use quote::{format_ident, quote};
use sea_query::{ColumnDef, ColumnSpec, ColumnType}; use sea_query::{ColumnDef, ColumnSpec, ColumnType};
use syn::punctuated::Punctuated;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct Column { pub struct Column {
@ -52,6 +53,20 @@ impl Column {
} }
} }
pub fn get_col_type_attrs(&self) -> Option<TokenStream> {
let col_type = match &self.col_type {
ColumnType::Float(Some(l)) => Some(format!("Float(Some({}))", l)),
ColumnType::Double(Some(l)) => Some(format!("Double(Some({}))", l)),
ColumnType::Decimal(Some((p, s))) => Some(format!("Decimal(Some(({}, {})))", p, s)),
ColumnType::Money(Some((p, s))) => Some(format!("Money(Some({}, {}))", p, s)),
ColumnType::Custom(iden) => {
Some(format!("Custom(\"{}\".to_owned())", iden.to_string()))
}
_ => None,
};
col_type.map(|ty| quote! { column_type = #ty })
}
pub fn get_def(&self) -> TokenStream { pub fn get_def(&self) -> TokenStream {
let mut col_def = match &self.col_type { let mut col_def = match &self.col_type {
ColumnType::Char(s) => match s { ColumnType::Char(s) => match s {

View File

@ -1,6 +1,8 @@
use crate::Entity; use crate::Entity;
use proc_macro2::TokenStream; use proc_macro2::TokenStream;
use quote::quote; use quote::quote;
use std::iter::FromIterator;
use syn::{punctuated::Punctuated, token::Comma};
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct EntityWriter { pub struct EntityWriter {
@ -332,13 +334,30 @@ impl EntityWriter {
.columns .columns
.iter() .iter()
.map(|col| { .map(|col| {
if !primary_keys.contains(&col.name) { let mut attrs: Punctuated<_, Comma> = Punctuated::new();
TokenStream::new() if primary_keys.contains(&col.name) {
} else { attrs.push(quote! { primary_key });
quote! { if !col.auto_increment {
#[sea_orm(primary_key)] attrs.push(quote! { auto_increment = false });
} }
} }
if let Some(ts) = col.get_col_type_attrs() {
attrs.extend(vec![ts]);
};
if !attrs.is_empty() {
let mut ts = TokenStream::new();
for (i, attr) in attrs.into_iter().enumerate() {
if i > 0 {
ts = quote! { #ts, };
}
ts = quote! { #ts #attr };
}
quote! {
#[sea_orm(#ts)]
}
} else {
TokenStream::new()
}
}) })
.collect(); .collect();
quote! { quote! {