fix bug if the value is a dict (#447)

neuralmagic · Feb 13, 2024 · 0e3d685 · 0e3d685
1 parent 91083c3
commit 0e3d685
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/src/sparsezoo/api/utils.py b/src/sparsezoo/api/utils.py
@@ -42,11 +42,12 @@ def map_keys(
     """
     mapped_dict = {}
     for key, value in dictionary.items():
-        if isinstance(value, List) or isinstance(value, Dict):
-            value_type = type(value)
-            mapped_dict[mapper(key)] = value_type(
+        if isinstance(value, List):
+            mapped_dict[mapper(key)] = [
                 map_keys(dictionary=sub_dict, mapper=mapper) for sub_dict in value
-            )
+            ]
+        elif isinstance(value, Dict):
+            mapped_dict[mapper(key)] = dict(map_keys(dictionary=value, mapper=mapper))
         else:
             mapped_dict[mapper(key)] = value
 

diff --git a/tests/sparsezoo/api/test_graphql.py b/tests/sparsezoo/api/test_graphql.py
@@ -121,6 +121,11 @@ def fetch(
             },
             "fields": None,
         },
+        {
+            "operation_body": "models",
+            "arguments": {"stub": "zoo:mobilenet_v2-1.0-imagenet-base"},
+            "fields": {"analysis": {"analysisId": None}},
+        },
     ],
 )
 def test_graphql_api_response(query_args: Dict[str, Any]):

diff --git a/tests/sparsezoo/api/test_query_parser.py b/tests/sparsezoo/api/test_query_parser.py
@@ -172,6 +172,18 @@
                 ),
             },
         ),
+        (
+            {
+                "operation_body": "models",
+                "arguments": {"stub": "zoo:mobilenet_v2-1.0-imagenet-base"},
+                "fields": {"analysis": {"analysisId": None}},
+            },
+            {
+                "operation_body": "models",
+                "arguments": '(stub: "zoo:mobilenet_v2-1.0-imagenet-base",)',
+                "fields": "analysis { analysisId } ",
+            },
+        ),
     ],
 )
 def test_query_parser(