Added multiple params as key support for query_map (#351)

polkascan · Aug 18, 2023 · 111ff53 · 111ff53
1 parent 3826ebe
commit 111ff53
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 16 deletions.
diff --git a/substrateinterface/base.py b/substrateinterface/base.py
@@ -786,8 +786,8 @@ def query_map(self, module: str, storage_function: str, params: Optional[list] =
         if len(param_types) == 0:
             raise ValueError('Given storage function is not a map')
 
-        if len(params) != len(param_types) - 1:
-            raise ValueError(f'Storage function map requires {len(param_types) -1} parameters, {len(params)} given')
+        if len(params) > len(param_types) - 1:
+            raise ValueError(f'Storage function map can accept max {len(param_types) - 1} parameters, {len(params)} given')
 
         # Generate storage key prefix
         storage_key = StorageKey.create_from_storage_function(
@@ -815,9 +815,9 @@ def query_map(self, module: str, storage_function: str, params: Optional[list] =
 
         def concat_hash_len(key_hasher: str) -> int:
             if key_hasher == "Blake2_128Concat":
-                return 32
-            elif key_hasher == "Twox64Concat":
                 return 16
+            elif key_hasher == "Twox64Concat":
+                return 8
             elif key_hasher == "Identity":
                 return 0
             else:
@@ -836,12 +836,27 @@ def concat_hash_len(key_hasher: str) -> int:
             for result_group in response['result']:
                 for item in result_group['changes']:
                     try:
-                        item_key = self.decode_scale(
-                            type_string=param_types[len(params)],
-                            scale_bytes='0x' + item[0][len(prefix) + concat_hash_len(key_hashers[len(params)]):],
+                        # Determine type string
+                        key_type_string = []
+                        for n in range(len(params), len(param_types)):
+                            key_type_string.append(f'[u8; {concat_hash_len(key_hashers[n])}]')
+                            key_type_string.append(param_types[n])
+
+                        item_key_obj = self.decode_scale(
+                            type_string=f"({', '.join(key_type_string)})",
+                            scale_bytes='0x' + item[0][len(prefix):],
                             return_scale_obj=True,
                             block_hash=block_hash
                         )
+
+                        # strip key_hashers to use as item key
+                        if len(param_types) - len(params) == 1:
+                            item_key = item_key_obj.value_object[1]
+                        else:
+                            item_key = tuple(
+                                item_key_obj.value_object[key + 1] for key in range(len(params), len(param_types) + 1, 2)
+                            )
+
                     except Exception:
                         if not ignore_decoding_errors:
                             raise

diff --git a/test/test_query_map.py b/test/test_query_map.py
@@ -213,13 +213,15 @@ def test_double_map_no_result(self):
         )
         self.assertEqual(era_stakers.records, [])
 
-    def test_double_map_missing_param(self):
-        with self.assertRaises(ValueError) as cm:
-            self.kusama_substrate.query_map(
-                module='Staking',
-                storage_function='ErasStakers'
-            )
-        self.assertEqual('Storage function map requires 1 parameters, 0 given', str(cm.exception))
+    def test_nested_keys(self):
+
+        result = self.kusama_substrate.query_map(
+            module='ConvictionVoting',
+            storage_function='VotingFor',
+            max_results=10
+        )
+        self.assertTrue(self.kusama_substrate.is_valid_ss58_address(result[0][0][0].value))
+        self.assertGreaterEqual(result[0][0][1], 0)
 
     def test_double_map_too_many_params(self):
         with self.assertRaises(ValueError) as cm:
@@ -228,7 +230,7 @@ def test_double_map_too_many_params(self):
                 storage_function='ErasStakers',
                 params=[21000000, 2]
             )
-        self.assertEqual('Storage function map requires 1 parameters, 2 given', str(cm.exception))
+        self.assertEqual('Storage function map can accept max 1 parameters, 2 given', str(cm.exception))
 
     def test_map_with_param(self):
         with self.assertRaises(ValueError) as cm:
@@ -237,7 +239,7 @@ def test_map_with_param(self):
                 storage_function='Account',
                 params=[2]
             )
-        self.assertEqual('Storage function map requires 0 parameters, 1 given', str(cm.exception))
+        self.assertEqual('Storage function map can accept max 0 parameters, 1 given', str(cm.exception))
 
 
 if __name__ == '__main__':