Support #[proptest] attribute for functions with explicit return types.

Author: frozenlib

src/proptest_fn.rs

@@ -1,9 +1,11 @@
+use std::mem::replace;
+
 use crate::syn_utils::{Arg, Args};
-use proc_macro2::TokenStream;
+use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens};
 use syn::{
     parse2, parse_quote, parse_str, spanned::Spanned, token, Block, Expr, Field, FieldMutability,
-    FnArg, Ident, ItemFn, LitStr, Pat, Result, Visibility,
+    FnArg, Ident, ItemFn, LitStr, Pat, Result, ReturnType, Visibility,
 };
 
 pub fn build_proptest(attr: TokenStream, mut item_fn: ItemFn) -> Result<TokenStream> {
@@ -35,11 +37,23 @@ pub fn build_proptest(attr: TokenStream, mut item_fn: ItemFn) -> Result<TokenStr
     if item_fn.sig.asyncness.is_none() {
         attr_args.r#async = None;
     }
+    let output = replace(&mut item_fn.sig.output, ReturnType::Default);
     let block = if let Some(a) = attr_args.r#async {
         item_fn.sig.asyncness = None;
-        a.apply(block)
+        a.apply(block, output)
     } else {
-        quote!(#block)
+        match output {
+            ReturnType::Default => quote!(#block),
+            ReturnType::Type(_, ty) => {
+                let f = Ident::new("__test_body", Span::mixed_site());
+                quote!({
+                    let #f = move || -> #ty {
+                        #block
+                    };
+                    ::std::result::Result::map_err(#f(), ::std::convert::Into::<TestCaseError>::into)?;
+                })
+            }
+        }
     };
     let block = quote! {
         {
@@ -136,28 +150,43 @@ enum Async {
     Expr(Expr),
 }
 impl Async {
-    fn apply(&self, block: &Block) -> TokenStream {
+    fn apply(&self, block: &Block, output: ReturnType) -> TokenStream {
+        let body;
+        let output_type;
+        let ret_expr;
+        match output {
+            ReturnType::Default => {
+                body = quote! {
+                    #block
+                    Ok(())
+                };
+                output_type =
+                    quote!(::core::result::Result<_, ::proptest::test_runner::TestCaseError> );
+                ret_expr = quote! { ret? };
+            }
+            ReturnType::Type(_, ty) => {
+                body = quote! { #block };
+                output_type = quote!(#ty);
+                ret_expr = quote! {
+                    ::std::result::Result::map_err(ret, ::std::convert::Into::<TestCaseError>::into)?
+                };
+            }
+        }
         match self {
             Async::Tokio => {
                 quote! {
-                    let ret: ::core::result::Result<_, ::proptest::test_runner::TestCaseError> =
+                    let ret: #output_type =
                         tokio::runtime::Runtime::new()
                             .unwrap()
-                            .block_on(async move {
-                                #block
-                                Ok(())
-                            });
-                    ret?;
+                            .block_on(async move { #body });
+                    #ret_expr;
                 }
             }
             Async::Expr(expr) => {
                 quote! {
-                    let ret: ::core::result::Result<(), ::proptest::test_runner::TestCaseError> =
-                    (#expr)(async move {
-                            #block
-                            Ok(())
-                        });
-                    ret?;
+                    let ret: #output_type =
+                    (#expr)(async move { #body });
+                    #ret_expr;
                 }
             }
         }

tests/proptest_fn.rs

@@ -116,3 +116,64 @@ async fn async_expr_test_prop_assert() {
 async fn async_expr_test_prop_assert_false() {
     prop_assert!(false);
 }
+
+struct NotImplError;
+
+#[derive(Debug)]
+struct CustomError(#[allow(dead_code)] TestCaseError);
+
+impl std::error::Error for CustomError {}
+impl std::fmt::Display for CustomError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "CustomError")
+    }
+}
+impl From<TestCaseError> for CustomError {
+    fn from(e: TestCaseError) -> Self {
+        CustomError(e)
+    }
+}
+impl From<NotImplError> for CustomError {
+    fn from(_: NotImplError) -> Self {
+        CustomError(TestCaseError::fail("Custom"))
+    }
+}
+
+#[proptest]
+fn custom_error_ok(#[strategy(1..10u8)] x: u8) -> Result<(), CustomError> {
+    assert!(x < 10);
+    not_impl_error_ok()?;
+    Ok(())
+}
+
+#[proptest]
+#[should_panic]
+fn custom_error_err(#[strategy(1..10u8)] x: u8) -> Result<(), CustomError> {
+    assert!(x < 10);
+    not_impl_error_err()?;
+    Ok(())
+}
+
+#[proptest(async = "tokio")]
+async fn custom_error_async_ok(#[strategy(1..10u8)] x: u8) -> Result<(), CustomError> {
+    assert!(x < 10);
+    not_impl_error_ok()?;
+    yield_now().await;
+    Ok(())
+}
+
+#[proptest(async = "tokio")]
+#[should_panic]
+async fn custom_error_async_err(#[strategy(1..10u8)] x: u8) -> Result<(), CustomError> {
+    assert!(x < 10);
+    not_impl_error_err()?;
+    yield_now().await;
+    Ok(())
+}
+
+fn not_impl_error_ok() -> Result<(), NotImplError> {
+    Ok(())
+}
+fn not_impl_error_err() -> Result<(), NotImplError> {
+    Err(NotImplError)
+}