Add more contract state and ABI retrieval methods

pbc_client/__init__.py

@@ -1,8 +1,10 @@
 import dataclasses
 import datetime
+import base64
 import email.utils
 import json
 import logging
+import pbcabi
 from collections.abc import Mapping
 from decimal import Decimal
 from typing import Any
@@ -117,6 +119,7 @@ class PbcClient:
         return Balances(date, mpc, byoc)
 
     def get_contract_state(self, address: str) -> tuple[dict, datetime.datetime]:
+        # TODO: Rename to get_contract_state_json
         url = URL_CONTRACT_STATE.format(
             hostname=self.hostname,
             shard=shard_id_for_address(address),
@@ -124,3 +127,46 @@ class PbcClient:
         )
         data: dict = {'path': []}
         return self.get_json(url, data=data)
+
+    def get_typed_contract_state(self, address: str) -> tuple[dict, datetime.datetime]:
+        """
+        Only suitable for non-governance WASM contracts.
+        """
+        file_abi = self.get_contract_abi(address)
+        state, server_time = self.get_contract_state(address)
+        state_bytes = base64.b64decode(state['state']['data'])
+        state_deserialized =  file_abi.contract.read_state(state_bytes)
+
+        return state_deserialized, server_time
+
+    def get_typed_contract_avl_tree(self, address: str, avl_tree_id: pbcabi.model.AvlTreeId) -> tuple[dict, datetime.datetime]:
+        file_abi = self.get_contract_abi(address)
+        state, server_time = self.get_contract_state(address)
+        for avl_tree in state['avlTrees']:
+            if avl_tree['key'] == avl_tree_id.avl_tree_id:
+                break
+
+        data = {}
+        for key_and_value in avl_tree['value']['avlTree']:
+            key_bytes = base64.b64decode(key_and_value['key']['data']['data'])
+            value_bytes = base64.b64decode(key_and_value['value']['data'])
+
+            key = file_abi.contract.read_state(key_bytes, avl_tree_id.type_spec.type_key)
+            value = file_abi.contract.read_state(value_bytes, avl_tree_id.type_spec.type_value)
+
+            data[key] = value
+
+            # TODO: Type
+
+        return data, server_time
+
+
+    def get_contract_abi(self, address: str) -> pbcabi.model.FileAbi:
+        url = URL_CONTRACT_STATE.format(
+            hostname=self.hostname,
+            shard=shard_id_for_address(address),
+            address=address,
+        )
+        meta, _ = self.get_json(url, method='GET')
+        abi_bytes = base64.b64decode(meta['abi'])
+        return pbcabi.model.FileAbi.read_from(abi_bytes)