shx_macros/
lib.rs

1use proc_macro2::{Group, TokenStream, TokenTree};
2use quote::{quote, ToTokens, TokenStreamExt};
3
4enum Arg {
5    Literal(String),
6    Expr(TokenStream),
7    Variadic(TokenStream),
8}
9
10enum ParseState {
11    Cmd,
12    Args,
13    Variadic(usize),
14    SetSink,
15    DoneSetSink,
16    SetSource,
17    DoneSetSource,
18}
19
20enum Sink {
21    File(String),
22    Expr(TokenStream),
23}
24
25enum Source {
26    File(String),
27    Expr(TokenStream),
28}
29
30struct CmdParser {
31    state: ParseState,
32    cmd: Option<String>,
33    args: Vec<Arg>,
34    sink: Option<Sink>,
35    source: Option<Source>,
36}
37
38struct Cmd {
39    cmd: String,
40    args: Vec<Arg>,
41    sink: Option<Sink>,
42    source: Option<Source>,
43}
44
45#[derive(Debug)]
46enum CmdTokenTree {
47    Value(String),
48    EndOfLine,
49    Expr(Group),
50    Dot,
51    Sink,
52    Source,
53}
54
55impl From<TokenTree> for CmdTokenTree {
56    fn from(value: TokenTree) -> Self {
57        match value {
58            TokenTree::Group(g) => CmdTokenTree::Expr(g),
59            TokenTree::Ident(value) => CmdTokenTree::Value(value.to_string()),
60            TokenTree::Punct(c) if c.as_char() == ';' => CmdTokenTree::EndOfLine,
61            TokenTree::Punct(c) if c.as_char() == '>' => CmdTokenTree::Sink,
62            TokenTree::Punct(c) if c.as_char() == '<' => CmdTokenTree::Source,
63            TokenTree::Punct(c) if c.as_char() == '.' => CmdTokenTree::Dot,
64            TokenTree::Punct(c) => panic!("Unexpected punctuation character: {c}"),
65            TokenTree::Literal(value) => {
66                let literal = litrs::Literal::from(value);
67                let value = match literal {
68                    litrs::Literal::Bool(b) => b.to_string(),
69                    litrs::Literal::Integer(i) => i.to_string(),
70                    litrs::Literal::Float(f) => f.to_string(),
71                    litrs::Literal::Char(_c) => {
72                        unimplemented!("Character literals are not implemented")
73                    }
74                    litrs::Literal::String(s) => s.into_value().into_owned(),
75                    litrs::Literal::Byte(_b) => unimplemented!("Byte literals are not implemented"),
76                    litrs::Literal::ByteString(_s) => {
77                        unimplemented!("Byte literals are not implemented")
78                    }
79                };
80                CmdTokenTree::Value(value)
81            }
82        }
83    }
84}
85
86#[must_use]
87enum ParseResult {
88    KeepGoing,
89    Done(Cmd),
90}
91
92impl Default for CmdParser {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98impl CmdParser {
99    pub const VALID_VARIADIC_DOTS_COUNT: usize = 3;
100
101    pub fn new() -> Self {
102        Self {
103            state: ParseState::Cmd,
104            cmd: None,
105            args: Vec::new(),
106            sink: None,
107            source: None,
108        }
109    }
110
111    pub fn feed(&mut self, token: TokenTree) -> ParseResult {
112        let token = CmdTokenTree::from(token);
113        match self.state {
114            ParseState::Cmd => {
115                let CmdTokenTree::Value(value) = token else {
116                    panic!("Unexpected command: {token:?}");
117                };
118                self.cmd = Some(value);
119                self.state = ParseState::Args;
120            }
121            ParseState::Args => match token {
122                CmdTokenTree::Value(v) => self.args.push(Arg::Literal(v)),
123                CmdTokenTree::EndOfLine => return ParseResult::Done(self.take()),
124                CmdTokenTree::Expr(g) => self.args.push(Arg::Expr(g.stream())),
125                CmdTokenTree::Sink => self.state = ParseState::SetSink,
126                CmdTokenTree::Source => self.state = ParseState::SetSource,
127                CmdTokenTree::Dot => self.state = ParseState::Variadic(1),
128            },
129            ParseState::Variadic(n) => {
130                self.state = match token {
131                    CmdTokenTree::Dot => ParseState::Variadic(n + 1),
132                    CmdTokenTree::Expr(g) if n == Self::VALID_VARIADIC_DOTS_COUNT => {
133                        (*self).args.push(Arg::Variadic(g.stream()));
134                        ParseState::Args
135                    }
136                    CmdTokenTree::Expr(_) => panic!(
137                        "Variadic expression dots count (x{n}) is unexpected, expect x{}",
138                        Self::VALID_VARIADIC_DOTS_COUNT
139                    ),
140                    other => panic!("Expected variadic expression token, but got: {other:?}"),
141                };
142            }
143            ParseState::SetSink => {
144                assert!(self.sink.is_none(), "Can't set the sink more than once");
145                match token {
146                    CmdTokenTree::Value(v) => self.sink = Some(Sink::File(v)),
147                    CmdTokenTree::Expr(g) => self.sink = Some(Sink::Expr(g.stream())),
148                    other => panic!("Unexpected token: {other:?}"),
149                }
150                self.state = ParseState::DoneSetSink;
151            }
152            ParseState::SetSource => {
153                assert!(self.source.is_none(), "Can't set the source more than once");
154                match token {
155                    CmdTokenTree::Value(v) => self.source = Some(Source::File(v)),
156                    CmdTokenTree::Expr(g) => {
157                        self.source = Some(Source::Expr(g.stream()));
158                    }
159                    other => panic!("Unexpected token: {other:?}"),
160                }
161                self.state = ParseState::DoneSetSource;
162            }
163            ParseState::DoneSetSink => match token {
164                CmdTokenTree::EndOfLine => return ParseResult::Done(self.take()),
165                CmdTokenTree::Source => self.state = ParseState::SetSource,
166                other => panic!("Unexpected token: {other:?}"),
167            },
168            ParseState::DoneSetSource => match token {
169                CmdTokenTree::EndOfLine => return ParseResult::Done(self.take()),
170                CmdTokenTree::Sink => self.state = ParseState::SetSink,
171                other => panic!("Unexpected token: {other:?}"),
172            },
173        }
174        ParseResult::KeepGoing
175    }
176
177    fn finish(mut self) -> Option<Cmd> {
178        if self.cmd.is_some() {
179            Some(self.take())
180        } else {
181            None
182        }
183    }
184
185    fn take(&mut self) -> Cmd {
186        let mut parser = std::mem::take(self);
187        Cmd {
188            cmd: parser.cmd.take().expect("Missing command"),
189            args: parser.args,
190            sink: parser.sink,
191            source: parser.source,
192        }
193    }
194}
195
196/// A command-running macro.
197///
198/// `cmd` is a macro for running external commands. It provides functionality to
199/// pipe the input and output to/from variables as well as using rust expressions
200/// as arguments to the program.
201///
202/// The format of a `cmd` call is like so:
203///
204/// ```ignore
205/// lex!( [prog] [arg]* [< {in_expr}]? [> {out_expr}]? [;]? )
206/// ```
207///
208/// Or you can create multiple commands on a single block
209///
210/// ```ignore
211/// lex! {
212///   [prog] [arg]* [< {in_expr}]? [> {out_expr}]? ;
213///   [prog] [arg]* [< {in_expr}]? [> {out_expr}]? ;
214///   [prog] [arg]* [< {in_expr}]? [> {out_expr}]? [;]?
215/// }
216/// ```
217///
218/// Arguments are allowed to take the form of identifiers (i.e. plain text),
219/// literals (numbers, quoted strings, characters, etc.), or rust expressions
220/// delimited by braces.
221///
222/// This macro doesn't execute the commands. It returns an iterator of `shx::Cmd` which
223/// can be executed. Alternatively, see `shx::sh` which executes the commands sequentially.
224///
225/// # Examples
226///
227/// ```ignore
228/// # use shx_macros::lex;
229/// # #[cfg(target_os = "linux")]
230/// # fn run() {
231/// let world = "world";
232/// let mut out = String::new();
233/// lex!(echo hello {world} > {&mut out}).for_each(|cmd| cmd.exec().unwrap());
234/// assert_eq!(out, "hello world\n");
235/// # }
236/// # run();
237/// ```
238///
239/// ```ignore
240/// # use shx_macros::lex;
241/// # #[cfg(target_os = "linux")]
242/// # fn run() {
243/// lex! {
244///   echo hello;
245///   sleep 5;
246///   echo world;
247/// }.for_each(|cmd| cmd.exec().unwrap()); // prints hello, waits 5 seconds, prints world.
248/// # }
249/// # run();
250/// ```
251///
252/// You can also use string literals as needed
253///
254/// ```ignore
255/// # use shx_macros::lex;
256/// # #[cfg(target_os = "linux")]
257/// # fn run() {
258/// let mut out = String::new();
259/// lex!(echo "hello world" > {&mut out}).for_each(|cmd| cmd.exec().unwrap());
260/// assert_eq!(out, "hello world\n");
261/// # }
262/// # run();
263/// ```
264#[proc_macro]
265pub fn lex(stream: proc_macro::TokenStream) -> proc_macro::TokenStream {
266    let stream: TokenStream = stream.into();
267    let stream = stream.into_iter();
268    let mut cmd_list: Vec<Cmd> = Vec::new();
269
270    let mut parser = CmdParser::new();
271    for token in stream {
272        match parser.feed(token) {
273            ParseResult::KeepGoing => {}
274            ParseResult::Done(sh) => cmd_list.push(sh),
275        }
276    }
277    if let Some(cmd) = parser.finish() {
278        cmd_list.push(cmd);
279    }
280
281    quote!(
282        {
283            let mut __commands = Vec::new();
284            #(
285                __commands.push({
286                    #cmd_list
287                });
288            )*
289            __commands.into_iter()
290        }
291    )
292    .into()
293}
294
295impl ToTokens for Cmd {
296    fn to_tokens(&self, tokens: &mut TokenStream) {
297        let Self {
298            cmd,
299            args,
300            sink,
301            source,
302            ..
303        } = self;
304
305        let args = args.iter().map(|arg| match arg {
306            Arg::Literal(s) => quote!({ __cmd.arg({ #s }) }),
307            Arg::Expr(ts) => quote!({ __cmd.arg({ #ts }) }),
308            Arg::Variadic(ts) => quote!({ __cmd.args({ #ts }) }),
309        });
310
311        tokens.append_all(quote! {
312            let mut __cmd = ::std::process::Command::new(#cmd);
313            #( #args; )*
314            let mut __builder = ::shx::CmdBuilder::new(__cmd);
315        });
316
317        match sink {
318            Some(Sink::File(path)) => {
319                unimplemented!("Writing command output to file is not yet implemented: {path:?}")
320            }
321            Some(Sink::Expr(expr)) => tokens.append_all(quote! {
322                __builder.sink({ #expr });
323            }),
324            None => {}
325        }
326        match source {
327            Some(Source::File(path)) => {
328                unimplemented!("Reading command input from file is not yet implemented: {path:?}");
329            }
330            Some(Source::Expr(expr)) => tokens.append_all(quote! {
331                __builder.source({ #expr });
332            }),
333            None => {}
334        }
335
336        tokens.append_all(quote! {
337            __builder.build()
338        })
339    }
340}