Allow adding generics to impl and trait name (#90)

* Add with_impl_generics

* Add tests

* fix clippy

* fix format

Author: branchseer

src/generate/generator.rs

@@ -309,4 +309,63 @@ mod test {
             .collect::<String>()
         );
     }
+
+    #[test]
+    fn impl_for_with_trait_generics() {
+        let mut generator = Generator::new(
+            Ident::new("StructOrEnum", Span::call_site()),
+            Generics::try_take(&mut token_stream("<'a>")).unwrap(),
+            None,
+        );
+        let _ = generator.impl_for("Foo").with_trait_generics(["&'a str"]);
+        let output = generator.finish().unwrap();
+        assert_eq!(
+            output
+                .into_iter()
+                .map(|v| v.to_string())
+                .collect::<String>(),
+            token_stream("impl<'a> Foo<&'a str> for StructOrEnum<'a> { }")
+                .map(|v| v.to_string())
+                .collect::<String>(),
+        );
+    }
+
+    #[test]
+    fn impl_for_with_impl_generics() {
+        //with simple generics
+        let mut generator = Generator::new(
+            Ident::new("StructOrEnum", Span::call_site()),
+            Generics::try_take(&mut token_stream("<T1, T2>")).unwrap(),
+            None,
+        );
+        let _ = generator.impl_for("Foo").with_impl_generics(["Bar"]);
+
+        let output = generator.finish().unwrap();
+        assert_eq!(
+            output
+                .into_iter()
+                .map(|v| v.to_string())
+                .collect::<String>(),
+            token_stream("impl<T1, T2, Bar> Foo for StructOrEnum<T1, T2> { }")
+                .map(|v| v.to_string())
+                .collect::<String>()
+        );
+        // with lifetimes
+        let mut generator = Generator::new(
+            Ident::new("StructOrEnum", Span::call_site()),
+            Generics::try_take(&mut token_stream("<'alpha, 'beta>")).unwrap(),
+            None,
+        );
+        let _ = generator.impl_for("Foo").with_impl_generics(["Bar"]);
+        let output = generator.finish().unwrap();
+        assert_eq!(
+            output
+                .into_iter()
+                .map(|v| v.to_string())
+                .collect::<String>(),
+            token_stream("impl<'alpha, 'beta, Bar> Foo for StructOrEnum<'alpha, 'beta> { }")
+                .map(|v| v.to_string())
+                .collect::<String>()
+        );
+    }
 }

src/generate/impl_for.rs

@@ -13,7 +13,8 @@ pub struct ImplFor<'a, P: Parent> {
     type_name: StringOrIdent,
     trait_name: Option<StringOrIdent>,
     lifetimes: Option<Vec<String>>,
-    generics: Option<Vec<String>>,
+    trait_generics: Option<Vec<String>>,
+    impl_generics: Vec<String>,
     consts: Vec<StreamBuilder>,
     custom_generic_constraints: Option<GenericConstraints>,
     impl_types: Vec<StreamBuilder>,
@@ -33,7 +34,8 @@ impl<'a, P: Parent> ImplFor<'a, P> {
             trait_name,
             type_name,
             lifetimes: None,
-            generics: None,
+            trait_generics: None,
+            impl_generics: vec![],
             consts: Vec::new(),
             custom_generic_constraints: None,
             impl_types: Vec::new(),
@@ -100,7 +102,30 @@ impl<'a, P: Parent> ImplFor<'a, P> {
         ITER: IntoIterator,
         ITER::Item: Into<String>,
     {
-        self.generics = Some(generics.into_iter().map(Into::into).collect());
+        self.trait_generics = Some(generics.into_iter().map(Into::into).collect());
+        self
+    }
+
+    /// Add generic parameters to the impl block.
+    ///```
+    /// # use virtue::prelude::Generator;
+    /// # let mut generator = Generator::with_name("Bar");
+    /// generator.impl_for("Foo")
+    ///          .with_impl_generics(["Baz"]);
+    /// # generator.assert_eq("impl < Baz > Foo for Bar { }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// impl<Baz> Foo for Bar { }
+    /// ```
+    pub fn with_impl_generics<ITER>(mut self, generics: ITER) -> Self
+    where
+        ITER: IntoIterator,
+        ITER::Item: Into<String>,
+    {
+        self.impl_generics = generics.into_iter().map(Into::into).collect();
         self
     }
 
@@ -282,20 +307,24 @@ impl<P: Parent> Drop for ImplFor<'_, P> {
 impl<P: Parent> ImplFor<'_, P> {
     fn generate_impl_definition(&mut self, builder: &mut StreamBuilder) {
         builder.ident_str("impl");
+
+        let impl_generics = self.impl_generics.as_slice();
         if let Some(lifetimes) = &self.lifetimes {
             if let Some(generics) = self.generator.generics() {
-                builder.append(generics.impl_generics_with_additional_lifetimes(lifetimes));
+                builder.append(generics.impl_generics_with_additional(lifetimes, impl_generics));
             } else {
-                append_lifetimes_and_generics(builder, lifetimes, &[]);
+                append_lifetimes_and_generics(builder, lifetimes, impl_generics);
             }
         } else if let Some(generics) = self.generator.generics() {
-            builder.append(generics.impl_generics());
+            builder.append(generics.impl_generics_with_additional(&[], impl_generics));
+        } else if !impl_generics.is_empty() {
+            append_lifetimes_and_generics(builder, &[], impl_generics)
         }
         if let Some(t) = &self.trait_name {
             builder.push_parsed(t.to_string()).unwrap();
 
             let lifetimes = self.lifetimes.as_deref().unwrap_or_default();
-            let generics = self.generics.as_deref().unwrap_or_default();
+            let generics = self.trait_generics.as_deref().unwrap_or_default();
             append_lifetimes_and_generics(builder, lifetimes, generics);
             builder.ident_str("for");
         }
@@ -333,7 +362,7 @@ fn append_lifetimes_and_generics(
         if idx > 0 || !lifetimes.is_empty() {
             builder.punct(',');
         }
-        builder.ident_str(gen);
+        builder.push_parsed(gen).unwrap();
     }
 
     builder.punct('>');

src/parse/body.rs

@@ -364,6 +364,7 @@ pub struct EnumVariant {
     /// - `Baz = 5`
     /// - `Baz(i32) = 5`
     /// - `Baz { a: i32} = 5`
+    ///
     /// In either case this value will be `Some(Literal::i32(5))`
     pub value: Option<Literal>,
     /// The attributes of this variant

src/parse/generics.rs

@@ -114,20 +114,39 @@ impl Generics {
         result
     }
 
-    pub(crate) fn impl_generics_with_additional_lifetimes(
+    pub(crate) fn impl_generics_with_additional(
         &self,
-        lifetime: &[String],
+        lifetimes: &[String],
+        types: &[String],
     ) -> StreamBuilder {
         let mut result = StreamBuilder::new();
-        for (idx, lt) in lifetime.iter().enumerate() {
-            result.punct(if idx == 0 { '<' } else { ',' });
+        result.punct('<');
+        let mut is_first = true;
+        for lt in lifetimes.iter() {
+            if !is_first {
+                result.punct(',');
+            } else {
+                is_first = false;
+            }
             result.lifetime_str(lt);
         }
 
         for generic in self.iter() {
-            result.punct(',');
+            if !is_first {
+                result.punct(',');
+            } else {
+                is_first = false;
+            }
             generic.append_to_result_with_constraints(&mut result);
         }
+        for ty in types {
+            if !is_first {
+                result.punct(',');
+            } else {
+                is_first = false;
+            }
+            result.ident_str(ty);
+        }
 
         result.punct('>');