Skip to main content

onebot_api_macros/
lib.rs

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