Merge pull request #16 from disintar/feature/storage-abi

Add storage abi support

Author: nikitunkun

setup.py

@@ -31,7 +31,7 @@ def finalize_options(self):
 
 setup(
     name="tonpy" if not IS_DEV else "tonpy-dev",
-    version="0.0.0.1.4c0" if not IS_DEV else "0.0.0.6.1a1",
+    version="0.0.0.1.4c0" if not IS_DEV else "0.0.0.6.2a1",
     author="Disintar LLP",
     author_email="andrey@head-labs.com",
     description="Types / API for TON blockchain",

src/tonpy/abi/getter.py

@@ -2,29 +2,10 @@
 
 from tonpy import StackEntry, add_tlb, Address, CellSlice
 from tonpy.abi.getter_cache import getter_cache
+from tonpy.abi.utils import parse_tlb_spec, supported_types
 from tonpy.tvm import TVM
 from loguru import logger
 
-supported_types = [
-    'Int8',
-    'Int16',
-    'Int32',
-    'Int64',
-    'Int128',
-    'Int256',
-    'UInt8',
-    'UInt16',
-    'UInt32',
-    'UInt64',
-    'UInt128',
-    'UInt256',
-    'String',
-    'FixedString(64)',
-    'Address',
-    'Boolean',
-    'Datetime'
-]
-
 
 class ABIGetterResultInstance:
     def __init__(self, instance):
@@ -85,31 +66,13 @@ def get_columns(self):
             }
         elif self.type in ['Slice', 'Cell', 'Continuation', 'Builder'] and self.instance.get('tlb', None):
             tlb = self.instance.get('tlb')
-
-            if 'parse' in tlb:
-                data = tlb.get('parse')
-                assert isinstance(data, list)
-
-                tmp = {}
-
-                if tlb['dump_with_types']:
-                    tmp[f'{self.dton_parse_prefix}{self.name}_type'] = 'String'
-
-                for item in data:
-                    path = item.get('path', None)
-                    assert path, "Missing path in TLB"
-
-                    path = path.split('.')
-                    if tmp.get(path[-1]):
-                        raise ValueError(f'Duplicate path {path[-1]}')
-
-                    dtype = item.get('labels', {}).get('dton_type', None)
-                    assert dtype in supported_types, f'Unsupported ABI type {dtype}'
-
-                    name = item.get('labels', {}).get('name', path[-1])
-                    tmp[f'{self.dton_parse_prefix}{self.name}_{name}'] = dtype
-
-                return tmp
+            return parse_tlb_spec(
+                None,
+                tlb,
+                prefix=self.dton_parse_prefix,
+                name_prefix=f"{self.name}_",
+                columns_only=True,
+            )
         elif self.type == 'Tuple':
             if not self.items or not len(self.items):
                 return {}
@@ -185,43 +148,18 @@ def parse_stack_item(self, stack_entry: StackEntry, tlb_sources, force_all: bool
                     data = to_parse.cell_unpack(item, True)
 
                 parsed_data = data.dump(with_types=tlb['dump_with_types'])
-
-                if 'parse' in tlb:
-                    data = tlb.get('parse')
-                    assert isinstance(data, list)
-
-                    tmp = {}
-                    if tlb['dump_with_types']:
-                        tmp[f'{self.dton_parse_prefix}{self.name}_type'] = parsed_data['type']
-
-                    for item in data:
-                        path = item.get('path', None)
-                        assert path, "Missing path in TLB"
-
-                        path = path.split('.')
-                        if tmp.get(path[-1]):
-                            raise ValueError(f'Duplicate path {path[-1]}')
-
-                        dtype = item.get('labels', {}).get('dton_type', None)
-                        assert dtype in supported_types, f'Unsupported ABI type {dtype}'
-
-                        name = item.get('labels', {}).get('name', path[-1])
-
-                        old = parsed_data.get(path[0], None)
-
-                        if len(path) > 1:
-                            for item in path[1:]:
-                                if old:
-                                    old = old.get(item, None)
-                        if old is not None:
-                            if dtype == 'FixedString(64)':
-                                old = hex(old).upper()[2:].zfill(64)
-
-                            tmp[f'{self.dton_parse_prefix}{self.name}_{name}'] = old
-
-                    return tmp
-
-                return {f"{self.dton_parse_prefix}{self.name}": stack_entry.get().to_boc()}
+                result = parse_tlb_spec(
+                    parsed_data,
+                    tlb,
+                    prefix=self.dton_parse_prefix,
+                    name_prefix=f"{self.name}_",
+                    columns_only=False,
+                )
+
+                if not result:
+                    return {f"{self.dton_parse_prefix}{self.name}": stack_entry.get().to_boc()}
+
+                return result
         elif self.dton_type in ['Int8', 'Int16', 'Int32', 'Int64', 'Int128', 'Int256']:
             return {
                 f"{self.dton_parse_prefix}{self.name}":

src/tonpy/abi/instance.py

@@ -55,7 +55,7 @@ def abi_for_getters(self, getters: List[int]) -> set[ABIInterfaceInstance]:
 
         return clear_parsers
 
-    def get_parsers(self, code_hash: str, getters: List[int]) -> set[ABIInterfaceInstance]:
+    def get_parsers(self, code_hash: str, getters: List[int] | None) -> set[ABIInterfaceInstance]:
         parsers = set()
 
         if code_hash in self.by_code_hash and self.by_code_hash[code_hash]:
@@ -157,3 +157,40 @@ async def aparse_getter_lazy(self, code_hash, get_tvm: Callable, getters: List[i
                 result.update()
 
         return result
+
+    def parse_storage(self, tvm: TVM) -> dict:
+        parsers = self.get_parsers(tvm.code_hash, getters=None)
+        result = {}
+        if len(parsers) > 0:
+            result['abi_interfaces'] = []
+            for parser in parsers:
+                # Always record interface if it declares storage, even if parse yields no fields
+                if getattr(parser, 'storage', None):
+                    result['abi_interfaces'].append(parser.name)
+                data = parser.parse_storage(tvm, self.tlb_sources)
+                if data:
+                    result.update(data)
+        return result
+
+    def parse_storage_lazy(self, code_hash, get_tvm: Callable) -> dict:
+        parsers = self.get_parsers(code_hash, getters=None)
+        result = {}
+        if len(parsers) > 0:
+            result['abi_interfaces'] = []
+            tvm = get_tvm()
+            for parser in parsers:
+                # Always record interface if it declares storage, even if parse yields no fields
+                if getattr(parser, 'storage', None):
+                    result['abi_interfaces'].append(parser.name)
+                data = parser.parse_storage(tvm, self.tlb_sources)
+                if data:
+                    for key, value in data.items():
+                        if key not in result:
+                            result[key] = value
+                        else:
+                            if result[key] is not None:
+                                result[key] = value
+                            else:
+                                logger.warning(f"Got multiple not null answers for storage field {key} in {parser.name}")
+            result.update()
+        return result

src/tonpy/abi/interface.py

@@ -1,9 +1,7 @@
-from typing import List
-
-from tonpy import TVM
+from tonpy import TVM, add_tlb
 from tonpy.abi.getter import ABIGetterInstance
+from tonpy.abi.utils import parse_tlb_spec
 from loguru import logger
-import asyncio
 import traceback
 
 class ABIInterfaceInstance:
@@ -18,8 +16,10 @@ def __init__(self, instance):
             self.dton_parse_prefix = self.instance['labels']['dton_parse_prefix']
 
         self.getters = []
+        self.storage = self.instance.get('storage')
 
-        for i in self.instance['get_methods']:
+        # get_methods could be absent when only storage is provided
+        for i in self.instance.get('get_methods', {}):
             for getter in self.instance['get_methods'][i]:
                 self.getters.append(ABIGetterInstance(getter))
 
@@ -29,13 +29,30 @@ def __hash__(self):
     def get_columns(self) -> dict:
         columns = {}
 
+        # getters columns
         for getter in self.getters:
             tmp = getter.get_columns()
             for c in tmp:
                 columns[f"{self.dton_parse_prefix}{c}"] = tmp[c]
 
+        # storage columns
+        columns.update(self.get_storage_columns())
+
         return columns
 
+    def get_storage_columns(self) -> dict:
+        if not self.storage:
+            return {}
+
+        tlb = self.storage
+        return parse_tlb_spec(
+            None,
+            tlb,
+            prefix=self.dton_parse_prefix,
+            name_prefix="",
+            columns_only=True,
+        )
+
     def parse_getters(self, tvm: TVM, tlb_sources):
         result = {}
 
@@ -53,6 +70,11 @@ def parse_getters(self, tvm: TVM, tlb_sources):
                 logger.warning(f"Can't parse {self.name}, (getter: {getter.method_name}): {e} {traceback.format_exc()}")
                 return {} # abi should work completely, with a result in each getter
 
+        # parse storage too
+        storage_parsed = self.parse_storage(tvm, tlb_sources)
+        if storage_parsed:
+            result.update(storage_parsed)
+
         return result
 
     async def aparse_getters(self, tvm: TVM, tlb_sources):
@@ -73,5 +95,45 @@ async def aparse_getters(self, tvm: TVM, tlb_sources):
                 logger.warning(f"Can't parse {self.name}, (getter: {getter.method_name}): {e} {traceback.format_exc()}")
                 return {} # abi should work completely, with a result in each getter
 
+        # parse storage too
+        storage_parsed = self.parse_storage(tvm, tlb_sources)
+        if storage_parsed:
+            result.update(storage_parsed)
+
         return result
 
+    def parse_storage(self, tvm: TVM, tlb_sources):
+        if not self.storage:
+            return {}
+
+        try:
+            tlb = self.storage
+            # resolve TLB text
+            tlb_text = None
+            if 'id' in tlb and tlb['id'] in tlb_sources:
+                tlb_text = tlb_sources[tlb['id']]['tlb']
+            elif 'file_path' in tlb and tlb['file_path'] in tlb_sources:
+                tlb_text = tlb_sources[tlb['file_path']]['tlb']
+
+            if tlb.get('use_block_tlb'):
+                tlb_text = f"{tlb_sources['block_tlb']}\n\n{tlb_text}"
+
+            env = {}
+            add_tlb(tlb_text, env)
+            to_parse = env[tlb['object']]()
+
+            # Unpack the contract storage cell
+            data = to_parse.cell_unpack(tvm.data, True)
+            dumped = data.dump(with_types=tlb.get('dump_with_types', False))
+
+            return parse_tlb_spec(
+                dumped,
+                tlb,
+                prefix=self.dton_parse_prefix,
+                name_prefix="",
+                columns_only=False,
+            )
+        except Exception as e:
+            logger.warning(f"Can't parse storage for {self.name}: {e} {traceback.format_exc()}")
+            return {}
+