Fix parsing storage with nested array fields (#207)

* Fix nested sets in tzkt

* Lint

* Cleanup

* changelog

* Refactor storage processing

* Linter fixes

* Fix tests

* Schemas, scripts

* Fix jsonschema preprocessing

* Cleanup, fix tests

* naming

* Naming

Co-authored-by: Michael Zaikin <[email protected]>

Author: droserasprout

CHANGELOG.md

@@ -4,10 +4,15 @@ Please use [this](https://docs.gitlab.com/ee/development/changelog.html) documen
 
 ## [unreleased]
 
+### Added
+
+* cli: Added `--keep-schemas` flag to `init` command to preserve JSONSchemas along with generated types.
+
 ### Fixed
 
 * demos: Tezos Domains and Homebase DAO demos were updated from edo2net to mainnet contracts.
 * hasura: Fixed missing relations for models with `ManyToManyField` fields.
+* tzkt: Fixed parsing storage with nested structures.
 
 ### Performance
 

scripts/init_tests.sh

@@ -1,10 +1,13 @@
 #$/bin/bash
 cd tests/test_dipdup
-for name in "asdf" "qwer" "hjkl" "zxcv" "rewq" "hen_subjkt" "kolibri_ovens"
+for name in "asdf" "qwer" "hjkl" "zxcv" "rewq" "hen_subjkt" "kolibri_ovens" "yupana"
 do
-    dipdup -c $name.yml init
+    dipdup -c $name.yml init --keep-schemas
     mkdir -p types/$name/
+    mkdir -p schemas/$name/
     mv $name/types/$name/storage.py types/$name/storage.py
+    touch types/$name/__init__.py
+    mv $name/schemas/$name/storage.json schemas/$name/storage.json
     mv $name/types/$name/parameter/set_delegate.py types/$name/set_delegate.py || true
     rm -r $name
 done

src/dipdup/cli.py

@@ -201,12 +201,13 @@ async def run(
 
 @cli.command(help='Generate missing callbacks and types')
 @click.option('--overwrite-types', is_flag=True, help='Regenerate existing types')
+@click.option('--keep-schemas', is_flag=True, help='Do not remove JSONSchemas after generating types')
 @click.pass_context
 @cli_wrapper
-async def init(ctx, overwrite_types: bool):
+async def init(ctx, overwrite_types: bool, keep_schemas: bool) -> None:
     config: DipDupConfig = ctx.obj.config
     dipdup = DipDup(config)
-    await dipdup.init(overwrite_types)
+    await dipdup.init(overwrite_types, keep_schemas)
 
 
 @cli.command(help='Migrate project to the new spec version')

src/dipdup/codegen.py

@@ -66,6 +66,8 @@ def preprocess_storage_jsonschema(schema: Dict[str, Any]) -> Dict[str, Any]:
     We resolve bigmaps from diffs so no need to include int in type signature."""
     if not isinstance(schema, dict):
         return schema
+    if 'oneOf' in schema:
+        schema['oneOf'] = [preprocess_storage_jsonschema(sub_schema) for sub_schema in schema['oneOf']]
     if 'properties' in schema:
         return {
             **schema,
@@ -104,14 +106,15 @@ def __init__(self, config: DipDupConfig, datasources: Dict[DatasourceConfigT, Da
         self._datasources = datasources
         self._schemas: Dict[TzktDatasourceConfig, Dict[str, Dict[str, Any]]] = {}
 
-    async def init(self, overwrite_types: bool = False) -> None:
+    async def init(self, overwrite_types: bool = False, keep_schemas: bool = False) -> None:
         self._logger.info('Initializing project')
         await self.create_package()
         await self.fetch_schemas()
         await self.generate_types(overwrite_types)
         await self.generate_hooks()
         await self.generate_handlers()
-        await self.cleanup()
+        if not keep_schemas:
+            await self.cleanup()
         await self.verify_package()
 
     async def docker_init(self, image: str, tag: str, env_file: str) -> None:

src/dipdup/datasources/tzkt/models.py

@@ -1,3 +1,4 @@
+import typing
 from contextlib import suppress
 from functools import lru_cache
 from itertools import groupby
@@ -6,11 +7,11 @@
 from typing import Iterable
 from typing import List
 from typing import Optional
+from typing import Tuple
 from typing import Type
 from typing import Union
 
 from pydantic.error_wrappers import ValidationError
-from pydantic.fields import FieldInfo
 from typing_extensions import get_args
 from typing_extensions import get_origin
 
@@ -21,74 +22,83 @@
 IntrospectionError = (KeyError, IndexError, AttributeError)
 
 
-def _extract_root_type(storage_type: Type) -> Type:
+def extract_root_outer_type(storage_type: Type) -> Type:
     """Extract Pydantic __root__ type"""
-    return storage_type.__fields__['__root__'].type_
+    root_field = storage_type.__fields__['__root__']
+    if root_field.allow_none:
+        return typing.Optional[root_field.type_]  # type: ignore
+    else:
+        return root_field.outer_type_
 
 
 @lru_cache(None)
-def _is_array(storage_type: Type) -> bool:
+def is_array_type(storage_type: Type) -> bool:
     """TzKT can return bigmaps as objects or as arrays of key-value objects. Guess it from storage type."""
     # NOTE: List[...]
     if get_origin(storage_type) == list:
         return True
 
-    # NOTE: Neither a list not Pydantic model, can't be an array
-    fields: Optional[Dict[str, FieldInfo]] = getattr(storage_type, '__fields__', None)
-    if fields is None:
-        return False
-
-    # NOTE: An item of TzKT array
-    if 'key' in fields and 'value' in fields:
-        return True
-
-    # NOTE: Pydantic model with __root__ field, dive into it
+    # NOTE: Pydantic model with __root__ field subclassing List
     with suppress(*IntrospectionError):
-        root_type = _extract_root_type(storage_type)
-        return _is_array(root_type)  # type: ignore
+        root_type = extract_root_outer_type(storage_type)
+        return is_array_type(root_type)  # type: ignore
 
     # NOTE: Something else
     return False
 
 
 @lru_cache(None)
-def _extract_list_types(storage_type: Type[Any]) -> Iterable[Type[Any]]:
-    """Extract list item types from field type"""
-    # NOTE: Pydantic model with __root__ field
-    with suppress(*IntrospectionError):
-        return (_extract_root_type(storage_type),)
+def get_list_elt_type(list_type: Type[Any]) -> Type[Any]:
+    """Extract list item type from list type"""
+    # NOTE: regular list
+    if get_origin(list_type) == list:
+        return get_args(list_type)[0]
 
-    # NOTE: Python list, return all args unpacking unions
-    with suppress(*IntrospectionError):
-        item_type = get_args(storage_type)[0]
-        if get_origin(item_type) == Union:
-            return get_args(item_type)
-        return (item_type,)
-
-    # NOTE: Something else
-    return ()
+    # NOTE: Pydantic model with __root__ field subclassing List
+    root_type = extract_root_outer_type(list_type)
+    return get_list_elt_type(root_type)  # type: ignore
 
 
 @lru_cache(None)
-def _extract_dict_types(storage_type: Type[Any], key: str) -> Iterable[Type[Any]]:
+def get_dict_value_type(dict_type: Type[Any], key: Optional[str] = None) -> Type[Any]:
     """Extract dict value types from field type"""
     # NOTE: Regular dict
-    if get_origin(storage_type) == dict:
-        return (get_args(storage_type)[1],)
+    if get_origin(dict_type) == dict:
+        return get_args(dict_type)[1]
+
+    # NOTE: Pydantic model with __root__ field subclassing Dict
+    with suppress(*IntrospectionError):
+        root_type = extract_root_outer_type(dict_type)
+        return get_dict_value_type(root_type, key)  # type: ignore
 
-    # NOTE: Unpack union args
-    if get_origin(storage_type) == Union:
-        return get_args(storage_type)
+    if key is None:
+        raise KeyError('Key name or alias is required for object introspection')
 
     # NOTE: Pydantic model, find corresponding field and return it's type
+    fields = dict_type.__fields__
+    for field in fields.values():
+        if key in (field.name, field.alias):
+            # NOTE: Pydantic does not preserve outer_type_ for Optional
+            if field.allow_none:
+                return typing.Optional[field.type_]  # type: ignore
+            else:
+                return field.outer_type_
+
+    # NOTE: typically when we try the wrong Union path
+    raise KeyError('Key not found')
+
+
+@lru_cache(None)
+def unwrap_union_type(union_type: Type) -> Tuple[bool, Tuple[Type, ...]]:
+    """Check if the type is either optional or union and return arg types if so"""
+    if get_origin(union_type) == Union:
+        return True, get_args(union_type)
+
     with suppress(*IntrospectionError):
-        fields = storage_type.__fields__
-        for field in fields.values():
-            if key in (field.name, field.alias):
-                return (field.type_,)
+        root_type = extract_root_outer_type(union_type)
+        return unwrap_union_type(root_type)  # type: ignore
 
-    # NOTE: Something else
-    return ()
+    return False, ()
 
 
 def _preprocess_bigmap_diffs(diffs: Iterable[Dict[str, Any]]) -> Dict[int, Iterable[Dict[str, Any]]]:
@@ -124,30 +134,32 @@ def _apply_bigmap_diffs(
         return dict_storage
 
 
-def _process_storage(
-    storage: Any,
-    storage_type: Type[StorageType],
-    bigmap_diffs: Dict[int, Iterable[Dict[str, Any]]],
-) -> Any:
+def _process_storage(storage: Any, storage_type: Type[Any], bigmap_diffs: Dict[int, Iterable[Dict[str, Any]]]) -> Any:
     """Replace bigmap pointers with actual data from diffs"""
+    # Check if Union or Optional (== Union[Any, NoneType])
+    is_union, arg_types = unwrap_union_type(storage_type)  # type: ignore
+    if is_union:
+        # NOTE: We have no way but trying every possible branch until first success
+        for arg_type in arg_types:
+            with suppress(*IntrospectionError):
+                return _process_storage(storage, arg_type, bigmap_diffs)
+
     # NOTE: Bigmap pointer, apply diffs
     if isinstance(storage, int) and type(storage) != storage_type:
-        is_array = _is_array(storage_type)
+        is_array = is_array_type(storage_type)  # type: ignore
         storage = _apply_bigmap_diffs(storage, bigmap_diffs, is_array)
 
     # NOTE: List, process recursively
     elif isinstance(storage, list):
+        elt_type = get_list_elt_type(storage_type)  # type: ignore
         for i, _ in enumerate(storage):
-            for item_type in _extract_list_types(storage_type):
-                with suppress(*IntrospectionError):
-                    storage[i] = _process_storage(storage[i], item_type, bigmap_diffs)
+            storage[i] = _process_storage(storage[i], elt_type, bigmap_diffs)
 
     # NOTE: Dict, process recursively
     elif isinstance(storage, dict):
         for key, value in storage.items():
-            for value_type in _extract_dict_types(storage_type, key):
-                with suppress(*IntrospectionError):
-                    storage[key] = _process_storage(value, value_type, bigmap_diffs)
+            value_type = get_dict_value_type(storage_type, key)  # type: ignore
+            storage[key] = _process_storage(value, value_type, bigmap_diffs)
 
     else:
         pass

src/dipdup/dipdup.py

@@ -323,15 +323,15 @@ def schema(self) -> Schema:
             raise DipDupException('Schema is not initialized')
         return self._schema
 
-    async def init(self, overwrite_types: bool = True) -> None:
+    async def init(self, overwrite_types: bool = False, keep_schemas: bool = False) -> None:
         """Create new or update existing dipdup project"""
         await self._create_datasources()
 
         async with AsyncExitStack() as stack:
             for datasource in self._datasources.values():
                 await stack.enter_async_context(datasource)
 
-            await self._codegen.init(overwrite_types)
+            await self._codegen.init(overwrite_types, keep_schemas)
 
     async def docker_init(self, image: str, tag: str, env_file: str) -> None:
         await self._codegen.docker_init(image, tag, env_file)

src/dipdup/exceptions.py

@@ -270,7 +270,7 @@ def _help(self) -> str:
             Failed to validate datasource message against generated type class.
 
             Expected type:
-            `{self.type_cls.__class__.__qualname__}`
+            `{self.type_cls.__name__}`
 
             Invalid data:
             {self.data}

tests/test_dipdup/asdf.yml

@@ -9,7 +9,7 @@ contracts:
 datasources:
   tzkt:
     kind: tzkt
-    url: https://api.tzkt.io
+    url: https://api.hangzhou2net.tzkt.io
 
 indexes:
   asdf:

tests/test_dipdup/hjkl.yml

@@ -9,7 +9,7 @@ contracts:
 datasources:
   tzkt:
     kind: tzkt
-    url: https://api.tzkt.io
+    url: https://api.hangzhou2net.tzkt.io
 
 indexes:
   hjkl:

tests/test_dipdup/qwer.yml

@@ -9,7 +9,7 @@ contracts:
 datasources:
   tzkt:
     kind: tzkt
-    url: https://api.tzkt.io
+    url: https://api.hangzhou2net.tzkt.io
 
 indexes:
   qwer: