Fix merging storage which contains specific nested structures (#159)

Author: droserasprout

scripts/init_demos.sh

@@ -1,6 +1,6 @@
 #$/bin/bash
 for name in `ls src | grep demo`
 do
-    dipdup -c src/$name/dipdup.yml init
+    dipdup -c src/$name/dipdup.yml init --overwrite-types
     dipdup -c src/$name/dipdup.yml -c src/$name/dipdup-docker.yml docker init
 done

src/dipdup/models.py

@@ -9,6 +9,7 @@
 from pydantic.dataclasses import dataclass
 from pydantic.error_wrappers import ValidationError
 from tortoise import Model, fields
+from typing_extensions import get_args
 
 from dipdup.enums import IndexStatus, IndexType
 from dipdup.exceptions import ConfigurationError, InvalidDataError, ReindexingReason
@@ -77,7 +78,6 @@ def _process_storage(
         prefix: str = None,
     ) -> Dict[str, Any]:
         for key, field in storage_type.__fields__.items():
-
             if key == '__root__':
                 continue
 
@@ -95,8 +95,16 @@ def _process_storage(
                     continue
                 raise ConfigurationError(f'Type `{storage_type.__name__}` is invalid: `{key}` field does not exists') from e
 
-            # FIXME: Pydantic bug. `BaseModel.type_` returns incorrect value when annotation is Dict[str, bool]
-            if field.type_ != field.outer_type_ and field.type_ == bool:
+            # FIXME: Pydantic bug? I have no idea how does it work, this workaround is just a guess.
+            # FIXME: `BaseModel.type_` returns incorrect value when annotation is Dict[str, bool], Dict[str, BaseModel], and possibly any other cases.
+            is_complex_type = field.type_ != field.outer_type_
+            try:
+                get_args(field.outer_type_)[1].__fields__
+                is_nested_dict_model = True
+            except Exception:
+                is_nested_dict_model = False
+
+            if is_complex_type and (field.type_ == bool or is_nested_dict_model):
                 annotation = field.outer_type_
             else:
                 annotation = field.type_

tests/test_dipdup/models.py

@@ -0,0 +1,47 @@
+# generated by datamodel-codegen:
+#   filename:  storage.json
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Extra
+
+
+class ResourceMap(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    id: str
+    rate: str
+
+
+class ResourceCollectorStorage(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    administrator: str
+    current_user: Optional[str]
+    default_start_time: str
+    generation_rate: str
+    managers: List[str]
+    metadata: Dict[str, str]
+    nft_registry: str
+    paused: bool
+    resource_map: Dict[str, ResourceMap]
+    resource_registry: str
+    tezotop_collection: Dict[str, str]
+
+
+# 'resource_map': {
+#     'type': 'object',
+#     'propertyNames': {'type': 'string', '$comment': 'string'},
+#     'additionalProperties': {
+#         'type': 'object',
+#         'properties': {'id': {'type': 'string', '$comment': 'nat'}, 'rate': {'type': 'string', '$comment': 'nat'}},
+#         'required': ['id', 'rate'],
+#         'additionalProperties': False,
+#         '$comment': 'pair',
+#     },
+#     '$comment': 'map',
+# },

tests/test_dipdup/test_models.py

@@ -1,11 +1,13 @@
+from datetime import datetime
 from unittest import TestCase
 
 from demo_tezos_domains.types.name_registry.storage import NameRegistryStorage
 from dipdup.models import OperationData
+from tests.test_dipdup.models import ResourceCollectorStorage
 
 
 class ModelsTest(TestCase):
-    def test_merged_storage(self):
+    def test_merged_storage(self) -> None:
         storage = {
             'store': {
                 'data': 15023,
@@ -76,7 +78,7 @@ def test_merged_storage(self):
             type='transaction',
             id=0,
             level=0,
-            timestamp=0,
+            timestamp=datetime.now(),
             hash='',
             counter=0,
             sender_address='',
@@ -88,3 +90,40 @@ def test_merged_storage(self):
         )
         merged_storage = operation_data.get_merged_storage(NameRegistryStorage)
         self.assertTrue('6672657175656e742d616e616c7973742e65646f' in merged_storage.store.records)
+
+    def test_merged_storage_dict_of_dicts(self) -> None:
+        storage = {
+            'paused': False,
+            'managers': ['tz1VPZyh4ZHjDDpgvznqQQXUCLcV7g91WGMz'],
+            'metadata': 43542,
+            'current_user': 'tz1VPZyh4ZHjDDpgvznqQQXUCLcV7g91WGMz',
+            'nft_registry': 'KT1SZ87ihAWc43YZxYjoRz8MQyAapUGbZigG',
+            'resource_map': {
+                'enr': {'id': '2', 'rate': '1875000'},
+                'mch': {'id': '3', 'rate': '625000'},
+                'min': {'id': '1', 'rate': '1250000'},
+                'uno': {'id': '0', 'rate': '1250'},
+            },
+            'administrator': 'tz1VPZyh4ZHjDDpgvznqQQXUCLcV7g91WGMz',
+            'generation_rate': '9',
+            'resource_registry': 'KT1SLaZNaDF7V6Lt8FXTbYNqkBS81gjHXMsP',
+            'default_start_time': '1630678200',
+            'tezotop_collection': 43543,
+        }
+        operation_data = OperationData(
+            storage=storage,
+            diffs=None,
+            type='transaction',
+            id=0,
+            level=0,
+            timestamp=datetime.now(),
+            hash='',
+            counter=0,
+            sender_address='',
+            target_address='',
+            initiator_address='',
+            amount=0,
+            status='',
+            has_internals=False,
+        )
+        operation_data.get_merged_storage(ResourceCollectorStorage)