extern crate proc_macro; use std::iter::Peekable; use proc_macro::token_stream::IntoIter; use proc_macro::{TokenStream, TokenTree}; fn expect_iden bool>(s: &mut Peekable, mut name: F) { let token = s.next().unwrap(); match token { TokenTree::Ident(iden) => { if !name(iden.to_string()) { panic!("expect identifier"); } } _ => panic!("expect identifier got {token:?}"), }; } fn skip_proc_macro(s: &mut Peekable) { while let Some(token) = s.peek() { match token { TokenTree::Punct(p) if p.to_string().as_str() == "#" => { s.next(); s.next(); } _ => return, } } } #[derive(Default, Debug)] struct VariantInfo { rs_name: String, db_name: String, } #[derive(Default, Debug)] struct EnumDef { name: String, variants: Vec, } #[proc_macro_derive(ActixAdminEnum)] pub fn derive_actix_admin_enum(item: TokenStream) -> TokenStream { let clone = item.clone(); let mut it = clone.into_iter().peekable(); skip_proc_macro(&mut it); expect_iden(&mut it, |iden| &iden == "pub"); expect_iden(&mut it, |iden| &iden == "enum"); let mut def = EnumDef::default(); expect_iden(&mut it, |iden| { def.name = iden; true }); let mut it = if let TokenTree::Group(g) = it.next().unwrap() { g.stream().into_iter().peekable() } else { panic!("expect enum body"); }; // Parse macros parse_enum_variants(&mut it, &mut def); let mut buffer = String::new(); buffer.push_str(&format!("impl std::str::FromStr for {} {{\n type Err = ();\n fn from_str(s: &str) -> Result {{\n match s {{\n", def.name)); for v in def.variants { buffer.push_str(&format!( " {:?} => Ok({}::{}),\n", v.db_name, def.name, v.rs_name )); } buffer.push_str(" _ => Err(()),\n"); buffer.push_str(" }\n }\n}"); // eprintln!("{buffer}"); buffer.as_str().parse().unwrap() } fn parse_enum_variants(body_it: &mut Peekable, def: &mut EnumDef) { while body_it.peek().is_some() { if let Some(variant) = parse_enum_variant(body_it) { def.variants.push(variant); } if let Some(TokenTree::Punct(p)) = body_it.peek() { if p.to_string().as_str() == "," { body_it.next(); } } } } fn parse_enum_variant(body_it: &mut Peekable) -> Option { let mut variant_info = VariantInfo::default(); while body_it.peek().is_some() { if let Some(db_name) = parse_enum_variant_macro(body_it) { variant_info.db_name = db_name; break; } if let Some(TokenTree::Punct(p)) = body_it.peek() { if p.to_string().as_str() == "," { body_it.next(); } } } let iden = body_it.next().expect("No variant name"); variant_info.rs_name = iden.to_string(); Some(variant_info) } fn parse_enum_variant_macro(body_it: &mut Peekable) -> Option { let (_punc, group) = ( body_it.next().expect("No # for macro"), body_it.next().expect("No macro body"), ); // eprintln!("{punc:#?} {group:#?}"); let mut it = if let TokenTree::Group(g) = group { // parse #[sea_orm] g.stream().into_iter().peekable() } else { panic!("Unexpected token after #: {group:?}"); }; let token = it.next(); if let Some(TokenTree::Ident(iden)) = token { if iden.to_string().as_str() != "sea_orm" { eprintln!("iden is not sea_orm"); return None; } } else { eprintln!("token should be iden but is {token:#?}"); return None; } let it = if let Some(TokenTree::Group(g)) = it.next() { g.stream().into_iter().peekable() } else { return None; }; let mut it = it.skip_while(|t| match t { TokenTree::Ident(id) => id.to_string().as_str() != "string_value", _ => true, }); let (_id, _punct, literal) = (it.next(), it.next(), it.next()); match literal { Some(TokenTree::Literal(l)) => Some(l.to_string().replace('\"', "")), _ => None, } }