[proptest-macro] add the ability to specify custom strategies

.github/workflows/rust.yml

@@ -46,6 +46,8 @@ jobs:
         run: cd proptest && cargo build --verbose
       - name: Run tests
         run: cd proptest && cargo test --verbose
+      - name: Run macro tests
+        run: cd proptest-macro && cargo test --verbose
       - name: Build coverage no-default-features
         if: ${{ matrix.build == 'stable' }}
         env:

proptest-macro/src/lib.rs

@@ -11,6 +11,7 @@ mod property_test;
 /// Using the `property_test` macro:
 ///
 /// ```
+/// # use proptest_macro::property_test;
 /// #[property_test]
 /// fn foo(x: i32) {
 ///     assert_eq!(x, x);
@@ -19,7 +20,7 @@ mod property_test;
 ///
 /// is roughly equivalent to:
 ///
-/// ```
+/// ```ignore
 /// proptest! {
 ///     #[test]
 ///     fn foo(x in any::<i32>()) {
@@ -35,19 +36,36 @@ mod property_test;
 /// of setting up the test harness and generating input values, allowing the user to focus
 /// on writing the test logic.
 ///
-/// # Attributes
+/// ## Attributes
 ///
 /// The `property_test` macro can take an optional `config` attribute, which allows you to
 /// customize the configuration of the `proptest` runner.
 ///
 /// E.g. running 100 cases:
 ///
-/// ```rust
+/// ```rust,ignore
 /// #[property_test(config = "ProptestConfig { cases: 100, .. ProptestConfig::default() }")]
 /// fn foo(x: i32) {
 ///     assert_eq!(x, x);
 /// }
 /// ```
+///
+/// ## Custom strategies
+///
+/// By default, [`property_test`] will use the `Arbitrary` impl for parameters. However, you can
+/// provide a custom `Strategy` with `#[strategy = <expr>]` on an argument:
+///
+/// ```
+/// # use proptest_macro::property_test;
+/// #[property_test]
+/// fn foo(#[strategy = "[0-9]*"] s: String) {
+///     for c in s.chars() {
+///         assert!(c.is_numeric());
+///     }
+/// }
+/// ```
+/// Multiple `#[strategy = <expr>]` attributes on an argument are not allowed.
+///
 #[proc_macro_attribute]
 pub fn property_test(attr: TokenStream, item: TokenStream) -> TokenStream {
     property_test::property_test(attr.into(), item.into()).into()

proptest-macro/src/property_test/codegen/arbitrary.rs

@@ -0,0 +1,112 @@
+use super::*;
+use quote::quote_spanned;
+
+/// Generate the arbitrary impl for the struct
+pub(super) fn gen_arbitrary_impl(
+    fn_name: &Ident,
+    args: &[Argument],
+) -> TokenStream {
+    if args.iter().all(|arg| arg.strategy.is_none()) {
+        no_custom_strategies(fn_name, args)
+    } else {
+        custom_strategies(fn_name, args)
+    }
+}
+
+// we can avoid boxing strategies if there are no custom strategies, since we have types written
+// out in function args
+//
+// If there are custom strategies, we can't write the type, because we're only provided the
+// expression for the strategy (e.g. `#[strategy = my_custom_strategy()]` doesn't tell us the
+// return type of `my_custom_strategy`). In these cases, we just use `BoxedStrategy<Self>`
+fn no_custom_strategies(fn_name: &Ident, args: &[Argument]) -> TokenStream {
+    let arg_types = args.iter().map(|arg| {
+        let ty = &arg.pat_ty.ty;
+        quote!(#ty,)
+    });
+
+    let arg_types = quote! { #(#arg_types)* };
+
+    let arg_names = args.iter().enumerate().map(|(index, arg)| {
+        let name = nth_field_name(arg.pat_ty.span(), index);
+        quote!(#name,)
+    });
+
+    let arg_names = quote! { #(#arg_names)* };
+
+    let strategy_type = quote! {
+        ::proptest::strategy::Map<::proptest::arbitrary::StrategyFor<(#arg_types)>, fn((#arg_types)) -> Self>
+    };
+
+    let strategy_expr = quote! {
+        use ::proptest::strategy::Strategy;
+        ::proptest::prelude::any::<(#arg_types)>().prop_map(|(#arg_names)| Self { #arg_names })
+    };
+
+    arbitrary_shared(fn_name, strategy_type, strategy_expr)
+}
+
+// if we have `fn foo(#[strategy = x] a: i32, b: i32) {}`, we want to generate something like this:
+// ```ignore
+// impl Arbitrary for FooArgs {
+//   type Parameters = ();
+//   type Strategy = BoxedStrategy<Self>;
+//
+//   fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
+//      (x, any::<i32>()).prop_map(|(a, b)| Self { a, b }).boxed()
+//   }
+// }
+// ```
+fn custom_strategies(fn_name: &Ident, args: &[Argument]) -> TokenStream {
+    let arg_strategies = args.iter().map(|arg| {
+        arg.strategy
+            .as_ref()
+            .map(|s| s.to_token_stream())
+            .unwrap_or_else(|| {
+                let ty = &arg.pat_ty.ty;
+                quote_spanned! {
+                    ty.span() => ::proptest::prelude::any::<#ty>()
+                }
+            })
+    });
+
+    let arg_names: TokenStream = args
+        .iter()
+        .enumerate()
+        .map(|(index, arg)| {
+            let name = nth_field_name(arg.pat_ty.span(), index);
+            quote!(#name,)
+        })
+        .collect();
+    let arg_names = &arg_names;
+
+    let strategy_expr = quote! {
+        use ::proptest::strategy::Strategy;
+        (#(#arg_strategies),*).prop_map(|(#arg_names)| Self { #arg_names }).boxed()
+    };
+
+    let strategy_type = quote! {
+        ::proptest::strategy::BoxedStrategy<Self>
+    };
+    arbitrary_shared(fn_name, strategy_type, strategy_expr)
+}
+
+/// shared code between both boxed and unboxed paths
+fn arbitrary_shared(
+    fn_name: &Ident,
+    strategy_type: TokenStream,
+    strategy_expr: TokenStream,
+) -> TokenStream {
+    let struct_name = struct_name(fn_name);
+
+    quote! {
+        impl ::proptest::prelude::Arbitrary for #struct_name {
+            type Parameters = ();
+            type Strategy = #strategy_type;
+
+            fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
+                #strategy_expr
+            }
+        }
+    }
+}

proptest-macro/src/property_test/codegen/mod.rs

@@ -1,9 +1,13 @@
 use proc_macro2::TokenStream;
 use quote::{quote, ToTokens};
-use syn::{parse_str, spanned::Spanned, Attribute, Ident, ItemFn, PatType};
+use syn::{parse_quote, spanned::Spanned, Attribute, Ident, ItemFn};
 
-use super::{options::Options, utils::strip_args};
+use super::{
+    options::Options,
+    utils::{strip_args, Argument},
+};
 
+mod arbitrary;
 mod test_body;
 
 /// Generate the modified test function
@@ -20,17 +24,18 @@ pub(super) fn generate(item_fn: ItemFn, options: Options) -> TokenStream {
     let (mut argless_fn, args) = strip_args(item_fn);
 
     let struct_tokens = generate_struct(&argless_fn.sig.ident, &args);
-    let arb_tokens = generate_arbitrary_impl(&argless_fn.sig.ident, &args);
+    let arb_tokens =
+        arbitrary::gen_arbitrary_impl(&argless_fn.sig.ident, &args);
 
-    let struct_and_tokens = quote! {
+    let struct_and_arb = quote! {
         #struct_tokens
         #arb_tokens
     };
 
     let new_body = test_body::body(
         *argless_fn.block,
         &args,
-        struct_and_tokens,
+        struct_and_arb,
         &argless_fn.sig.ident,
         &argless_fn.sig.output,
         &options,
@@ -43,12 +48,12 @@ pub(super) fn generate(item_fn: ItemFn, options: Options) -> TokenStream {
 }
 
 /// Generate the inner struct that represents the arguments of the function
-fn generate_struct(fn_name: &Ident, args: &[PatType]) -> TokenStream {
+fn generate_struct(fn_name: &Ident, args: &[Argument]) -> TokenStream {
     let struct_name = struct_name(fn_name);
 
     let fields = args.iter().enumerate().map(|(index, arg)| {
-        let field_name = nth_field_name(&arg.pat, index);
-        let ty = &arg.ty;
+        let field_name = nth_field_name(&arg.pat_ty.pat, index);
+        let ty = &arg.pat_ty.ty;
 
         quote! { #field_name: #ty, }
     });
@@ -61,37 +66,6 @@ fn generate_struct(fn_name: &Ident, args: &[PatType]) -> TokenStream {
     }
 }
 
-/// Generate the arbitrary impl for the struct
-fn generate_arbitrary_impl(fn_name: &Ident, args: &[PatType]) -> TokenStream {
-    let struct_name = struct_name(fn_name);
-
-    let arg_types = args.iter().map(|arg| {
-        let ty = &arg.ty;
-        quote!(#ty,)
-    });
-
-    let arg_types = quote! { #(#arg_types)* };
-
-    let arg_names = args.iter().enumerate().map(|(index, ty)| {
-        let name = nth_field_name(ty.span(), index);
-        quote!(#name,)
-    });
-
-    let arg_names = quote! { #(#arg_names)* };
-
-    quote! {
-        impl ::proptest::prelude::Arbitrary for #struct_name {
-            type Parameters = ();
-            type Strategy = ::proptest::strategy::Map<::proptest::arbitrary::StrategyFor<(#arg_types)>, fn((#arg_types)) -> Self>;
-
-            fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
-                use ::proptest::strategy::Strategy;
-                ::proptest::prelude::any::<(#arg_types)>().prop_map(|(#arg_names)| Self { #arg_names })
-            }
-        }
-    }
-}
-
 /// Convert the name of a function to the name of a struct representing its args
 ///
 /// E.g. `some_function` -> `SomeFunctionArgs`
@@ -111,19 +85,14 @@ fn nth_field_name(span: impl Spanned, index: usize) -> Ident {
     Ident::new(&format!("field{index}"), span.span())
 }
 
-/// I couldn't find a better way to get just the `#[test]` attribute since [`syn::Attribute`]
-/// doesn't implement `Parse`
 fn test_attr() -> Attribute {
-    let mut f: ItemFn = parse_str("#[test] fn foo() {}").unwrap();
-    f.attrs.pop().unwrap()
+    parse_quote! { #[test] }
 }
 
 #[cfg(test)]
 mod tests {
-    use quote::ToTokens;
-    use syn::{parse2, parse_str, ItemStruct};
-
     use super::*;
+    use syn::{parse2, parse_quote, parse_str, ItemStruct};
 
     /// Simple helper that parses a function, and validates that the struct name and fields are
     /// correct
@@ -180,31 +149,18 @@ mod tests {
 
     #[test]
     fn generates_arbitrary_impl() {
-        let f: ItemFn = parse_str("fn foo(x: i32, y: u8) {}").unwrap();
+        let f: ItemFn = parse_quote! { fn foo(x: i32, y: u8) {} };
         let (f, args) = strip_args(f);
-        let arb = generate_arbitrary_impl(&f.sig.ident, &args);
-
-        let expected = quote! {
-            impl ::proptest::prelude::Arbitrary for FooArgs {
-                type Parameters = ();
-                type Strategy = ::proptest::strategy::Map<::proptest::arbitrary::StrategyFor<(i32, u8,)>, fn((i32, u8,)) -> Self>;
-
-                fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
-                    use ::proptest::strategy::Strategy;
-
-                    ::proptest::prelude::any::<(i32, u8,)>().prop_map(|(field0, field1,)| Self { field0, field1, })
-                }
-
-            }
-        };
+        let arb = arbitrary::gen_arbitrary_impl(&f.sig.ident, &args);
 
-        assert_eq!(arb.to_string(), expected.to_string());
+        insta::assert_snapshot!(arb.to_string());
     }
 }
 
 #[cfg(test)]
 mod snapshot_tests {
     use super::*;
+    use syn::parse_str;
 
     macro_rules! snapshot_test {
         ($name:ident) => {

proptest-macro/src/property_test/codegen/snapshots/proptest_macro__property_test__codegen__tests__generates_arbitrary_impl.snap

@@ -0,0 +1,5 @@
+---
+source: proptest-macro/src/property_test/codegen/mod.rs
+expression: arb.to_string()
+---
+impl :: proptest :: prelude :: Arbitrary for FooArgs { type Parameters = () ; type Strategy = :: proptest :: strategy :: Map < :: proptest :: arbitrary :: StrategyFor < (i32 , u8 ,) > , fn ((i32 , u8 ,)) -> Self > ; fn arbitrary_with (() : Self :: Parameters) -> Self :: Strategy { use :: proptest :: strategy :: Strategy ; :: proptest :: prelude :: any :: < (i32 , u8 ,) > () . prop_map (| (field0 , field1 ,) | Self { field0 , field1 , }) } }

proptest-macro/src/property_test/codegen/test_body.rs

@@ -1,32 +1,31 @@
 use proc_macro2::TokenStream;
 use quote::{quote, ToTokens};
 use syn::{
-    parse2, spanned::Spanned, Block, Expr, Ident, PatType, ReturnType, Type,
-    TypeTuple,
+    parse2, spanned::Spanned, Block, Expr, Ident, ReturnType, Type, TypeTuple,
 };
 
-use crate::property_test::options::Options;
+use crate::property_test::{options::Options, utils::Argument};
 
 use super::{nth_field_name, struct_name};
 
 /// Generate the new test body by putting the struct and arbitrary impl at the start, then adding
 /// the usual glue that `proptest!` adds
 pub(super) fn body(
     block: Block,
-    args: &[PatType],
+    args: &[Argument],
     struct_and_impl: TokenStream,
     fn_name: &Ident,
     ret_ty: &ReturnType,
-    options: &Options, 
+    options: &Options,
 ) -> Block {
     let struct_name = struct_name(fn_name);
 
     let errors = &options.errors;
 
     // convert each arg to `field0: x`
     let struct_fields = args.iter().enumerate().map(|(index, arg)| {
-        let pat = &arg.pat;
-        let field_name = nth_field_name(arg.pat.span(), index);
+        let pat = &arg.pat_ty.pat;
+        let field_name = nth_field_name(arg.pat_ty.pat.span(), index);
         quote!(#field_name: #pat,)
     });
 
@@ -66,6 +65,7 @@ pub(super) fn body(
     } );
 
     // unwrap here is fine because the double braces create a block
+    // std::fs::write("/tmp/pt-debug", tokens.to_string());
     parse2(tokens).unwrap()
 }
 
@@ -104,4 +104,3 @@ fn make_config(config: Option<&Expr>) -> TokenStream {
         };
     }
 }
-

proptest-macro/src/property_test/mod.rs

@@ -22,10 +22,10 @@ macro_rules! parse {
 }
 
 pub fn property_test(attr: TokenStream, item: TokenStream) -> TokenStream {
-    let item_fn = parse!(item);
+    let mut item_fn = parse!(item);
     let options = parse!(attr);
 
-    if let Err(compile_error) = validate(&item_fn) {
+    if let Err(compile_error) = validate(&mut item_fn) {
         return compile_error;
     }