support HashMap and BTreeMap

3 files changed, 25 insertions, 7 deletions

diff --git a/src/lib.rs b/src/lib.rs
index a5253d3..5e3653d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -413,6 +413,11 @@ impl TypeTransform {
                 let inner = inner.convert(quote!(x));
                 quote! {#arg.into_iter().map(|x| #inner).collect()}
             }
+            TypeTransform::CollectionMap(key, value) => {
+                let key = key.convert(quote!(k));
+                let value = value.convert(quote!(v));
+                quote! {#arg.into_iter().map(|(k,v)| (#key, #value)).collect()}
+            }
             TypeTransform::NoOp => arg,
         }
     }
diff --git a/src/parse_trait_sig.rs b/src/parse_trait_sig.rs
index c924a71..4c0b4fd 100644
--- a/src/parse_trait_sig.rs
+++ b/src/parse_trait_sig.rs
@@ -18,6 +18,7 @@ pub enum TypeTransform {
     Box(BoxType),
     Map(Box<TypeTransform>),
     Collection(Box<TypeTransform>),
+    CollectionMap(Box<TypeTransform>, Box<TypeTransform>),
     Iterator(BoxType, Box<TypeTransform>),
     Result(Box<TypeTransform>, Box<TypeTransform>),
 }
diff --git a/src/transform.rs b/src/transform.rs
index 4dc88f0..4625f4c 100644
--- a/src/transform.rs
+++ b/src/transform.rs
@@ -33,6 +33,12 @@ fn is_supported_collection(ident: &Ident) -> bool {
         || ident == "BTreeSet"
 }
 
+fn is_supported_collection_map(ident: &Ident) -> bool {
+    // collections added here must implement IntoIterator & FromIterator
+    // when adding a type here don't forget to document it in the README
+    ident == "HashMap" || ident == "BTreeMap"
+}
+
 impl AssocTypeConversions<'_> {
     pub fn parse_type_path(&self, type_: &mut Type) -> Result<TypeTransform, TransformError> {
         if !iter_type(type_).any(match_assoc_type) {
@@ -123,14 +129,20 @@ impl AssocTypeConversions<'_> {
                         if let (GenericArgument::Type(arg1), GenericArgument::Type(arg2)) =
                             (args_iter.next().unwrap(), args_iter.next().unwrap())
                         {
-                            if (iter_type(arg1).any(match_assoc_type)
-                                || iter_type(arg2).any(match_assoc_type))
-                                && last_seg.ident == "Result"
+                            if iter_type(arg1).any(match_assoc_type)
+                                || iter_type(arg2).any(match_assoc_type)
                             {
-                                return Ok(TypeTransform::Result(
-                                    self.parse_type_path(arg1)?.into(),
-                                    self.parse_type_path(arg2)?.into(),
-                                ));
+                                if last_seg.ident == "Result" {
+                                    return Ok(TypeTransform::Result(
+                                        self.parse_type_path(arg1)?.into(),
+                                        self.parse_type_path(arg2)?.into(),
+                                    ));
+                                } else if is_supported_collection_map(&last_seg.ident) {
+                                    return Ok(TypeTransform::CollectionMap(
+                                        self.parse_type_path(arg1)?.into(),
+                                        self.parse_type_path(arg2)?.into(),
+                                    ));
+                                }
                             }
                         }
                     }