diff --git a/pbc_client/__init__.py b/pbc_client/__init__.py index 4640d6b..f9b5384 100644 --- a/pbc_client/__init__.py +++ b/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)