cli_rs_command_gen/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::{Ident, Span};
3use quote::quote;
4
5const MAX_COMMANDS: u16 = 6;
6
7#[proc_macro]
8pub fn command(input: TokenStream) -> TokenStream {
9    let n: u16 = input.to_string().parse().unwrap();
10    if n > MAX_COMMANDS {
11        panic!("Max command size is {MAX_COMMANDS}");
12    }
13
14    let command_name = Ident::new(&format!("Command{n}"), Span::call_site());
15    let callback_name = Ident::new(&format!("Callback{n}"), Span::call_site());
16
17    // T1, T2, ... TN
18    let generics: &Vec<Ident> = &(1..=n)
19        .into_iter()
20        .map(|num| Ident::new(&format!("T{num}"), Span::call_site()))
21        .collect();
22
23    let ins: &Vec<Ident> = &(1..=n)
24        .into_iter()
25        .map(|num| Ident::new(&format!("in{num}"), Span::call_site()))
26        .collect();
27
28    // T1: Input, T2: Input, ... TN: Input
29    let generic_with_types = quote!(#(#generics: Input),*);
30
31    let where_generic_with_types = if n >= 1 {
32        quote! {
33            where
34                #generic_with_types
35        }
36    } else {
37        quote!()
38    };
39
40    // pub in1: T1,\npub in2: T2, ...
41    let struct_definition = quote! {
42        #(pub #ins: #generics,)*
43    };
44
45    let symbol_vec = quote! {
46        vec![#(&mut self.#ins),*]
47    };
48
49    let handler_call = quote! {
50        handler(#(&self.#ins),*)
51    };
52
53    let input_fn = if n == MAX_COMMANDS {
54        quote! {}
55    } else {
56        let next_command = Ident::new(&format!("Command{}", n + 1), Span::call_site());
57        let next_generic = Ident::new(&format!("T{}", n + 1), Span::call_site());
58        let next_in = Ident::new(&format!("in{}", n + 1), Span::call_site());
59        let ins_copy = &ins;
60
61        let in_transfer = quote! {
62            #(#ins: self.#ins_copy,)*
63        };
64
65        quote! {
66            pub fn input<#next_generic: Input>(self, #next_in: #next_generic) -> #next_command<'a, #(#generics,)* #next_generic> {
67                #next_command {
68                    docs: self.docs,
69                    handler: None,
70
71                    #in_transfer
72
73                    #next_in,
74
75                    subcommands: self.subcommands,
76                }
77            }
78        }
79    };
80
81    let command_0_fns = if n == 0 {
82        quote! {
83            pub fn name(name: &str) -> Self {
84                Self {
85                    docs: DocInfo {
86                        name: name.to_string(),
87                        ..Default::default()
88                    },
89                    subcommands: vec![],
90                    handler: None,
91                }
92            }
93
94            pub fn with_completions(self) -> Self {
95                let name = self.docs.name.clone();
96
97                self.subcommand(
98                    Self::name("completions")
99                        .description("generate completions for a given shell")
100                        .input(Arg::<CompletionMode>::name("shell").completor(|prompt| {
101                            Ok(["bash".to_string(), "zsh".to_string(), "fish".to_string()]
102                                .into_iter()
103                                .filter(|sh| sh.starts_with(prompt))
104                                .collect())
105                        }))
106                        .handler(move |shell| {
107                            shell.get().print_completion(&name);
108                            Ok(())
109                        }),
110                )
111            }
112
113            pub fn version(mut self, version: &str) -> Self {
114                self.docs.version = Some(version.to_string());
115                self
116            }
117
118            pub fn description(mut self, description: &str) -> Self {
119                self.docs.description = Some(description.to_string());
120                self
121            }
122        }
123    } else {
124        quote! {}
125    };
126
127    quote! {
128
129        type #callback_name<'a, #(#generics),* > = Box<dyn FnMut(#(&#generics),*) -> CliResult<()> + 'a>;
130        pub struct #command_name<'a, #generic_with_types> {
131            pub docs: DocInfo,
132
133            pub subcommands: Vec<Box<dyn Cmd + 'a>>,
134            pub handler: Option<#callback_name<'a, #( #generics),*>>,
135
136            #struct_definition
137        }
138
139        impl<'a, #( #generics ),*> ParserInfo for #command_name<'a, #( #generics ),*>
140            #where_generic_with_types
141        {
142            fn docs(&self) -> &DocInfo {
143                &self.docs
144            }
145
146            fn symbols(&mut self) -> Vec<&mut dyn Input> {
147                #symbol_vec
148            }
149
150            fn subcommand_docs(&self) -> Vec<DocInfo> {
151                self.subcommands.iter().map(|s| s.docs().clone()).collect()
152            }
153
154            fn call_handler(&mut self) -> CliResult<()> {
155                if let Some(handler) = &mut self.handler {
156                    #handler_call
157                } else {
158                    Err(CliError::from(format!(
159                        "No handler hooked up to {}",
160                        self.docs.cmd_path()
161                    )))
162                }
163            }
164
165            fn push_parent(&mut self, parents: &[String]) {
166                self.docs.parents.extend_from_slice(parents);
167            }
168
169            fn complete_subcommand(&mut self, sub_idx: usize, tokens: &[String]) -> Result<Vec<CompOut>, CliError> {
170                self.subcommands[sub_idx].complete_args(tokens)
171            }
172
173            fn parse_subcommand(&mut self, sub_idx: usize, tokens: &[String]) -> Result<(), CliError> {
174                self.subcommands[sub_idx].parse_args(tokens)
175            }
176        }
177
178        impl<'a, #generic_with_types> #command_name<'a, #( #generics ),*> {
179
180            #command_0_fns
181
182            #input_fn
183
184            pub fn handler<F>(mut self, handler: F) -> Self
185            where
186                F: FnMut(#(&#generics),*) -> CliResult<()> + 'a,
187            {
188                self.handler = Some(Box::new(handler));
189                self
190            }
191
192            pub fn subcommand<C: Cmd + 'a>(mut self, mut sub: C) -> Self {
193                sub.push_parent(&self.docs.parents);
194                sub.push_parent(&[self.docs.name.clone()]);
195                self.subcommands.push(Box::new(sub));
196                self
197            }
198        }
199    }.into()
200}