fix extract function

dipdup-io · Jan 29, 2025 · d91a7bd · d91a7bd
1 parent b466aae
commit d91a7bd
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 43 deletions.
diff --git a/src/dipdup/runtimes.py b/src/dipdup/runtimes.py
@@ -270,46 +270,26 @@ def parse(value: Any, type_: str) -> Any:
 
 
 def extract_subsquid_payload(data: Any) -> Any:
-    if isinstance(data, list):
+    if isinstance(data, list | tuple):
         return tuple(extract_subsquid_payload(item) for item in data)
 
     if isinstance(data, dict):
-        if 'interior' in data:
-            # Handle interior dictionary
-            result = {'parents': data['parents']}
-            interior = data['interior']
-
-            if '__kind' in interior:
-                kind = interior['__kind']
-                if 'value' in interior:
-                    value = interior['value']
-                    if isinstance(value, list):
-                        # Handle list of values
-                        value = tuple(
-                            (
-                                {
-                                    item['__kind']: (
-                                        int(item['value'])
-                                        if isinstance(item['value'], str)
-                                        else (
-                                            tuple(item['value']) if isinstance(item['value'], list) else item['value']
-                                        )
-                                    )
-                                }
-                                if item.get('value')
-                                else item['__kind']
-                            )
-                            for item in value
-                        )
-                    elif isinstance(value, str):
-                        value = int(value)
-                    result['interior'] = {kind: value}
-                else:
-                    # If 'value' is not present, just use the kind
-                    result['interior'] = kind
-            else:
-                result['interior'] = extract_subsquid_payload(interior)
-            return result
+
+        if '__kind' in data:
+            kind = data['__kind']
+            if 'value' in data:
+                value = data['value']
+                if isinstance(value, list):
+                    # Handle list of values
+                    value = tuple(extract_subsquid_payload(item) for item in value)
+                elif isinstance(value, str):
+                    value = int(value)
+                return {kind: value}
+            # NOTE: Special case
+            if 'key' in data:
+                return {kind: data['key']}
+            # If 'value' is not present, just use the kind
+            return kind
 
         # Process other dictionaries
         return {key: extract_subsquid_payload(value) for key, value in data.items()}

diff --git a/tests/test_datasources/test_substrate.py b/tests/test_datasources/test_substrate.py
@@ -1,8 +1,6 @@
-
 from dipdup.runtimes import extract_subsquid_payload
 
-
-path = [
+path_1 = [
     [
         {
             'parents': 0,
@@ -32,8 +30,35 @@
         122612710,
     ],
 ]
+path_2 = [
+    {
+        'interior': {
+            '__kind': 'X3',
+            'value': [
+                {'__kind': 'Parachain', 'value': 1000},
+                {'__kind': 'GeneralIndex', 'value': 50},
+                {'__kind': 'GeneralIndex', 'value': 42069},
+            ],
+        },
+        'parents': 1,
+    }
+]
+
+path_3 = [
+    {
+        'interior': {
+            '__kind': 'X3',
+            'value': [
+                {'__kind': 'Parachain', 'value': 2004},
+                {'__kind': 'PalletInstance', 'value': 110},
+                {'__kind': 'AccountKey20', 'key': 39384093},
+            ],
+        },
+        'parents': 1,
+    }
+]
 
-processed_path = (
+processed_path_1 = (
     (
         {
             'parents': 0,
@@ -55,7 +80,35 @@
     ),
 )
 
+processed_path_2 = (
+    {
+        'parents': 1,
+        'interior': {
+            'X3': (
+                {'Parachain': 1000},
+                {'GeneralIndex': 50},
+                {'GeneralIndex': 42069},
+            ),
+        },
+    },
+)
+
+processed_path_3 = (
+    {
+        'parents': 1,
+        'interior': {
+            'X3': (
+                {'Parachain': 2004},
+                {'PalletInstance': 110},
+                {'AccountKey20': 39384093},
+            ),
+        },
+    },
+)
+
 
-def test_extract_subsquid_payload():
+def test_extract_subsquid_payload() -> None:
 
-    assert extract_subsquid_payload(path) == processed_path
+    assert extract_subsquid_payload(path_1) == processed_path_1
+    assert extract_subsquid_payload(path_2) == processed_path_2
+    assert extract_subsquid_payload(path_3) == processed_path_3