Skip to content

Commit

Permalink
clean up custom command registration
Browse files Browse the repository at this point in the history
  • Loading branch information
tao-guo committed Nov 23, 2023
1 parent a150ce2 commit 9c2a420
Show file tree
Hide file tree
Showing 5 changed files with 15 additions and 59 deletions.
5 changes: 2 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -228,11 +228,10 @@ run_cmd!(info "This is an infomation message")?;
```

#### Macros to register your own commands
Declare your function with `#[export_cmd(..)]` attribute, and import it with [`use_custom_cmd!`](https://docs.rs/cmd_lib/latest/cmd_lib/macro.use_custom_cmd.html) macro:
Declare your function with the right signature, and register it with [`use_custom_cmd!`](https://docs.rs/cmd_lib/latest/cmd_lib/macro.use_custom_cmd.html) macro:

```rust
#[export_cmd(my_cmd)]
fn foo(env: &mut CmdEnv) -> CmdResult {
fn my_cmd(env: &mut CmdEnv) -> CmdResult {
let msg = format!("msg from foo(), args: {:?}", env.args());
writeln!(env.stderr(), "{}", msg)?;
writeln!(env.stdout(), "bar")
Expand Down
52 changes: 6 additions & 46 deletions macros/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use proc_macro2::{Span, TokenStream, TokenTree};
use proc_macro2::{TokenStream, TokenTree};
use proc_macro_error::{abort, proc_macro_error};
use quote::{quote, ToTokens};
use quote::quote;

/// Mark main function to log error result by default.
///
Expand Down Expand Up @@ -40,50 +40,10 @@ pub fn main(
.into()
}

/// Export the function as an command to be run by [`run_cmd!`] or [`run_fun!`].
///
/// ```
/// # use cmd_lib::*;
/// # use std::io::Write;
/// #[export_cmd(my_cmd)]
/// fn foo(env: &mut CmdEnv) -> CmdResult {
/// let msg = format!("msg from foo(), args: {:?}\n", env.args());
/// writeln!(env.stderr(), "{}", msg)?;
/// writeln!(env.stdout(), "bar")
/// }
///
/// use_custom_cmd!(my_cmd);
/// run_cmd!(my_cmd)?;
/// println!("get result: {}", run_fun!(my_cmd)?);
/// # Ok::<(), std::io::Error>(())
/// ```
/// Here we export function `foo` as `my_cmd` command.
#[proc_macro_attribute]
pub fn export_cmd(
attr: proc_macro::TokenStream,
item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
let cmd_name = attr.to_string();
let export_cmd_fn = syn::Ident::new(&format!("export_cmd_{}", cmd_name), Span::call_site());

let orig_function: syn::ItemFn = syn::parse2(item.into()).unwrap();
let fn_ident = &orig_function.sig.ident;

let mut new_functions = orig_function.to_token_stream();
new_functions.extend(quote! (
fn #export_cmd_fn() {
export_cmd(#cmd_name, #fn_ident);
}
));
new_functions.into()
}

/// Import user registered custom command.
/// ```
/// # use cmd_lib::*;
/// #[export_cmd(my_cmd)]
/// fn foo(env: &mut CmdEnv) -> CmdResult {
/// fn my_cmd(env: &mut CmdEnv) -> CmdResult {
/// let msg = format!("msg from foo(), args: {:?}\n", env.args());
/// writeln!(env.stderr(), "{}", msg)?;
/// writeln!(env.stdout(), "bar")
Expand All @@ -105,15 +65,15 @@ pub fn use_custom_cmd(item: proc_macro::TokenStream) -> proc_macro::TokenStream
abort!(t, "only comma is allowed");
}
} else if let TokenTree::Ident(cmd) = t {
let cmd_fn = syn::Ident::new(&format!("export_cmd_{}", cmd), Span::call_site());
cmd_fns.push(cmd_fn);
let cmd_name = cmd.to_string();
cmd_fns.push(quote!(&#cmd_name, #cmd));
} else {
abort!(t, "expect a list of comma separated commands");
}
}

quote! (
#(#cmd_fns();)*
#(::cmd_lib::register_cmd(#cmd_fns);)*
)
.into()
}
Expand Down
9 changes: 4 additions & 5 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -251,13 +251,12 @@
//! ```
//!
//! ### Macros to register your own commands
//! Declare your function with `#[export_cmd(..)]` attribute, and import it with [`use_custom_cmd!`](https://docs.rs/cmd_lib/latest/cmd_lib/macro.use_custom_cmd.html) macro:
//! Declare your function with the right signature, and register it with [`use_custom_cmd!`](https://docs.rs/cmd_lib/latest/cmd_lib/macro.use_custom_cmd.html) macro:
//!
//! ```
//! # use cmd_lib::*;
//! # use std::io::Write;
//! #[export_cmd(my_cmd)]
//! fn foo(env: &mut CmdEnv) -> CmdResult {
//! fn my_cmd(env: &mut CmdEnv) -> CmdResult {
//! let msg = format!("msg from foo(), args: {:?}", env.args());
//! writeln!(env.stderr(), "{}", msg)?;
//! writeln!(env.stdout(), "bar")
Expand Down Expand Up @@ -363,7 +362,7 @@
//!
pub use cmd_lib_macros::{
cmd_die, export_cmd, main, run_cmd, run_fun, spawn, spawn_with_output, use_custom_cmd,
cmd_die, main, run_cmd, run_fun, spawn, spawn_with_output, use_custom_cmd,
};
/// Return type for [`run_fun!()`] macro.
pub type FunResult = std::io::Result<String>;
Expand All @@ -378,7 +377,7 @@ pub use logger::try_init_default_logger;
pub use main_error::MainError;
pub use main_error::MainResult;
#[doc(hidden)]
pub use process::{export_cmd, AsOsStr, Cmd, CmdString, Cmds, GroupCmds, Redirect};
pub use process::{register_cmd, AsOsStr, Cmd, CmdString, Cmds, GroupCmds, Redirect};
pub use process::{set_debug, set_pipefail, CmdEnv};

mod builtins;
Expand Down
2 changes: 1 addition & 1 deletion src/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ lazy_static! {
}

#[doc(hidden)]
pub fn export_cmd(cmd: &'static str, func: FnFun) {
pub fn register_cmd(cmd: &'static str, func: FnFun) {
CMD_MAP.lock().unwrap().insert(OsString::from(cmd), func);
}

Expand Down
6 changes: 2 additions & 4 deletions tests/test_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,15 +188,13 @@ fn test_proc_env() {
#[test]
fn test_export_cmd() {
use std::io::Write;
#[export_cmd(my_cmd)]
fn foo(env: &mut CmdEnv) -> CmdResult {
fn my_cmd(env: &mut CmdEnv) -> CmdResult {
let msg = format!("msg from foo(), args: {:?}", env.args());
writeln!(env.stderr(), "{}", msg)?;
writeln!(env.stdout(), "bar")
}

#[export_cmd(my_cmd2)]
fn foo2(env: &mut CmdEnv) -> CmdResult {
fn my_cmd2(env: &mut CmdEnv) -> CmdResult {
let msg = format!("msg from foo2(), args: {:?}", env.args());
writeln!(env.stderr(), "{}", msg)?;
writeln!(env.stdout(), "bar2")
Expand Down

0 comments on commit 9c2a420

Please sign in to comment.