Skip to main content

onebot_api_macros/
lib.rs

1mod selector;
2
3use proc_macro::TokenStream;
4
5#[proc_macro_derive(Selector, attributes(selector))]
6pub fn derive_selector(input: TokenStream) -> TokenStream {
7	selector::derive_selector(input)
8}
9
10use quote::quote;
11use syn::{
12	Attribute, Expr, FnArg, ImplItem, ImplItemFn, ItemImpl, Meta, Pat, Type, parse_macro_input,
13};
14
15#[derive(Default)]
16struct ApiAction {
17	extract: Option<String>,
18	response_type: Option<Type>,
19	renames: Vec<(String, String)>,
20}
21
22fn take_api_attr(attrs: &mut Vec<Attribute>) -> Option<Attribute> {
23	let idx = attrs.iter().position(|a| a.path().is_ident("api"))?;
24	Some(attrs.remove(idx))
25}
26
27fn parse_string_meta(nv: &syn::MetaNameValue) -> Option<String> {
28	let Expr::Lit(lit) = &nv.value else {
29		return None;
30	};
31	let syn::Lit::Str(s) = &lit.lit else {
32		return None;
33	};
34	Some(s.value())
35}
36
37fn parse_type_meta(nv: &syn::MetaNameValue) -> Option<Type> {
38	let Expr::Path(expr_path) = &nv.value else {
39		return None;
40	};
41	Some(Type::Path(syn::TypePath {
42		qself: None,
43		path: expr_path.path.clone(),
44	}))
45}
46
47fn parse_map_list(list: &syn::MetaList) -> Vec<(String, String)> {
48	let Ok(inner) = list
49		.parse_args_with(syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated)
50	else {
51		return Vec::new();
52	};
53
54	let mut renames = Vec::new();
55	for m in inner {
56		let Meta::NameValue(nv) = m else {
57			continue;
58		};
59		let Some(param_name) = nv.path.get_ident().map(|i| i.to_string()) else {
60			continue;
61		};
62		if let Some(value) = parse_string_meta(&nv) {
63			renames.push((param_name, value));
64		}
65	}
66	renames
67}
68
69fn parse_api_meta(meta: Meta, action: &mut ApiAction) {
70	match meta {
71		Meta::NameValue(nv) if nv.path.is_ident("extract") => {
72			action.extract = parse_string_meta(&nv);
73		}
74		Meta::NameValue(nv) if nv.path.is_ident("response") => {
75			action.response_type = parse_type_meta(&nv);
76		}
77		Meta::List(list) if list.path.is_ident("map") => {
78			action.renames = parse_map_list(&list);
79		}
80		_ => {}
81	}
82}
83
84fn take_api_action(attrs: &mut Vec<Attribute>) -> ApiAction {
85	let attr = match take_api_attr(attrs) {
86		Some(a) => a,
87		None => return ApiAction::default(),
88	};
89
90	let meta_list = match attr.meta.require_list() {
91		Ok(list) => list,
92		Err(_) => return ApiAction::default(),
93	};
94
95	let nested = match meta_list
96		.parse_args_with(syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated)
97	{
98		Ok(n) => n,
99		Err(_) => return ApiAction::default(),
100	};
101
102	let mut action = ApiAction::default();
103	for meta in nested {
104		parse_api_meta(meta, &mut action);
105	}
106
107	action
108}
109
110fn generate_method_body(method: &ImplItemFn, action: &ApiAction) -> proc_macro2::TokenStream {
111	let action_name = method.sig.ident.to_string();
112	let action_lit = syn::LitStr::new(&action_name, proc_macro2::Span::call_site());
113
114	let json_entries: Vec<proc_macro2::TokenStream> = method
115		.sig
116		.inputs
117		.iter()
118		.filter_map(|arg| {
119			if let FnArg::Typed(pat_type) = arg
120				&& let Pat::Ident(pat_ident) = &*pat_type.pat
121			{
122				let param_name = pat_ident.ident.to_string();
123				let json_key = action
124					.renames
125					.iter()
126					.find(|(p, _)| p == &param_name)
127					.map(|(_, k)| k.clone())
128					.unwrap_or_else(|| param_name.clone());
129
130				let param_ident = &pat_ident.ident;
131				let json_key_lit = syn::LitStr::new(&json_key, proc_macro2::Span::call_site());
132
133				return Some(quote! {
134					#json_key_lit: #param_ident,
135				});
136			}
137			None
138		})
139		.collect();
140
141	if let Some(ref extract_field) = action.extract {
142		let response_type = action
143			.response_type
144			.as_ref()
145			.expect("#[api(extract = \"...\")] 需要同时指定 #[api(response = Type)] 以确定中间响应类型");
146		let extract_ident = syn::Ident::new(extract_field, proc_macro2::Span::call_site());
147
148		quote!({
149			let params = ::serde_json::json!({
150				#(#json_entries)*
151			});
152			let response: #response_type = self.send_and_parse(#action_lit, params).await?;
153			Ok(response.#extract_ident)
154		})
155	} else {
156		quote!({
157			let params = ::serde_json::json!({
158				#(#json_entries)*
159			});
160			self.send_and_parse(#action_lit, params).await
161		})
162	}
163}
164
165#[proc_macro_attribute]
166pub fn api_sender(_attr: TokenStream, item: TokenStream) -> TokenStream {
167	let mut impl_block = parse_macro_input!(item as ItemImpl);
168
169	let mut new_items: Vec<ImplItem> = Vec::with_capacity(impl_block.items.len());
170
171	for item in std::mem::take(&mut impl_block.items) {
172		if let ImplItem::Fn(mut method) = item {
173			let api_action = take_api_action(&mut method.attrs);
174			let body = generate_method_body(&method, &api_action);
175			method.block = syn::parse2(body).expect("生成方法体失败");
176			new_items.push(ImplItem::Fn(method));
177		} else {
178			new_items.push(item);
179		}
180	}
181
182	impl_block.items = new_items;
183
184	quote!(#impl_block).into()
185}