1#![deny(
4 missing_docs,
5 rust_2018_idioms, unused,
7 unused_import_braces,
8 unused_lifetimes,
9 unused_qualifications,
10 warnings,
11)]
12
13use {
14 itertools::Itertools as _,
15 proc_macro::TokenStream,
16 proc_macro2::Span,
17 quote::{
18 quote,
19 quote_spanned,
20 },
21 syn::{
22 *,
23 punctuated::Punctuated,
24 spanned::Spanned as _,
25 },
26};
27
28#[proc_macro_attribute]
38pub fn command(args: TokenStream, item: TokenStream) -> TokenStream {
39 let args = parse_macro_input!(args with Punctuated::<Meta, Token![,]>::parse_terminated);
40 let varargs = match args.into_iter().at_most_one() {
41 Ok(None) => false,
42 Ok(Some(arg)) if arg.path().is_ident("varargs") => true,
43 _ => return quote!(compile_error!("unexpected bitbar::command arguments");).into(),
44 };
45 let command_fn = parse_macro_input!(item as ItemFn);
46 let vis = &command_fn.vis;
47 let asyncness = &command_fn.sig.asyncness;
48 let command_name = &command_fn.sig.ident;
49 let command_name_str = command_name.to_string();
50 let wrapper_name = Ident::new(&format!("bitbar_{command_name}_wrapper"), Span::call_site());
51 let awaitness = asyncness.as_ref().map(|_| quote!(.await));
52 let (wrapper_body, command_params, command_args) = if varargs {
53 (
54 quote!(::bitbar::CommandOutput::report(#command_name(args)#awaitness, #command_name_str)),
55 quote!(::std::iter::Iterator::collect(::std::iter::Iterator::chain(::std::iter::once(::std::string::ToString::to_string(#command_name_str)), args))),
56 quote!(_: ::bitbar::flavor::SwiftBar, args: ::std::vec::Vec<::std::string::String>),
57 )
58 } else {
59 let mut wrapper_params = Vec::default();
60 let mut wrapped_args = Vec::default();
61 let mut command_params = Vec::default();
62 let mut command_args = Vec::default();
63 for (arg_idx, arg) in command_fn.sig.inputs.iter().enumerate() {
64 match arg {
65 FnArg::Receiver(_) => return quote_spanned! {arg.span()=>
66 compile_error("unexpected `self` parameter in bitbar::command");
67 }.into(),
68 FnArg::Typed(PatType { ty, .. }) => {
69 let ident = Ident::new(&format!("arg{}", arg_idx), arg.span());
70 wrapper_params.push(quote_spanned! {arg.span()=>
71 #ident
72 });
73 wrapped_args.push(quote_spanned! {arg.span()=>
74 match #ident.parse() {
75 ::core::result::Result::Ok(arg) => arg,
76 ::core::result::Result::Err(e) => {
77 ::bitbar::notify_error(
78 &::std::format!("{}: error parsing parameter {}: {}", #command_name_str, #arg_idx, e),
79 &::std::format!("{e:?}"),
80 );
81 ::std::process::exit(1)
82 }
83 }
84 });
85 command_params.push(quote_spanned! {arg.span()=>
86 #ident.to_string()
87 });
88 command_args.push(quote_spanned! {arg.span()=>
89 #ident: &#ty
90 });
91 }
92 }
93 }
94 if command_args.len() > 5 {
95 command_args.insert(0, quote!(_: ::bitbar::flavor::SwiftBar));
96 }
97 (
98 quote! {
99 match &*args {
100 [#(#wrapper_params),*] => ::bitbar::CommandOutput::report(#command_name(#(#wrapped_args),*)#awaitness, #command_name_str),
101 _ => {
102 ::bitbar::notify("wrong number of command arguments");
103 ::std::process::exit(1)
104 }
105 }
106 },
107 quote!(::std::vec![
108 ::std::string::ToString::to_string(#command_name_str),
109 #(#command_params,)*
110 ]),
111 quote!(#(#command_args),*),
112 )
113 };
114 #[cfg(not(feature = "tokio"))] let (wrapper_ret, wrapper_body) = (quote!(), wrapper_body);
115 #[cfg(feature = "tokio")] let (wrapper_ret, wrapper_body) = (
116 quote!(-> ::std::pin::Pin<::std::boxed::Box<dyn ::std::future::Future<Output = ()>>>),
117 quote!(::std::boxed::Box::pin(async move { #wrapper_body })),
118 );
119 TokenStream::from(quote! {
120 fn #wrapper_name(args: ::std::vec::Vec<::std::string::String>) #wrapper_ret {
121 #command_fn
122
123 #wrapper_body
124 }
125
126 #vis fn #command_name(#command_args) -> ::std::io::Result<::bitbar::attr::Params> {
127 ::std::io::Result::Ok(
128 ::bitbar::attr::Params::new(::std::env::current_exe()?.into_os_string().into_string().expect("non-UTF-8 plugin path"), #command_params)
129 )
130 }
131 })
132}
133
134#[proc_macro_attribute]
144pub fn fallback_command(_: TokenStream, item: TokenStream) -> TokenStream {
145 let fallback_fn = parse_macro_input!(item as ItemFn);
146 let asyncness = &fallback_fn.sig.asyncness;
147 let fn_name = &fallback_fn.sig.ident;
148 let wrapper_name = Ident::new(&format!("bitbar_{fn_name}_wrapper"), Span::call_site());
149 let awaitness = asyncness.as_ref().map(|_| quote!(.await));
150 let wrapper_body = quote! {
151 ::bitbar::CommandOutput::report(#fn_name(cmd.clone(), args)#awaitness, &cmd);
152 };
153 #[cfg(not(feature = "tokio"))] let (wrapper_ret, wrapper_body) = (quote!(), wrapper_body);
154 #[cfg(feature = "tokio")] let (wrapper_ret, wrapper_body) = (
155 quote!(-> ::std::pin::Pin<::std::boxed::Box<dyn ::std::future::Future<Output = ()>>>),
156 quote!(::std::boxed::Box::pin(async move { #wrapper_body })),
157 );
158 TokenStream::from(quote! {
159 fn #wrapper_name(cmd: ::std::string::String, args: ::std::vec::Vec<::std::string::String>) #wrapper_ret {
160 #fallback_fn
161
162 #wrapper_body
163 }
164 })
165}
166
167#[proc_macro_attribute]
179pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
180 let args = parse_macro_input!(args with Punctuated::<Meta, Token![,]>::parse_terminated);
181 let mut error_template_image = quote!(::core::option::Option::None);
182 let mut fallback_lit = None;
183 let mut subcommand_names = Vec::default();
184 let mut subcommand_fns = Vec::default();
185 for arg in args {
186 if arg.path().is_ident("commands") {
187 match arg.require_list() {
188 Ok(list) => match list.parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated) {
189 Ok(nested) => for cmd in nested {
190 match cmd.require_path_only() {
191 Ok(path) => if let Some(ident) = path.get_ident() {
192 subcommand_names.push(ident.to_string());
193 subcommand_fns.push(Ident::new(&format!("bitbar_{ident}_wrapper"), ident.span()));
194 } else {
195 return quote_spanned! {cmd.span()=>
196 compile_error!("bitbar subcommands must be simple identifiers");
197 }.into()
198 },
199 Err(e) => return e.into_compile_error().into(),
200 }
201 },
202 Err(e) => return e.into_compile_error().into(),
203 }
204 Err(e) => return e.into_compile_error().into(),
205 }
206 } else if arg.path().is_ident("error_template_image") {
207 match arg.require_name_value() {
208 Ok(MetaNameValue { value, .. }) => if let Expr::Lit(ExprLit { lit: Lit::Str(lit), .. }) = value {
209 error_template_image = quote!(::core::option::Option::Some(::bitbar::attr::Image::from(&include_bytes!(#lit)[..])));
210 } else {
211 return quote_spanned! {value.span()=>
212 compile_error!("error_template_image value must be a string literal");
213 }.into()
214 },
215 Err(e) => return e.into_compile_error().into(),
216 }
217 } else if arg.path().is_ident("fallback_command") {
218 match arg.require_name_value() {
219 Ok(MetaNameValue { value, .. }) => if let Expr::Lit(ExprLit { lit: Lit::Str(lit), .. }) = value {
220 fallback_lit = Some(Ident::new(&format!("bitbar_{}_wrapper", lit.value()), lit.span()));
221 } else {
222 return quote_spanned! {value.span()=>
223 compile_error!("fallback_command value must be a string literal");
224 }.into()
225 },
226 Err(e) => return e.into_compile_error().into(),
227 }
228 } else {
229 return quote_spanned! {arg.span()=>
230 compile_error!("unexpected bitbar::main attribute argument");
231 }.into()
232 }
233 }
234 let main_fn = parse_macro_input!(item as ItemFn);
235 let asyncness = &main_fn.sig.asyncness;
236 let inner_params = &main_fn.sig.inputs;
237 let inner_args = if inner_params.len() >= 1 {
238 quote!(::bitbar::Flavor::check())
239 } else {
240 quote!()
241 };
242 #[cfg(not(feature = "tokio"))] let (cmd_awaitness, wrapper_body) = (
243 quote!(),
244 quote!(::bitbar::MainOutput::main_output(main_inner(#inner_args), #error_template_image);),
245 );
246 #[cfg(feature = "tokio")] let awaitness = asyncness.as_ref().map(|_| quote!(.await));
247 #[cfg(feature = "tokio")] let (cmd_awaitness, wrapper_body) = (
248 quote!(.await),
249 quote!(::bitbar::AsyncMainOutput::main_output(main_inner(#inner_args)#awaitness, #error_template_image).await;),
250 );
251 let fallback = if let Some(fallback_lit) = fallback_lit {
252 quote!(#fallback_lit(subcommand, args.collect())#cmd_awaitness)
253 } else {
254 quote! {{
255 ::bitbar::notify(format!("no such subcommand: {}", subcommand));
256 ::std::process::exit(1)
257 }}
258 };
259 let wrapper_body = quote!({
260 let mut args = ::std::env::args();
262 let _ = args.next().expect("missing program name");
263 if let ::core::option::Option::Some(subcommand) = args.next() {
264 match &*subcommand {
265 #(
266 #subcommand_names => #subcommand_fns(args.collect())#cmd_awaitness,
267 )*
268 _ => #fallback,
269 }
270 } else {
271 #wrapper_body
272 }
273 });
274 #[cfg(feature = "tokio")] let wrapper_body = quote!({
275 ::bitbar::tokio::runtime::Builder::new_multi_thread()
276 .enable_all()
277 .build()
278 .unwrap()
279 .block_on(async #wrapper_body)
280 });
281 let ret = main_fn.sig.output;
282 let inner_body = main_fn.block;
283 TokenStream::from(quote! {
284 #asyncness fn main_inner(#inner_params) #ret #inner_body
285
286 fn main() #wrapper_body
287 })
288}