pbc-client/pbc_client/__init__.py

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

import requests
from frozendict import frozendict

logger = logging.getLogger(__name__)


HOSTNAME_MAINNET = 'reader.partisiablockchain.com'
HOSTNAME_TESTNET = 'node1.testnet.partisiablockchain.com'

URL_ACCOUNT_PLUGIN = 'https://{hostname}/{shard}blockchain/accountPlugin/local'
URL_ACCOUNT_PLUGIN_GLOBAL = 'https://{hostname}/{shard}blockchain/accountPlugin/global'
URL_CONTRACT_STATE = 'https://{hostname}/{shard}blockchain/contracts/{address}?requireContractState=false'


MPC_DECIMALS = 10000


def shard_id_for_address(address: str) -> str:
    # Very rough implementation
    if address is None:
        msg = 'Address must not be None'
        raise TypeError(msg)
    if address.endswith('a'):
        return 'shards/Shard0/'
    if address.endswith('2'):
        return 'shards/Shard1/'
    return 'shards/Shard2/'


@dataclasses.dataclass(frozen=True)
class Balances:
    update_time: datetime.datetime
    mpc: Decimal
    byoc: Mapping[str, Decimal]


@dataclasses.dataclass(frozen=True)
class PbcClient:
    session: requests.Session
    hostname: str = HOSTNAME_MAINNET

    def get_json(
        self,
        url: str,
        data: Mapping[str, Any] = frozendict(),
        method='POST',
    ) -> tuple[Any, datetime.datetime]:
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
        }

        response = self.session.request(
            method,
            url,
            headers=headers,
            data=json.dumps(data),
        )
        response.raise_for_status()
        date_text = response.headers.get('last-modified') or response.headers.get(
            'date',
        )
        date = email.utils.parsedate_to_datetime(date_text)
        json_data = response.json()
        if json_data is None:
            msg = 'No result data for ' + url
            raise Exception(msg)
        return (json_data, date)

    def determine_coins(self) -> list[dict[str, Any]]:
        data: dict = {'path': []}

        url: str = URL_ACCOUNT_PLUGIN_GLOBAL.format(
            hostname=self.hostname,
            shard='',
        )

        json_data, date = self.get_json(url, data=data)
        return json_data['coins']['coins']

    def get_account_balances(self, address: str) -> Balances:
        coins = self.determine_coins()

        url = URL_ACCOUNT_PLUGIN.format(
            hostname=self.hostname,
            shard=shard_id_for_address(address),
        )

        data: dict = {
            'path': [
                {'type': 'field', 'name': 'accounts'},
                {'type': 'avl', 'keyType': 'BLOCKCHAIN_ADDRESS', 'key': address},
            ],
        }
        account_data, date = self.get_json(url, data=data)

        byoc: dict[str, Decimal] = {}
        mpc = Decimal(account_data['mpcTokens']) / MPC_DECIMALS

        for coin_idx, amount_data in enumerate(account_data['accountCoins']):
            coin_data = coins[coin_idx]
            byoc_balance = Decimal(amount_data['balance'])
            denominator = Decimal(coin_data['conversionRate']['denominator'])
            native_balance = byoc_balance / denominator
            byoc[coin_data['symbol']] = native_balance
            del coin_idx, coin_data

        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),
            address=address,
        )
        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)