cmdstruct_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{
4    parse_macro_input, spanned::Spanned, AttrStyle, Attribute, Data, DataStruct, DeriveInput,
5    Field, Fields, FieldsNamed, Ident, LitStr, MetaList, Path, Type,
6};
7
8type Result<T> = std::result::Result<T, syn::Error>;
9
10#[proc_macro_derive(Command, attributes(command, arg))]
11pub fn command(input: TokenStream) -> TokenStream {
12    let derive_input = parse_macro_input!(input as DeriveInput);
13    match Command::parse(derive_input) {
14        Ok(command) => command.into(),
15        Err(err) => err.into_compile_error().into(),
16    }
17}
18
19enum Executable {
20    Const(String),
21    Function(Path),
22}
23
24struct CommandAttributes {
25    executable: Executable,
26}
27
28impl CommandAttributes {
29    fn parse(derive_input: &DeriveInput) -> Result<Self> {
30        let mut executable = None;
31        for attr in &derive_input.attrs {
32            if attr.path().is_ident("command") {
33                match &attr.meta {
34                    syn::Meta::List(MetaList {
35                        path: _,
36                        delimiter: _,
37                        tokens: _,
38                    }) => {
39                        attr.parse_nested_meta(|meta| {
40                            if meta.path.is_ident("executable") {
41                                let value = meta.value()?;
42                                let s: LitStr = value.parse()?;
43                                executable = Some(Executable::Const(s.value()));
44                                Ok(())
45                            } else if meta.path.is_ident("executable_fn") {
46                                let value = meta.value()?;
47                                let s: Path = value.parse()?;
48                                executable = Some(Executable::Function(s));
49                                Ok(())
50                            } else {
51                                return Err(syn::Error::new(attr.span(), "Unsupported attribute"));
52                            }
53                        })?;
54                    }
55                    _ => {}
56                }
57            }
58        }
59        if let Some(executable) = executable {
60            Ok(Self { executable })
61        } else {
62            Err(syn::Error::new(
63                derive_input.span(),
64                "No 'executable' defined for 'command'",
65            ))
66        }
67    }
68}
69
70struct Command {
71    attributes: CommandAttributes,
72    ident: Ident,
73    args: Vec<Arg>,
74}
75
76impl Command {
77    fn parse(derive_input: DeriveInput) -> Result<Command> {
78        let attributes = CommandAttributes::parse(&derive_input)?;
79
80        let args = match derive_input.data {
81            Data::Struct(DataStruct {
82                struct_token: _,
83                fields:
84                    Fields::Named(FieldsNamed {
85                        brace_token: _,
86                        mut named,
87                    }),
88                semi_token: _,
89            }) => named.iter_mut().filter_map(collect_arg).collect(),
90            _ => Err(syn::Error::new(
91                derive_input.span(),
92                "Only structs with named fields supported.",
93            )),
94        }?;
95        Ok(Command {
96            attributes,
97            ident: derive_input.ident.clone(),
98            args,
99        })
100    }
101}
102
103enum ArgType {
104    Option { name: String },
105    Flag { name: String },
106    Positional,
107}
108
109#[allow(dead_code)]
110struct Arg {
111    arg_type: ArgType,
112    ident: Ident,
113    ty: Type,
114}
115
116type ArgResult = Result<(Option<Attribute>, Option<ArgType>)>;
117
118fn parse_arg_with_attributes(attr: Attribute) -> ArgResult {
119    let mut arg_type = None;
120    attr.parse_nested_meta(|meta| {
121        if meta.path.is_ident("option") {
122            if arg_type.is_none() {
123                let value = meta.value()?;
124                let s: LitStr = value.parse()?;
125                arg_type = Some(ArgType::Option { name: s.value() });
126                Ok(())
127            } else {
128                Err(meta.error("Only one argument type allowed."))
129            }
130        } else if meta.path.is_ident("flag") {
131            if arg_type.is_none() {
132                let value = meta.value()?;
133                let s: LitStr = value.parse()?;
134                arg_type = Some(ArgType::Flag { name: s.value() });
135                Ok(())
136            } else {
137                Err(meta.error("Only one argument type allowed."))
138            }
139        } else {
140            Err(meta.error("Unrecognized arg"))
141        }
142    })
143    .map(|_| arg_type.map_or((Some(attr), None), |arg_type| (None, Some(arg_type))))
144}
145
146fn map_to_attr_or_arg(attr: Attribute) -> ArgResult {
147    match attr.style {
148        AttrStyle::Outer => match &attr.meta {
149            syn::Meta::List(list) if list.path.is_ident("arg") => parse_arg_with_attributes(attr),
150            syn::Meta::Path(path) if path.is_ident("arg") => Ok((None, Some(ArgType::Positional))),
151            _ => Ok((Some(attr), None)),
152        },
153        _ => Ok((Some(attr), None)),
154    }
155}
156
157fn collect_arg(field: &mut Field) -> Option<Result<Arg>> {
158    if let Some(ident) = &field.ident {
159        let arg_results: Result<Vec<_>> = field
160            .attrs
161            .clone()
162            .into_iter()
163            .map(map_to_attr_or_arg)
164            .collect();
165        match arg_results {
166            Ok(results) => {
167                let unzipped: (Vec<_>, Vec<_>) = results.into_iter().unzip();
168                match unzipped {
169                    (attrs, arg_types) => {
170                        let attrs: Vec<_> = attrs.into_iter().filter_map(|attr| attr).collect();
171                        let mut arg_types: Vec<_> = arg_types
172                            .into_iter()
173                            .filter_map(|arg_type| arg_type)
174                            .collect();
175                        field.attrs = attrs;
176                        match arg_types.len() {
177                            1 => Some(Ok(Arg {
178                                arg_type: arg_types.remove(0),
179                                ident: ident.clone(),
180                                ty: field.ty.clone(),
181                            })),
182                            0 => None,
183                            _ => Some(Err(syn::Error::new(field.span(), "Too many args"))),
184                        }
185                    }
186                }
187            }
188            Err(err) => Some(Err(err)),
189        }
190    } else {
191        None
192    }
193}
194
195fn append_arg_tokens(arg: &Arg) -> proc_macro2::TokenStream {
196    let ident = &arg.ident;
197    match &arg.arg_type {
198        ArgType::Option { name } => quote! {
199            cmdstruct::Arg::append_option(&self.#ident, #name, &mut command);
200        },
201        ArgType::Flag { name } => quote! {
202            if self.#ident {
203                command.arg(#name);
204            }
205        },
206        ArgType::Positional => quote! {
207            cmdstruct::Arg::append_arg(&self.#ident, &mut command);
208        },
209    }
210}
211
212impl Into<TokenStream> for Command {
213    fn into(self) -> TokenStream {
214        let args: Vec<_> = self.args.iter().map(append_arg_tokens).collect();
215        let executable = match &self.attributes.executable {
216            Executable::Const(executable) => quote! { #executable },
217            Executable::Function(func) => quote! { #func(&self) },
218        };
219        let struct_ident = &self.ident;
220        let impls_combined = quote! {
221
222            impl cmdstruct::Command for #struct_ident {
223
224                fn command(&self) -> std::process::Command {
225                    let mut command = std::process::Command::new(#executable);
226                    #(#args)*
227                    command
228                }
229            }
230        };
231        impls_combined.into()
232    }
233}