oswilno/crates/oswilno-actix-admin/src/lib.rs

extern crate proc_macro;
use std::iter::Peekable;

use proc_macro::token_stream::IntoIter;
use proc_macro::{TokenStream, TokenTree};

fn expect_iden<F: FnMut(String) -> bool>(s: &mut Peekable<IntoIter>, 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<IntoIter>) {
    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<VariantInfo>,
}

#[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<Self, ()> {{\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<IntoIter>, 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<IntoIter>) -> Option<VariantInfo> {
    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<IntoIter>) -> Option<String> {
    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,
    }
}