feat: add support for generic traits

1 files changed, 46 insertions, 15 deletions

diff --git a/macros/src/expectation.rs b/macros/src/expectation.rs
index 4fc1451..d35ef97 100644
--- a/macros/src/expectation.rs
+++ b/macros/src/expectation.rs
@@ -50,7 +50,8 @@ use crate::util::{create_path, create_unit_type_tuple};
 pub struct Expectation
 {
     ident: Ident,
-    generics: Generics,
+    method_generics: Generics,
+    generic_params: Punctuated<GenericParam, Token![,]>,
     receiver: Option<Receiver>,
     mock: Ident,
     arg_types: Vec<Type>,
@@ -60,12 +61,24 @@ pub struct Expectation
 
 impl Expectation
 {
-    pub fn new(mock: &Ident, item_method: &ImplItemMethod) -> Self
+    pub fn new(
+        mock: &Ident,
+        item_method: &ImplItemMethod,
+        generic_params: Punctuated<GenericParam, Token![,]>,
+    ) -> Self
     {
         let ident = create_expectation_ident(mock, &item_method.sig.ident);
 
-        let phantom_fields =
-            Self::create_phantom_fields(&item_method.sig.generics.params);
+        let phantom_fields = Self::create_phantom_fields(
+            &item_method
+                .sig
+                .generics
+                .params
+                .clone()
+                .into_iter()
+                .chain(generic_params.clone())
+                .collect(),
+        );
 
         let receiver =
             item_method
@@ -91,7 +104,8 @@ impl Expectation
 
         Self {
             ident,
-            generics: item_method.sig.generics.clone(),
+            method_generics: item_method.sig.generics.clone(),
+            generic_params,
             receiver,
             mock: mock.clone(),
             arg_types,
@@ -158,9 +172,17 @@ impl ToTokens for Expectation
 {
     fn to_tokens(&self, tokens: &mut TokenStream)
     {
-        let generic_params = &self.generics.params;
+        let generics = {
+            let mut generics = self.method_generics.clone();
 
-        let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
+            generics.params.extend(self.generic_params.clone());
+
+            generics
+        };
+
+        let generic_params = &generics.params;
+
+        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
 
         let bogus_generics = create_bogus_generics(generic_params);
 
@@ -186,7 +208,7 @@ impl ToTokens for Expectation
             vis: Visibility::new_pub_crate(),
             struct_token: <Token![struct]>::default(),
             ident: self.ident.clone(),
-            generics: self.generics.clone().without_where_clause(),
+            generics: generics.clone().without_where_clause(),
             fields: Fields::Named(FieldsNamed {
                 brace_token: Brace::default(),
                 named: [Field {
@@ -206,13 +228,7 @@ impl ToTokens for Expectation
                     ))),
                 }]
                 .into_iter()
-                .chain(phantom_fields.iter().map(|phantom_field| Field {
-                    attrs: vec![],
-                    vis: Visibility::Inherited,
-                    ident: Some(phantom_field.field.clone()),
-                    colon_token: Some(<Token![:]>::default()),
-                    ty: Type::Path(phantom_field.type_path.clone()),
-                }))
+                .chain(phantom_fields.iter().cloned().map(Field::from))
                 .collect(),
             }),
             semi_token: None,
@@ -285,6 +301,7 @@ pub fn create_expectation_ident(mock: &Ident, method: &Ident) -> Ident
     format_ident!("{mock}Expectation_{method}")
 }
 
+#[derive(Clone)]
 struct PhantomField
 {
     field: Ident,
@@ -303,6 +320,20 @@ impl ToTokens for PhantomField
     }
 }
 
+impl From<PhantomField> for Field
+{
+    fn from(phantom_field: PhantomField) -> Self
+    {
+        Self {
+            attrs: vec![],
+            vis: Visibility::Inherited,
+            ident: Some(phantom_field.field.clone()),
+            colon_token: Some(<Token![:]>::default()),
+            ty: Type::Path(phantom_field.type_path),
+        }
+    }
+}
+
 fn create_phantom_field_ident(ident: &Ident, kind: &PhantomFieldKind) -> Ident
 {
     match kind {