Merge pull request #1429 from CosmWasm/query-responses-generics

Handle `QueryMsg` with generics when deriving the `QueryResponses` trait
CosmWasm · Sep 19, 2022 · c468bcb · c468bcb
2 parents 9753cfa + a20d5f9
commit c468bcb
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,8 +20,11 @@ and this project adheres to
 
 - cosmwasm-vm: Bump `MODULE_SERIALIZATION_VERSION` to "v4" because the module
   serialization format changed between Wasmer 2.2 and 2.3 ([#1426]).
+- cosmwasm-schema: The `QueryResponses` derive macro now supports `QueryMsg`s
+  with generics. ([#1429])
 
 [#1426]: https://github.com/CosmWasm/cosmwasm/issues/1426
+[#1429]: https://github.com/CosmWasm/cosmwasm/pull/1429
 
 ## [1.1.1] - 2022-09-15
 

diff --git a/packages/schema-derive/src/query_responses.rs b/packages/schema-derive/src/query_responses.rs
@@ -7,10 +7,13 @@ pub fn query_responses_derive_impl(input: ItemEnum) -> ItemImpl {
     queries.sort();
     let mappings = mappings.map(parse_tuple);
 
+    // Handle generics if the type has any
+    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
+
     parse_quote! {
         #[automatically_derived]
         #[cfg(not(target_arch = "wasm32"))]
-        impl ::cosmwasm_schema::QueryResponses for #ident {
+        impl #impl_generics ::cosmwasm_schema::QueryResponses for #ident #type_generics #where_clause {
             fn response_schemas_impl() -> ::std::collections::BTreeMap<String, ::cosmwasm_schema::schemars::schema::RootSchema> {
                 ::std::collections::BTreeMap::from([
                     #( #mappings, )*
@@ -113,6 +116,95 @@ mod tests {
         );
     }
 
+    #[test]
+    fn generics() {
+        let input: ItemEnum = parse_quote! {
+            #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema, QueryResponses)]
+            #[serde(rename_all = "snake_case")]
+            pub enum QueryMsg<T> {
+                #[returns(bool)]
+                Foo,
+                #[returns(u32)]
+                Bar(T),
+            }
+        };
+
+        let input2: ItemEnum = parse_quote! {
+            #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema, QueryResponses)]
+            #[serde(rename_all = "snake_case")]
+            pub enum QueryMsg<T: std::fmt::Debug + SomeTrait> {
+                #[returns(bool)]
+                Foo,
+                #[returns(u32)]
+                Bar { data: T },
+            }
+        };
+
+        let input3: ItemEnum = parse_quote! {
+            #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema, QueryResponses)]
+            #[serde(rename_all = "snake_case")]
+            pub enum QueryMsg<T>
+                where T: std::fmt::Debug + SomeTrait,
+            {
+                #[returns(bool)]
+                Foo,
+                #[returns(u32)]
+                Bar { data: T },
+            }
+        };
+
+        let result = query_responses_derive_impl(input);
+        dbg!(&result);
+        assert_eq!(
+            result,
+            parse_quote! {
+                #[automatically_derived]
+                #[cfg(not(target_arch = "wasm32"))]
+                impl<T> ::cosmwasm_schema::QueryResponses for QueryMsg<T> {
+                    fn response_schemas_impl() -> ::std::collections::BTreeMap<String, ::cosmwasm_schema::schemars::schema::RootSchema> {
+                        ::std::collections::BTreeMap::from([
+                            ("foo".to_string(), ::cosmwasm_schema::schema_for!(bool)),
+                            ("bar".to_string(), ::cosmwasm_schema::schema_for!(u32)),
+                        ])
+                    }
+                }
+            }
+        );
+        assert_eq!(
+            query_responses_derive_impl(input2),
+            parse_quote! {
+                #[automatically_derived]
+                #[cfg(not(target_arch = "wasm32"))]
+                impl<T: std::fmt::Debug + SomeTrait> ::cosmwasm_schema::QueryResponses for QueryMsg<T> {
+                    fn response_schemas_impl() -> ::std::collections::BTreeMap<String, ::cosmwasm_schema::schemars::schema::RootSchema> {
+                        ::std::collections::BTreeMap::from([
+                            ("foo".to_string(), ::cosmwasm_schema::schema_for!(bool)),
+                            ("bar".to_string(), ::cosmwasm_schema::schema_for!(u32)),
+                        ])
+                    }
+                }
+            }
+        );
+        let a = query_responses_derive_impl(input3);
+        assert_eq!(
+            a,
+            parse_quote! {
+                #[automatically_derived]
+                #[cfg(not(target_arch = "wasm32"))]
+                impl<T> ::cosmwasm_schema::QueryResponses for QueryMsg<T>
+                    where T: std::fmt::Debug + SomeTrait,
+                {
+                    fn response_schemas_impl() -> ::std::collections::BTreeMap<String, ::cosmwasm_schema::schemars::schema::RootSchema> {
+                        ::std::collections::BTreeMap::from([
+                            ("foo".to_string(), ::cosmwasm_schema::schema_for!(bool)),
+                            ("bar".to_string(), ::cosmwasm_schema::schema_for!(u32)),
+                        ])
+                    }
+                }
+            }
+        );
+    }
+
     #[test]
     #[should_panic(expected = "missing return type for query: Supply")]
     fn missing_return() {

diff --git a/packages/schema/tests/idl.rs b/packages/schema/tests/idl.rs
@@ -106,3 +106,43 @@ fn test_query_responses() {
     // Find the "balance" query in responses
     api.get("responses").unwrap().get("balance").unwrap();
 }
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, QueryResponses)]
+#[serde(rename_all = "snake_case")]
+pub enum QueryMsgWithGenerics<T: std::fmt::Debug>
+where
+    T: JsonSchema,
+{
+    #[returns(u128)]
+    QueryData { data: T },
+}
+
+#[test]
+fn test_query_responses_generics() {
+    let api_str = generate_api! {
+        instantiate: InstantiateMsg,
+        query: QueryMsgWithGenerics<u32>,
+    }
+    .render()
+    .to_string()
+    .unwrap();
+
+    let api: Value = serde_json::from_str(&api_str).unwrap();
+    let queries = api
+        .get("query")
+        .unwrap()
+        .get("oneOf")
+        .unwrap()
+        .as_array()
+        .unwrap();
+
+    // Find the "balance" query in the queries schema
+    assert_eq!(queries.len(), 1);
+    assert_eq!(
+        queries[0].get("required").unwrap().get(0).unwrap(),
+        "query_data"
+    );
+
+    // Find the "balance" query in responses
+    api.get("responses").unwrap().get("query_data").unwrap();
+}