f cleanup macro

Author: devrandom

ext-test-macro/src/lib.rs

@@ -1,7 +1,7 @@
 use proc_macro::TokenStream;
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
-use syn::{parse2, ItemFn, ItemMod};
+use syn::{parse_macro_input, Item};
 
 /// An exposed test.  This is a test that will run locally and also be
 /// made available to other crates that want to run it in their own context.
@@ -40,34 +40,41 @@ use syn::{parse2, ItemFn, ItemMod};
 /// is on.
 #[proc_macro_attribute]
 pub fn xtest(attrs: TokenStream, item: TokenStream) -> TokenStream {
-	let input = syn::parse_macro_input!(item as TokenStream2);
-	let attrs = syn::parse_macro_input!(attrs as TokenStream2);
+	let attrs = parse_macro_input!(attrs as TokenStream2);
+	let input = parse_macro_input!(item as Item);
 
-	let expanded = if is_module_definition(input.clone()) {
-		let cfg = if attrs.is_empty() {
-			quote! { #[cfg(test)] }
-		} else {
-			quote! { #[cfg(any(test, #attrs))] }
-		};
-		quote! {
-			#cfg
-			#input
+	let expanded = match input {
+		Item::Mod(item_mod) => {
+			let cfg = if attrs.is_empty() {
+				quote! { #[cfg(test)] }
+			} else {
+				quote! { #[cfg(any(test, #attrs))] }
+			};
+			quote! {
+				#cfg
+				#item_mod
+			}
 		}
-	} else if is_function_definition(input.clone()) {
-		quote! {
-			#[cfg_attr(test, ::core::prelude::v1::test)]
-			#input
+		Item::Fn(item_fn) => {
+			let cfg_attr = if attrs.is_empty() {
+				quote! { #[cfg(test)] }
+			} else {
+				quote! { #[cfg(any(test, #attrs))] }
+			};
+			quote! {
+				#cfg_attr
+				#item_fn
+			}
+		}
+		_ => {
+			return syn::Error::new_spanned(
+				input,
+				"xtest can only be applied to functions or modules",
+			)
+			.to_compile_error()
+			.into();
 		}
-	} else {
-		panic!("xtest can only be applied to functions or modules");
 	};
-	expanded.into()
-}
-
-fn is_module_definition(input: TokenStream2) -> bool {
-	parse2::<ItemMod>(input).is_ok()
-}
 
-fn is_function_definition(input: TokenStream2) -> bool {
-	parse2::<ItemFn>(input).is_ok()
+	TokenStream::from(expanded)
 }