Make tokio-macros attributes more IDE friendly

tokio-rs · Oct 9, 2021 · 57892f5 · 57892f5
1 parent d9b2dc8
commit 57892f5
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 19 deletions.
diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
@@ -1,6 +1,10 @@
 use proc_macro::TokenStream;
 use proc_macro2::Span;
 use quote::{quote, quote_spanned, ToTokens};
+use syn::parse::Parser;
+
+// syn::AttributeArgs does not implement syn::Parse
+type AttributeArgs = syn::punctuated::Punctuated<syn::NestedMeta, syn::Token![,]>;
 
 #[derive(Clone, Copy, PartialEq)]
 enum RuntimeFlavor {
@@ -186,7 +190,7 @@ fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Erro
 
 fn parse_knobs(
     mut input: syn::ItemFn,
-    args: syn::AttributeArgs,
+    args: AttributeArgs,
     is_test: bool,
     rt_multi_thread: bool,
 ) -> Result<TokenStream, syn::Error> {
@@ -357,32 +361,36 @@ fn parse_knobs(
 }
 
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
-pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
-    let input = syn::parse_macro_input!(item as syn::ItemFn);
-    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
+pub(crate) fn main(
+    args: TokenStream,
+    item: TokenStream,
+    rt_multi_thread: bool,
+) -> syn::Result<TokenStream> {
+    let input: syn::ItemFn = syn::parse(item)?;
+    let args = AttributeArgs::parse_terminated.parse(args)?;
 
     if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
         let msg = "the main function cannot accept arguments";
-        return syn::Error::new_spanned(&input.sig.ident, msg)
-            .to_compile_error()
-            .into();
+        return Err(syn::Error::new_spanned(&input.sig.ident, msg));
     }
 
-    parse_knobs(input, args, false, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
+    parse_knobs(input, args, false, rt_multi_thread)
 }
 
-pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
-    let input = syn::parse_macro_input!(item as syn::ItemFn);
-    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
+pub(crate) fn test(
+    args: TokenStream,
+    item: TokenStream,
+    rt_multi_thread: bool,
+) -> syn::Result<TokenStream> {
+    let input: syn::ItemFn = syn::parse(item)?;
+    let args = AttributeArgs::parse_terminated.parse(args)?;
 
     for attr in &input.attrs {
         if attr.path.is_ident("test") {
             let msg = "second test attribute is supplied";
-            return syn::Error::new_spanned(&attr, msg)
-                .to_compile_error()
-                .into();
+            return Err(syn::Error::new_spanned(&attr, msg));
         }
     }
 
-    parse_knobs(input, args, true, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
+    parse_knobs(input, args, true, rt_multi_thread)
 }
diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
@@ -185,7 +185,11 @@ use proc_macro::TokenStream;
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
-    entry::main(args, item, true)
+    entry::main(args, item.clone(), true).unwrap_or_else(|err| {
+        let mut output: TokenStream = err.to_compile_error().into();
+        output.extend(item);
+        output
+    })
 }
 
 /// Marks async function to be executed by selected runtime. This macro helps set up a `Runtime`
@@ -231,7 +235,11 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
-    entry::main(args, item, false)
+    entry::main(args, item.clone(), false).unwrap_or_else(|err| {
+        let mut output: TokenStream = err.to_compile_error().into();
+        output.extend(item);
+        output
+    })
 }
 
 /// Marks async function to be executed by runtime, suitable to test environment
@@ -277,7 +285,11 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
 /// available as `tokio` in the module where this macro is expanded.
 #[proc_macro_attribute]
 pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
-    entry::test(args, item, true)
+    entry::test(args, item.clone(), true).unwrap_or_else(|err| {
+        let mut output: TokenStream = err.to_compile_error().into();
+        output.extend(item);
+        output
+    })
 }
 
 /// Marks async function to be executed by runtime, suitable to test environment
@@ -299,7 +311,11 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
 /// available as `tokio` in the module where this macro is expanded.
 #[proc_macro_attribute]
 pub fn test_rt(args: TokenStream, item: TokenStream) -> TokenStream {
-    entry::test(args, item, false)
+    entry::test(args, item.clone(), false).unwrap_or_else(|err| {
+        let mut output: TokenStream = err.to_compile_error().into();
+        output.extend(item);
+        output
+    })
 }
 
 /// Always fails with the error message below.