From c01ce5bc70d3636c5b1d1c689a29fd10412b0e1f Mon Sep 17 00:00:00 2001
From: Tao Guo <tao.guo.code@gmail.com>
Date: Tue, 31 Oct 2023 03:58:32 +0000
Subject: [PATCH] Add `cmd_lib::main` attribute macro to log main() error by
 default

---
 README.md                 |  6 ++----
 examples/dd_test.rs       |  1 +
 examples/pipes.rs         |  1 +
 examples/rust_cookbook.rs |  2 ++
 examples/tetris.rs        |  1 +
 macros/src/lib.rs         | 39 +++++++++++++++++++++++++++++++++++++--
 src/lib.rs                |  8 +++-----
 7 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 04fc197..bf54833 100644
--- a/README.md
+++ b/README.md
@@ -91,16 +91,14 @@ run_cmd!(du -ah $dir | sort -hr | head -n 10)?;
 // if any command fails, just return Err(...)
 let file = "/tmp/f";
 let keyword = "rust";
-if run_cmd! {
+run_cmd! {
     cat ${file} | grep ${keyword};
     echo "bad cmd" >&2;
     ignore ls /nofile;
     date;
     ls oops;
     cat oops;
-}.is_err() {
-    // your error handling code
-}
+}?;
 ```
 
 - run_fun! --> FunResult
diff --git a/examples/dd_test.rs b/examples/dd_test.rs
index 2e59ba1..1ad259f 100644
--- a/examples/dd_test.rs
+++ b/examples/dd_test.rs
@@ -33,6 +33,7 @@ struct Opt {
     file: String,
 }
 
+#[cmd_lib::main]
 fn main() -> MainResult {
     let Opt {
         block_size,
diff --git a/examples/pipes.rs b/examples/pipes.rs
index 0e7af83..0f44de7 100644
--- a/examples/pipes.rs
+++ b/examples/pipes.rs
@@ -338,6 +338,7 @@ fn rand() -> i32 {
     run_fun!(bash -c r"echo $RANDOM").unwrap().parse().unwrap()
 }
 
+#[cmd_lib::main]
 fn main() -> MainResult {
     // simple pre-check of TERM, tput's error message should be enough
     let term = std::env::var("TERM").unwrap();
diff --git a/examples/rust_cookbook.rs b/examples/rust_cookbook.rs
index 2aac829..b7f0fe6 100644
--- a/examples/rust_cookbook.rs
+++ b/examples/rust_cookbook.rs
@@ -4,6 +4,8 @@
 //
 use cmd_lib::*;
 use std::io::{BufRead, BufReader};
+
+#[cmd_lib::main]
 fn main() -> MainResult {
     cmd_lib::set_pipefail(false); // do not fail due to pipe errors
 
diff --git a/examples/tetris.rs b/examples/tetris.rs
index f1293ed..7266d9c 100644
--- a/examples/tetris.rs
+++ b/examples/tetris.rs
@@ -556,6 +556,7 @@ fn cmd_exit() {
     std::process::exit(0);
 }
 
+#[cmd_lib::main]
 fn main() -> MainResult {
     #[rustfmt::skip]
     let old_cfg = run_fun!(stty -g)?; // let's save terminal state ...
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index e04727e..d59dca5 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -2,6 +2,42 @@ use proc_macro2::{Span, TokenStream, TokenTree};
 use proc_macro_error::{abort, proc_macro_error};
 use quote::{quote, ToTokens};
 
+/// Mark main function to log error result by default
+///
+/// ```
+/// # use cmd_lib::*;
+///
+/// #[cmd_lib::main]
+/// fn main() -> MainResult {
+///     run_cmd!(bad_cmd)?;
+///     Ok(())
+/// }
+/// // output:
+/// // [ERROR] FATAL: Running ["bad_cmd"] failed: No such file or directory (os error 2)
+/// ```
+#[proc_macro_attribute]
+pub fn main(
+    _args: proc_macro::TokenStream,
+    item: proc_macro::TokenStream,
+) -> proc_macro::TokenStream {
+    let orig_function: syn::ItemFn = syn::parse2(item.into()).unwrap();
+    let orig_main_block = orig_function.block;
+
+    quote! (
+        fn main() {
+            __cmd_main().unwrap_or_else(|err| {
+                ::cmd_lib::error!("FATAL: {err:?}");
+                std::process::exit(1);
+            });
+        }
+
+        fn __cmd_main() -> MainResult {
+            #orig_main_block
+        }
+    )
+    .into()
+}
+
 /// Export the function as an command to be run by `run_cmd!` or `run_fun!`
 ///
 /// ```
@@ -211,8 +247,7 @@ pub fn cmd_die(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let msg = parse_msg(input.into());
     quote!({
         use ::cmd_lib::error;
-        use ::cmd_lib::AsOsStr;
-        error!("FATAL: {}", #msg);
+        ::cmd_lib::error!("FATAL: {}", #msg);
         std::process::exit(1)
     })
     .into()
diff --git a/src/lib.rs b/src/lib.rs
index 8b2043e..b677a4d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -99,16 +99,14 @@
 //! // if any command fails, just return Err(...)
 //! let file = "/tmp/f";
 //! let keyword = "rust";
-//! if run_cmd! {
+//! run_cmd! {
 //!     cat ${file} | grep ${keyword};
 //!     echo "bad cmd" >&2;
 //!     ignore ls /nofile;
 //!     date;
 //!     ls oops;
 //!     cat oops;
-//! }.is_err() {
-//!     // your error handling code
-//! }
+//! }?;
 //! # Ok::<(), std::io::Error>(())
 //! ```
 //!
@@ -356,7 +354,7 @@
 //!
 
 pub use cmd_lib_macros::{
-    cmd_die, export_cmd, run_cmd, run_fun, spawn, spawn_with_output, use_custom_cmd,
+    cmd_die, export_cmd, 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>;