support result

fast · Jul 22, 2023 · 42fd3bd · 42fd3bd
1 parent cc5fef9
commit 42fd3bd
Show file tree

Hide file tree

Showing 9 changed files with 206 additions and 47 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "logcall"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 authors = ["andylokandy <[email protected]>"]
 description = "An attribute macro that logs the return value from function call."

diff --git a/README.md b/README.md
@@ -14,11 +14,25 @@ fn foo(a: usize) -> usize {
     a + 1
 }
 
+#[logcall(err = "error")]
+fn bar(a: usize) -> Result<usize, usize> {
+    Err(a + 1)
+}
+
+#[logcall(ok = "info", err = "error")]
+fn baz(a: usize) -> Result<usize, usize> {
+    Ok(a + 1)
+}
+
 fn main() {
     env_logger::builder().filter_level(log::LevelFilter::Info).init();
     foo(1);
+    bar(1).ok();
+    baz(1).ok();
 }
 
 // prints:
-// [2023-07-21T12:57:43Z INFO  main] foo() => 2
+// [2023-07-22T06:55:10Z INFO  main] foo() => 2
+// [2023-07-22T06:55:10Z ERROR main] bar() => Err(2)
+// [2023-07-22T06:55:10Z INFO  main] baz() => Ok(2)
 ```
diff --git a/examples/main.rs b/examples/main.rs
@@ -5,9 +5,19 @@ fn foo(a: usize) -> usize {
     a + 1
 }
 
+#[logcall(err = "error")]
+fn bar(a: usize) -> Result<usize, usize> {
+    Err(a + 1)
+}
+
+#[logcall(ok = "info", err = "error")]
+fn baz(a: usize) -> Result<usize, usize> {
+    Ok(a + 1)
+}
+
 fn main() {
-    env_logger::builder()
-        .filter_level(log::LevelFilter::Info)
-        .init();
+    env_logger::builder().filter_level(log::LevelFilter::Info).init();
     foo(1);
+    bar(1).ok();
+    baz(1).ok();
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -17,23 +17,60 @@ use syn::spanned::Spanned;
 use syn::Ident;
 use syn::*;
 
-struct Args {
-    level: String,
+enum Args {
+    Simple {
+        level: String,
+    },
+    Result {
+        ok_level: Option<String>,
+        err_level: Option<String>,
+    },
 }
 
 impl Args {
     fn parse(input: AttributeArgs) -> Args {
-        if input.len() > 1 {
-            abort_call_site!("too many arguments");
+        match input.as_slice() {
+            [NestedMeta::Lit(Lit::Str(s))] => Args::Simple { level: s.value() },
+            [NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+                path,
+                lit: Lit::Str(s),
+                ..
+            }))] if path.is_ident("ok") => Args::Result {
+                ok_level: Some(s.value()),
+                err_level: None,
+            },
+            [NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+                path,
+                lit: Lit::Str(s),
+                ..
+            }))] if path.is_ident("err") => Args::Result {
+                ok_level: None,
+                err_level: Some(s.value()),
+            },
+            [NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+                path,
+                lit: Lit::Str(s),
+                ..
+            })), NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+                path: path2,
+                lit: Lit::Str(s2),
+                ..
+            }))]
+            | [NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+                path: path2,
+                lit: Lit::Str(s2),
+                ..
+            })), NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+                path,
+                lit: Lit::Str(s),
+                ..
+            }))] if path.is_ident("ok") && path2.is_ident("err") => Args::Result {
+                ok_level: Some(s.value()),
+                err_level: Some(s2.value()),
+            },
+            [] => abort_call_site!("missing arguments"),
+            _ => abort_call_site!("invalid arguments"),
         }
-
-        let level = if let Some(NestedMeta::Lit(Lit::Str(s))) = input.first() {
-            s.value()
-        } else {
-            abort_call_site!("invalid argument");
-        };
-
-        Args { level }
     }
 }
 
@@ -127,32 +164,94 @@ fn gen_block(
     fn_name: &str,
     args: Args,
 ) -> proc_macro2::TokenStream {
-    // Generate the instrumented function body.
-    // If the function is an `async fn`, this will wrap it in an async block.
-    if async_context {
-        let log = gen_log(&args.level, fn_name, "__ret_value");
-        let block = quote_spanned!(block.span()=>
-            async move {
-                let __ret_value = #block;
-                #log;
-                __ret_value
+    match args {
+        Args::Simple { level } =>  {
+            // Generate the instrumented function body.
+            // If the function is an `async fn`, this will wrap it in an async block.
+            if async_context {
+                let log = gen_log(&level, fn_name, "__ret_value");
+                let block = quote_spanned!(block.span()=>
+                    async move {
+                        let __ret_value = #block;
+                        #log;
+                        __ret_value
+                    }
+                );
+
+                if async_keyword {
+                    quote_spanned!(block.span()=>
+                        #block.await
+                    )
+                } else {
+                    block
+                }
+            } else {
+                let log = gen_log(&level, fn_name, "__ret_value");
+                quote_spanned!(block.span()=>
+                    let __ret_value = #block;
+                    #log;
+                    __ret_value
+                )
+            }
+        }
+        Args::Result { ok_level, err_level } => {
+            let ok_arm = if let Some(ok_level) = ok_level {
+                let log_ok = gen_log(&ok_level, fn_name, "__ret_value");
+                quote_spanned!(block.span()=>
+                    __ret_value@Ok(_) => {
+                        #log_ok;
+                        __ret_value
+                    }
+                )
+            } else {
+                quote_spanned!(block.span()=>
+                    Ok(__ret_value) => Ok(__ret_value),
+                )
+            };
+            let err_arm = if let Some(err_level) = err_level {
+                let log_err = gen_log(&err_level, fn_name, "__ret_value");
+                quote_spanned!(block.span()=>
+                    __ret_value@Err(_) => {
+                        #log_err;
+                        __ret_value
+                    }
+                )
+            } else {
+                quote_spanned!(block.span()=>
+                    Err(__ret_value) => Err(__ret_value),
+                )
+            };
+
+            // Generate the instrumented function body.
+            // If the function is an `async fn`, this will wrap it in an async block.
+            if async_context {
+                let block = quote_spanned!(block.span()=>
+                    async move {
+                        let __ret_value = #block;
+                        match __ret_value {
+                            #ok_arm
+                            #err_arm
+                        }
+                    }
+                );
+
+                if async_keyword {
+                    quote_spanned!(block.span()=>
+                        #block.await
+                    )
+                } else {
+                    block
+                }
+            } else {
+                quote_spanned!(block.span()=>
+                    let __ret_value = #block;
+                    match __ret_value {
+                        #ok_arm
+                        #err_arm
+                    }
+                )
             }
-        );
-
-        if async_keyword {
-            quote_spanned!(block.span()=>
-                #block.await
-            )
-        } else {
-            block
         }
-    } else {
-        let log = gen_log(&args.level, fn_name, "__ret_value");
-        quote_spanned!(block.span()=>
-            let __ret_value = #block;
-            #log;
-            __ret_value
-        )
     }
 }
 

diff --git a/tests/ui/err/has-no-argument.stderr b/tests/ui/err/has-no-argument.stderr
@@ -1,4 +1,4 @@
-error: invalid argument
+error: missing arguments
  --> tests/ui/err/has-no-argument.rs:1:1
   |
 1 | #[logcall::logcall]

diff --git a/tests/ui/err/has-too-many-arguments.stderr b/tests/ui/err/has-too-many-arguments.stderr
@@ -1,4 +1,4 @@
-error: too many arguments
+error: invalid arguments
  --> tests/ui/err/has-too-many-arguments.rs:1:1
   |
 1 | #[logcall::logcall("info", "error")]

diff --git a/tests/ui/ok/async-in-trait.rs b/tests/ui/ok/async-in-trait.rs
@@ -2,16 +2,16 @@
 #![allow(unused_mut)]
 
 trait MyTrait {
-    async fn work(&self) -> usize;
+    async fn work(&self) -> Result<usize, usize>;
 }
 
 struct MyStruct;
 
 impl MyTrait for MyStruct {
     #[logcall::logcall("debug")]
-    #[logcall::logcall("debug")]
-    async fn work(&self) -> usize {
-        1
+    #[logcall::logcall(ok = "debug", err = "error")]
+    async fn work(&self) -> Result<usize, usize> {
+        Ok(1)
     }
 }
 

diff --git a/tests/ui/ok/async.rs b/tests/ui/ok/async.rs
@@ -3,7 +3,25 @@ async fn f(a: u32) -> u32 {
     a
 }
 
+#[logcall::logcall(ok = "info")]
+async fn g(a: u32) -> Result<u32, u32> {
+    Ok(a)
+}
+
+#[logcall::logcall(err = "info")]
+async fn h(a: u32) -> Result<u32, u32> {
+    Ok(a)
+}
+
+#[logcall::logcall(ok = "info", err = "info")]
+async fn i(a: u32) -> Result<u32, u32> {
+    Ok(a)
+}
+
 #[tokio::main]
 async fn main() {
     f(1).await;
+    g(1).await.ok();
+    h(1).await.ok();
+    i(1).await.ok();
 }
diff --git a/tests/ui/ok/sync.rs b/tests/ui/ok/sync.rs
@@ -3,6 +3,24 @@ fn f(a: u32) -> u32 {
     a
 }
 
+#[logcall::logcall(ok = "info")]
+fn g(a: u32) -> Result<u32, u32> {
+    Ok(a)
+}
+
+#[logcall::logcall(err = "info")]
+fn h(a: u32) -> Result<u32, u32> {
+    Ok(a)
+}
+
+#[logcall::logcall(ok = "info", err = "info")]
+fn i(a: u32) -> Result<u32, u32> {
+    Ok(a)
+}
+
 fn main() {
     f(1);
+    g(1).ok();
+    h(1).ok();
+    i(1).ok();
 }