123 lines
3.4 KiB
Rust
123 lines
3.4 KiB
Rust
extern crate proc_macro2;
|
|
extern crate quote;
|
|
extern crate syn;
|
|
|
|
extern crate proc_macro;
|
|
|
|
use parse::{Parse, ParseStream};
|
|
use proc_macro::{Span, TokenStream};
|
|
use punctuated::Punctuated;
|
|
// 用于处理宏输入
|
|
use quote::quote;
|
|
// 用于生成代码
|
|
use syn::*;
|
|
use syn::{parse_macro_input, ItemFn};
|
|
|
|
#[derive(Debug)]
|
|
struct TaskArgs {
|
|
cron: String,
|
|
interval: u64,
|
|
time_zone: String,
|
|
}
|
|
|
|
// 实现 Parse trait 以支持解析参数
|
|
impl Parse for TaskArgs {
|
|
fn parse(input: ParseStream) -> Result<Self> {
|
|
let args = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
|
|
|
|
let mut task_args = TaskArgs {
|
|
cron: String::from(""),
|
|
interval: 0,
|
|
time_zone: String::from(""),
|
|
};
|
|
|
|
for arg in args {
|
|
if let Meta::NameValue(meta) = arg {
|
|
match meta.path.get_ident().unwrap().to_string().as_str() {
|
|
"cron" => {
|
|
if let Expr::Lit(ExprLit { lit, .. }) = meta.value {
|
|
if let Lit::Str(value) = lit {
|
|
task_args.cron = value.value();
|
|
}
|
|
}
|
|
}
|
|
"interval" => {
|
|
if let Expr::Lit(ExprLit { lit, .. }) = meta.value {
|
|
if let Lit::Int(value) = lit {
|
|
task_args.interval = value.base10_parse()?;
|
|
}
|
|
}
|
|
}
|
|
"time_zone" => {
|
|
if let Expr::Lit(ExprLit { lit, .. }) = meta.value {
|
|
if let Lit::Str(value) = lit {
|
|
task_args.time_zone = value.value();
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
if task_args.cron.is_empty() && task_args.interval <= 0 {
|
|
return Err(Error::new(
|
|
Span::call_site().into(),
|
|
"必须设置有效的cron表达式或者interval参数",
|
|
));
|
|
}
|
|
|
|
Ok(task_args)
|
|
}
|
|
}
|
|
|
|
pub fn gen_task(attr: TokenStream, item: TokenStream) -> TokenStream {
|
|
let func = parse_macro_input!(item as ItemFn);
|
|
// 解析宏的参数
|
|
let TaskArgs {
|
|
cron,
|
|
interval,
|
|
time_zone,
|
|
} = parse_macro_input!(attr as TaskArgs);
|
|
|
|
// 将 method 转换为小写并生成标识符
|
|
// let vis = func.vis.clone();
|
|
let ident = func.sig.ident.clone();
|
|
|
|
let result = quote! {
|
|
let mut _result = library::task::Task {
|
|
job: std::sync::Arc::new(|| Box::pin(#ident::#ident ())),
|
|
cron: None,
|
|
interval: None,
|
|
time_zone: None,
|
|
};
|
|
if !#cron.is_empty() {
|
|
_result.cron = Some(#cron.to_string());
|
|
}
|
|
if #interval > 0 {
|
|
_result.interval = Some(#interval);
|
|
}
|
|
if !#time_zone.is_empty() {
|
|
_result.time_zone = Some(#time_zone.to_string());
|
|
}
|
|
_result
|
|
};
|
|
|
|
let generated = quote! {
|
|
#[allow(non_camel_case_types)]
|
|
struct #ident;
|
|
impl #ident {
|
|
#func
|
|
}
|
|
impl library::task::TaskMethod for #ident {
|
|
fn ge_task(&self) -> library::task::Task {
|
|
#result
|
|
}
|
|
}
|
|
|
|
::library::submit_task!(#ident);
|
|
};
|
|
|
|
TokenStream::from(generated)
|
|
}
|