Add proc macro to implement models
Signed-off-by: Lee Smet <lee.smet@hotmail.com>
This commit is contained in:
14
heromodels/heromodels-derive/Cargo.toml
Normal file
14
heromodels/heromodels-derive/Cargo.toml
Normal file
@@ -0,0 +1,14 @@
|
||||
[package]
|
||||
name = "heromodels-derive"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
description = "Derive macros for heromodels"
|
||||
authors = ["Your Name <your.email@example.com>"]
|
||||
|
||||
[lib]
|
||||
proc-macro = true
|
||||
|
||||
[dependencies]
|
||||
syn = { version = "2.0", features = ["full", "extra-traits", "parsing"] }
|
||||
quote = "1.0"
|
||||
proc-macro2 = "1.0"
|
158
heromodels/heromodels-derive/src/lib.rs
Normal file
158
heromodels/heromodels-derive/src/lib.rs
Normal file
@@ -0,0 +1,158 @@
|
||||
use proc_macro::TokenStream;
|
||||
use quote::{quote, format_ident};
|
||||
use syn::{parse_macro_input, DeriveInput, Data, Fields};
|
||||
|
||||
/// Convert a string to snake_case
|
||||
fn to_snake_case(s: &str) -> String {
|
||||
let mut result = String::new();
|
||||
for (i, c) in s.char_indices() {
|
||||
if i > 0 && c.is_uppercase() {
|
||||
result.push('_');
|
||||
}
|
||||
result.push(c.to_lowercase().next().unwrap());
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Convert a string to PascalCase
|
||||
fn to_pascal_case(s: &str) -> String {
|
||||
let mut result = String::new();
|
||||
let mut capitalize_next = true;
|
||||
|
||||
for c in s.chars() {
|
||||
if c == '_' {
|
||||
capitalize_next = true;
|
||||
} else if capitalize_next {
|
||||
result.push(c.to_uppercase().next().unwrap());
|
||||
capitalize_next = false;
|
||||
} else {
|
||||
result.push(c);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
|
||||
/// Implements the Model trait and generates Index trait implementations for fields marked with #[index].
|
||||
#[proc_macro_attribute]
|
||||
pub fn model(_attr: TokenStream, item: TokenStream) -> TokenStream {
|
||||
// Parse the input tokens into a syntax tree
|
||||
let input = parse_macro_input!(item as DeriveInput);
|
||||
|
||||
// Extract struct name
|
||||
let struct_name = &input.ident;
|
||||
|
||||
// Convert struct name to snake_case for db_prefix
|
||||
let name_str = struct_name.to_string();
|
||||
let db_prefix = to_snake_case(&name_str);
|
||||
|
||||
// Extract fields with /// @index doc comment
|
||||
let mut indexed_fields = Vec::new();
|
||||
|
||||
if let Data::Struct(data_struct) = &input.data {
|
||||
if let Fields::Named(fields_named) = &data_struct.fields {
|
||||
for field in &fields_named.named {
|
||||
for attr in &field.attrs {
|
||||
if attr.path().is_ident("doc") {
|
||||
let meta = attr.meta.clone().try_into().unwrap();
|
||||
if let syn::Meta::NameValue(name_value) = meta {
|
||||
if let syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Str(lit_str), .. }) = name_value.value {
|
||||
let doc_str = lit_str.value();
|
||||
if doc_str.trim().starts_with("@index") {
|
||||
if let Some(field_name) = &field.ident {
|
||||
indexed_fields.push(field_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Generate Model trait implementation
|
||||
let db_keys_impl = if indexed_fields.is_empty() {
|
||||
quote! {
|
||||
fn db_keys(&self) -> Vec<IndexKey> {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let field_keys = indexed_fields.iter().map(|field_name| {
|
||||
let name_str = field_name.to_string();
|
||||
quote! {
|
||||
IndexKey {
|
||||
name: #name_str,
|
||||
value: self.#field_name.to_string(),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
quote! {
|
||||
fn db_keys(&self) -> Vec<IndexKey> {
|
||||
vec![
|
||||
#(#field_keys),*
|
||||
]
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let model_impl = quote! {
|
||||
impl Model for #struct_name {
|
||||
fn db_prefix() -> &'static str {
|
||||
#db_prefix
|
||||
}
|
||||
|
||||
fn get_id(&self) -> u32 {
|
||||
self.base_data.id
|
||||
}
|
||||
|
||||
fn base_data_mut(&mut self) -> &mut BaseModelData {
|
||||
&mut self.base_data
|
||||
}
|
||||
|
||||
#db_keys_impl
|
||||
}
|
||||
};
|
||||
|
||||
// Generate Index trait implementations
|
||||
let mut index_impls = proc_macro2::TokenStream::new();
|
||||
|
||||
for field_name in &indexed_fields {
|
||||
let name_str = field_name.to_string();
|
||||
|
||||
// Convert field name to PascalCase for struct name
|
||||
let struct_name_str = to_pascal_case(&name_str);
|
||||
let index_struct_name = format_ident!("{}", struct_name_str);
|
||||
|
||||
// Default to str for key type
|
||||
let index_impl = quote! {
|
||||
pub struct #index_struct_name;
|
||||
|
||||
impl Index for #index_struct_name {
|
||||
type Model = #struct_name;
|
||||
type Key = str;
|
||||
|
||||
fn key() -> &'static str {
|
||||
#name_str
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
index_impls.extend(index_impl);
|
||||
}
|
||||
|
||||
// Combine the original struct with the generated implementations
|
||||
let expanded = quote! {
|
||||
#input
|
||||
|
||||
#model_impl
|
||||
|
||||
#index_impls
|
||||
};
|
||||
|
||||
// Return the generated code
|
||||
expanded.into()
|
||||
}
|
31
heromodels/heromodels-derive/tests/test_model_macro.rs
Normal file
31
heromodels/heromodels-derive/tests/test_model_macro.rs
Normal file
@@ -0,0 +1,31 @@
|
||||
use heromodels::model;
|
||||
use heromodels::models::core::model::{BaseModelData, Model, Index};
|
||||
|
||||
#[model]
|
||||
struct TestUser {
|
||||
base_data: BaseModelData,
|
||||
|
||||
#[index]
|
||||
username: String,
|
||||
|
||||
#[index]
|
||||
is_active: bool,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_basic_model() {
|
||||
assert_eq!(TestUser::db_prefix(), "test_user");
|
||||
|
||||
let user = TestUser {
|
||||
base_data: BaseModelData::new(1),
|
||||
username: "test".to_string(),
|
||||
is_active: true,
|
||||
};
|
||||
|
||||
let keys = user.db_keys();
|
||||
assert_eq!(keys.len(), 2);
|
||||
assert_eq!(keys[0].name, "username");
|
||||
assert_eq!(keys[0].value, "test");
|
||||
assert_eq!(keys[1].name, "is_active");
|
||||
assert_eq!(keys[1].value, "true");
|
||||
}
|
Reference in New Issue
Block a user