cmd_lib_macros/
lib.rs

1use proc_macro2::{TokenStream, TokenTree};
2use proc_macro_error2::{abort, proc_macro_error};
3use quote::quote;
4
5/// Mark main function to log error result by default.
6///
7/// ```no_run
8/// # use cmd_lib::*;
9///
10/// #[cmd_lib::main]
11/// fn main() -> CmdResult {
12///     run_cmd!(bad_cmd)?;
13///     Ok(())
14/// }
15/// // output:
16/// // [ERROR] FATAL: Running ["bad_cmd"] failed: No such file or directory (os error 2)
17/// ```
18#[proc_macro_attribute]
19pub fn main(
20    _args: proc_macro::TokenStream,
21    item: proc_macro::TokenStream,
22) -> proc_macro::TokenStream {
23    let orig_function: syn::ItemFn = syn::parse2(item.into()).unwrap();
24    let orig_main_return_type = orig_function.sig.output;
25    let orig_main_block = orig_function.block;
26
27    quote! (
28        fn main() {
29            fn cmd_lib_main() #orig_main_return_type {
30                #orig_main_block
31            }
32
33            cmd_lib_main().unwrap_or_else(|err| {
34                ::cmd_lib::error!("FATAL: {err}");
35                std::process::exit(1);
36            });
37        }
38
39    )
40    .into()
41}
42
43/// Import user registered custom command.
44/// ```no_run
45/// # use cmd_lib::*;
46/// # use std::io::Write;
47/// fn my_cmd(env: &mut CmdEnv) -> CmdResult {
48///     let msg = format!("msg from foo(), args: {:?}", env.get_args());
49///     writeln!(env.stderr(), "{msg}")?;
50///     writeln!(env.stdout(), "bar")
51/// }
52///
53/// use_custom_cmd!(my_cmd);
54/// run_cmd!(my_cmd)?;
55/// # Ok::<(), std::io::Error>(())
56/// ```
57/// Here we import the previous defined `my_cmd` command, so we can run it like a normal command.
58#[proc_macro]
59#[proc_macro_error]
60pub fn use_custom_cmd(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
61    let item: proc_macro2::TokenStream = item.into();
62    let mut cmd_fns = vec![];
63    for t in item {
64        if let TokenTree::Punct(ref ch) = t {
65            if ch.as_char() != ',' {
66                abort!(t, "only comma is allowed");
67            }
68        } else if let TokenTree::Ident(cmd) = t {
69            let cmd_name = cmd.to_string();
70            cmd_fns.push(quote!(&#cmd_name, #cmd));
71        } else {
72            abort!(t, "expect a list of comma separated commands");
73        }
74    }
75
76    quote! (
77        #(::cmd_lib::register_cmd(#cmd_fns);)*
78    )
79    .into()
80}
81
82/// Run commands, returning [`CmdResult`](../cmd_lib/type.CmdResult.html) to check status.
83/// ```no_run
84/// # use cmd_lib::run_cmd;
85/// let msg = "I love rust";
86/// run_cmd!(echo $msg)?;
87/// run_cmd!(echo "This is the message: $msg")?;
88///
89/// // pipe commands are also supported
90/// run_cmd!(du -ah . | sort -hr | head -n 10)?;
91///
92/// // or a group of commands
93/// // if any command fails, just return Err(...)
94/// let file = "/tmp/f";
95/// let keyword = "rust";
96/// if run_cmd! {
97///     cat ${file} | grep ${keyword};
98///     echo "bad cmd" >&2;
99///     ignore ls /nofile;
100///     date;
101///     ls oops;
102///     cat oops;
103/// }.is_err() {
104///     // your error handling code
105/// }
106/// # Ok::<(), std::io::Error>(())
107/// ```
108#[proc_macro]
109#[proc_macro_error]
110pub fn run_cmd(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
111    let cmds = lexer::Lexer::new(input.into()).scan().parse(false);
112    quote! ({
113        use ::cmd_lib::AsOsStr;
114        #cmds.run_cmd()
115    })
116    .into()
117}
118
119/// Run commands, returning [`FunResult`](../cmd_lib/type.FunResult.html) to capture output and to check status.
120/// ```no_run
121/// # use cmd_lib::run_fun;
122/// let version = run_fun!(rustc --version)?;
123/// println!("Your rust version is {}", version);
124///
125/// // with pipes
126/// let n = run_fun!(echo "the quick brown fox jumped over the lazy dog" | wc -w)?;
127/// println!("There are {} words in above sentence", n);
128/// # Ok::<(), std::io::Error>(())
129/// ```
130#[proc_macro]
131#[proc_macro_error]
132pub fn run_fun(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
133    let cmds = lexer::Lexer::new(input.into()).scan().parse(false);
134    quote! ({
135        use ::cmd_lib::AsOsStr;
136        #cmds.run_fun()
137    })
138    .into()
139}
140
141/// Run commands with/without pipes as a child process, returning [`CmdChildren`](../cmd_lib/struct.CmdChildren.html) result.
142/// ```no_run
143/// # use cmd_lib::*;
144///
145/// let mut handle = spawn!(ping -c 10 192.168.0.1)?;
146/// // ...
147/// if handle.wait().is_err() {
148///     // ...
149/// }
150/// # Ok::<(), std::io::Error>(())
151#[proc_macro]
152#[proc_macro_error]
153pub fn spawn(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
154    let cmds = lexer::Lexer::new(input.into()).scan().parse(true);
155    quote! ({
156        use ::cmd_lib::AsOsStr;
157        #cmds.spawn(false)
158    })
159    .into()
160}
161
162/// Run commands with/without pipes as a child process, returning [`FunChildren`](../cmd_lib/struct.FunChildren.html) result.
163/// ```no_run
164/// # use cmd_lib::*;
165/// let mut procs = vec![];
166/// for _ in 0..4 {
167///     let proc = spawn_with_output!(
168///         sudo bash -c "dd if=/dev/nvmen0 of=/dev/null bs=4096 skip=0 count=1024 2>&1"
169///         | awk r#"/copied/{print $(NF-1) " " $NF}"#
170///     )?;
171///     procs.push(proc);
172/// }
173///
174/// for (i, mut proc) in procs.into_iter().enumerate() {
175///     let bandwidth = proc.wait_with_output()?;
176///     info!("thread {i} bandwidth: {bandwidth} MB/s");
177/// }
178/// # Ok::<(), std::io::Error>(())
179/// ```
180#[proc_macro]
181#[proc_macro_error]
182pub fn spawn_with_output(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
183    let cmds = lexer::Lexer::new(input.into()).scan().parse(true);
184    quote! ({
185        use ::cmd_lib::AsOsStr;
186        #cmds.spawn_with_output()
187    })
188    .into()
189}
190
191#[proc_macro]
192#[proc_macro_error]
193/// Log a fatal message at the error level, and exit process.
194///
195/// e.g:
196/// ```no_run
197/// # use cmd_lib::*;
198/// let file = "bad_file";
199/// cmd_die!("could not open file: $file");
200/// // output:
201/// // [ERROR] FATAL: could not open file: bad_file
202/// ```
203/// format should be string literals, and variable interpolation is supported.
204/// Note that this macro is just for convenience. The process will exit with 1 and print
205/// "FATAL: ..." messages to error console. If you want to exit with other code, you
206/// should probably define your own macro or functions.
207pub fn cmd_die(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
208    let msg = parse_msg(input.into());
209    quote!({
210        ::cmd_lib::error!("FATAL: {} at {}:{}", #msg, file!(), line!());
211        std::process::exit(1)
212    })
213    .into()
214}
215
216fn parse_msg(input: TokenStream) -> TokenStream {
217    let mut iter = input.into_iter();
218    let mut output = TokenStream::new();
219    let mut valid = false;
220    if let Some(ref tt) = iter.next() {
221        if let TokenTree::Literal(lit) = tt {
222            let s = lit.to_string();
223            if s.starts_with('\"') || s.starts_with('r') {
224                let str_lit = lexer::scan_str_lit(lit);
225                output.extend(quote!(#str_lit));
226                valid = true;
227            }
228        }
229        if !valid {
230            abort!(tt, "invalid format: expect string literal");
231        }
232        if let Some(tt) = iter.next() {
233            abort!(
234                tt,
235                "expect string literal only, found extra {}",
236                tt.to_string()
237            );
238        }
239    }
240    output
241}
242
243mod lexer;
244mod parser;
cmd_lib_macros/lib.rs

cmd_lib_macros/
lib.rs
