mm-sol 0.5.9__py3-none-any.whl → 0.6.0__py3-none-any.whl

mm_sol/rpc.py

@@ -1,232 +1,81 @@
+import json
+from collections.abc import Sequence
 from typing import Any
 
-from mm_std import Err, Ok, Result, hr
-from pydantic import BaseModel, ConfigDict, Field
+import websockets
+from mm_crypto_utils import Nodes, Proxies, retry_with_node_and_proxy
+from mm_std import Result, http_request
 
-DEFAULT_MAINNET_RPC = "https://api.mainnet-beta.solana.com"
-DEFAULT_TESTNET_RPC = "https://api.testnet.solana.com"
 
-
-class EpochInfo(BaseModel):
-    model_config = ConfigDict(populate_by_name=True)
-
-    epoch: int
-    absolute_slot: int = Field(..., alias="absoluteSlot")
-    block_height: int = Field(..., alias="blockHeight")
-    slot_index: int = Field(..., alias="slotIndex")
-    slots_in_epoch: int = Field(..., alias="slotsInEpoch")
-    transaction_count: int = Field(..., alias="transactionCount")
-
-    @property
-    def progress(self) -> float:
-        return round(self.slot_index / self.slots_in_epoch * 100, 2)
-
-
-class ClusterNode(BaseModel):
-    pubkey: str
-    version: str | None
-    gossip: str | None
-    rpc: str | None
-
-
-class VoteAccount(BaseModel):
-    class EpochCredits(BaseModel):
-        epoch: int
-        credits: int
-        previous_credits: int
-
-    validator: str
-    vote: str
-    commission: int
-    stake: int
-    credits: list[EpochCredits]
-    epoch_vote_account: bool
-    root_slot: int
-    last_vote: int
-    delinquent: bool
-
-
-class BlockProduction(BaseModel):
-    class Leader(BaseModel):
-        address: str
-        produced: int
-        skipped: int
-
-    slot: int
-    first_slot: int
-    last_slot: int
-    leaders: list[Leader]
-
-    @property
-    def total_produced(self) -> int:
-        return sum(leader.produced for leader in self.leaders)
-
-    @property
-    def total_skipped(self) -> int:
-        return sum(leader.skipped for leader in self.leaders)
-
-
-class StakeActivation(BaseModel):
-    state: str
-    active: int
-    inactive: int
-
-
-def rpc_call(
-    *,
+async def rpc_call(
     node: str,
     method: str,
-    params: list[Any],
+    params: Sequence[object],
+    timeout: float,
+    proxy: str | None,
     id_: int = 1,
-    timeout: float = 10,
-    proxy: str | None = None,
 ) -> Result[Any]:
     data = {"jsonrpc": "2.0", "method": method, "params": params, "id": id_}
     if node.startswith("http"):
-        return _http_call(node, data, timeout, proxy)
-    raise NotImplementedError("ws is not implemented")
+        return await _http_call(node, data, timeout, proxy)
+    return await _ws_call(node, data, timeout)
 
 
-def _http_call(node: str, data: dict[str, object], timeout: float, proxy: str | None) -> Result[Any]:
-    res = hr(node, method="POST", proxy=proxy, timeout=timeout, params=data, json_params=True)
+async def _http_call(node: str, data: dict[str, object], timeout: float, proxy: str | None) -> Result[Any]:
+    res = await http_request(node, method="POST", proxy=proxy, timeout=timeout, json=data)
+    if res.is_error():
+        return res.to_result_failure()
     try:
-        if res.is_error():
-            return res.to_err_result()
-
-        err = res.json.get("error", {}).get("message", "")
+        parsed_body = res.parse_json_body()
+        err = parsed_body.get("error", {}).get("message", "")
         if err:
-            return res.to_err_result(f"service_error: {err}")
-        if "result" in res.json:
-            return res.to_ok_result(res.json["result"])
-
-        return res.to_err_result("unknown_response")
+            return res.to_result_failure(f"service_error: {err}")
+        if "result" in parsed_body:
+            return res.to_result_success(parsed_body["result"])
+        return res.to_result_failure("unknown_response")
     except Exception as e:
-        return res.to_err_result(f"exception: {e!s}")
-
-
-def get_balance(node: str, address: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
-    """Returns balance in lamports"""
-    return rpc_call(node=node, method="getBalance", params=[address], timeout=timeout, proxy=proxy).and_then(lambda r: r["value"])
-
-
-def get_block_height(node: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
-    """Returns balance in lamports"""
-    return rpc_call(node=node, method="getBlockHeight", params=[], timeout=timeout, proxy=proxy)
-
-
-def get_slot(node: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
-    return rpc_call(node=node, method="getSlot", params=[], timeout=timeout, proxy=proxy)
-
-
-def get_epoch_info(node: str, epoch: int | None = None, timeout: float = 10, proxy: str | None = None) -> Result[EpochInfo]:
-    params = [epoch] if epoch else []
-    return rpc_call(node=node, method="getEpochInfo", params=params, timeout=timeout, proxy=proxy).and_then(
-        lambda r: EpochInfo(**r),
-    )
-
-
-def get_health(node: str, timeout: float = 10, proxy: str | None = None) -> Result[bool]:
-    return rpc_call(node=node, method="getHealth", params=[], timeout=timeout, proxy=proxy).and_then(lambda r: r == "ok")
-
-
-def get_cluster_nodes(node: str, timeout: float = 30, proxy: str | None = None) -> Result[list[ClusterNode]]:
-    return rpc_call(node=node, method="getClusterNodes", timeout=timeout, proxy=proxy, params=[]).and_then(
-        lambda r: [ClusterNode(**n) for n in r],
-    )
+        return res.to_result_failure(e)
 
 
-def get_vote_accounts(node: str, timeout: float = 30, proxy: str | None = None) -> Result[list[VoteAccount]]:
-    res = rpc_call(node=node, method="getVoteAccounts", timeout=timeout, proxy=proxy, params=[])
-    if res.is_err():
-        return res
+async def _ws_call(node: str, data: dict[str, object], timeout: float) -> Result[Any]:
+    response = None
     try:
-        data = res.unwrap()
-        result: list[VoteAccount] = []
-        for a in data["current"]:
-            result.append(  # noqa: PERF401
-                VoteAccount(
-                    validator=a["nodePubkey"],
-                    vote=a["votePubkey"],
-                    commission=a["commission"],
-                    stake=a["activatedStake"],
-                    credits=[
-                        VoteAccount.EpochCredits(epoch=c[0], credits=c[1], previous_credits=c[2]) for c in a["epochCredits"]
-                    ],
-                    delinquent=False,
-                    epoch_vote_account=a["epochVoteAccount"],
-                    root_slot=a["rootSlot"],
-                    last_vote=a["lastVote"],
-                ),
-            )
-        for a in data["delinquent"]:
-            result.append(  # noqa: PERF401
-                VoteAccount(
-                    validator=a["nodePubkey"],
-                    vote=a["votePubkey"],
-                    commission=a["commission"],
-                    stake=a["activatedStake"],
-                    credits=[
-                        VoteAccount.EpochCredits(epoch=c[0], credits=c[1], previous_credits=c[2]) for c in a["epochCredits"]
-                    ],
-                    delinquent=True,
-                    epoch_vote_account=a["epochVoteAccount"],
-                    root_slot=a["rootSlot"],
-                    last_vote=a["lastVote"],
-                ),
-            )
-        return Ok(result, res.data)
+        async with websockets.connect(node, open_timeout=timeout) as ws:
+            await ws.send(json.dumps(data))
+            response = json.loads(await ws.recv())
+
+        err = response.get("error", {}).get("message", "")
+        if err:
+            return Result.failure(f"service_error: {err}", {"response": response})
+        if "result" in response:
+            return Result.success(response["result"], {"response": response})
+        return Result.failure("unknown_response", {"response": response})
+    except TimeoutError:
+        return Result.failure("timeout", {"response": response})
     except Exception as e:
-        return Err(e, res.data)
+        return Result.failure(e, {"response": response})
 
 
-def get_leader_scheduler(
-    node: str,
-    slot: int | None = None,
-    timeout: float = 10,
-    proxy: str | None = None,
-) -> Result[dict[str, list[int]]]:
-    return rpc_call(
-        node=node,
-        method="getLeaderSchedule",
-        timeout=timeout,
-        proxy=proxy,
-        params=[slot],
-    )
+async def get_block_height(node: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
+    return await rpc_call(node=node, method="getBlockHeight", params=[], timeout=timeout, proxy=proxy)
 
 
-def get_block_production(node: str, timeout: float = 60, proxy: str | None = None) -> Result[BlockProduction]:
-    res = rpc_call(node=node, method="getBlockProduction", timeout=timeout, proxy=proxy, params=[])
-    if res.is_err():
-        return res
-    try:
-        res_ok = res.unwrap()
-        slot = res_ok["context"]["slot"]
-        first_slot = res_ok["value"]["range"]["firstSlot"]
-        last_slot = res_ok["value"]["range"]["lastSlot"]
-        leaders = []
-        for address, (leader, produced) in res.ok["value"]["byIdentity"].items():  # type: ignore[index]
-            leaders.append(BlockProduction.Leader(address=address, produced=produced, skipped=leader - produced))
-        return Ok(BlockProduction(slot=slot, first_slot=first_slot, last_slot=last_slot, leaders=leaders), res.data)
-    except Exception as e:
-        return Err(e, data=res.data)
-
-
-def get_stake_activation(node: str, address: str, timeout: float = 60, proxy: str | None = None) -> Result[StakeActivation]:
-    return rpc_call(node=node, method="getStakeActivation", timeout=timeout, proxy=proxy, params=[address]).and_then(
-        lambda ok: StakeActivation(**ok),
+async def get_balance(node: str, address: str, timeout: float = 5, proxy: str | None = None) -> Result[int]:
+    """Returns balance in lamports"""
+    return (await rpc_call(node=node, method="getBalance", params=[address], timeout=timeout, proxy=proxy)).map(
+        lambda r: r["value"]
     )
 
 
-def get_transaction(
-    node: str,
-    signature: str,
-    max_supported_transaction_version: int | None = None,
-    encoding: str = "json",
-    timeout: float = 60,
-    proxy: str | None = None,
-) -> Result[dict[str, object] | None]:
-    if max_supported_transaction_version is not None:
-        params = [signature, {"maxSupportedTransactionVersion": max_supported_transaction_version, "encoding": encoding}]
-    else:
-        params = [signature, encoding]
-    return rpc_call(node=node, method="getTransaction", timeout=timeout, proxy=proxy, params=params)
+async def get_balance_with_retries(
+    retries: int, nodes: Nodes, proxies: Proxies, *, address: str, timeout: float = 5
+) -> Result[int]:
+    """
+    Retry get_balance with different nodes and proxies
+    """
+    return await retry_with_node_and_proxy(
+        retries,
+        nodes,
+        proxies,
+        lambda node, proxy: get_balance(node=node, address=address, timeout=timeout, proxy=proxy),
+    )

mm_sol/rpc_sync.py

@@ -0,0 +1,232 @@
+from typing import Any
+
+import pydash
+from mm_std import Result, http_request_sync
+from pydantic import BaseModel, ConfigDict, Field
+
+DEFAULT_MAINNET_RPC = "https://api.mainnet-beta.solana.com"
+DEFAULT_TESTNET_RPC = "https://api.testnet.solana.com"
+
+
+class EpochInfo(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    epoch: int
+    absolute_slot: int = Field(..., alias="absoluteSlot")
+    block_height: int = Field(..., alias="blockHeight")
+    slot_index: int = Field(..., alias="slotIndex")
+    slots_in_epoch: int = Field(..., alias="slotsInEpoch")
+    transaction_count: int = Field(..., alias="transactionCount")
+
+    @property
+    def progress(self) -> float:
+        return round(self.slot_index / self.slots_in_epoch * 100, 2)
+
+
+class ClusterNode(BaseModel):
+    pubkey: str
+    version: str | None
+    gossip: str | None
+    rpc: str | None
+
+
+class VoteAccount(BaseModel):
+    class EpochCredits(BaseModel):
+        epoch: int
+        credits: int
+        previous_credits: int
+
+    validator: str
+    vote: str
+    commission: int
+    stake: int
+    credits: list[EpochCredits]
+    epoch_vote_account: bool
+    root_slot: int
+    last_vote: int
+    delinquent: bool
+
+
+class BlockProduction(BaseModel):
+    class Leader(BaseModel):
+        address: str
+        produced: int
+        skipped: int
+
+    slot: int
+    first_slot: int
+    last_slot: int
+    leaders: list[Leader]
+
+    @property
+    def total_produced(self) -> int:
+        return sum(leader.produced for leader in self.leaders)
+
+    @property
+    def total_skipped(self) -> int:
+        return sum(leader.skipped for leader in self.leaders)
+
+
+class StakeActivation(BaseModel):
+    state: str
+    active: int
+    inactive: int
+
+
+def rpc_call(
+    *,
+    node: str,
+    method: str,
+    params: list[Any],
+    id_: int = 1,
+    timeout: float = 10,
+    proxy: str | None = None,
+) -> Result[Any]:
+    data = {"jsonrpc": "2.0", "method": method, "params": params, "id": id_}
+    if node.startswith("http"):
+        return _http_call(node, data, timeout, proxy)
+    raise NotImplementedError("ws is not implemented")
+
+
+def _http_call(node: str, data: dict[str, object], timeout: float, proxy: str | None) -> Result[Any]:
+    res = http_request_sync(node, method="POST", proxy=proxy, timeout=timeout, json=data)
+    try:
+        if res.is_error():
+            return res.to_result_failure()
+
+        json_body = res.parse_json_body()
+        err = pydash.get(json_body, "error.message")
+        if err:
+            return res.to_result_failure(f"service_error: {err}")
+        if "result" in json_body:
+            return res.to_result_success(json_body["result"])
+
+        return res.to_result_failure("unknown_response")
+    except Exception as e:
+        return res.to_result_failure(e)
+
+
+def get_balance(node: str, address: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
+    """Returns balance in lamports"""
+    return rpc_call(node=node, method="getBalance", params=[address], timeout=timeout, proxy=proxy).map(lambda r: r["value"])
+
+
+def get_block_height(node: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
+    """Returns balance in lamports"""
+    return rpc_call(node=node, method="getBlockHeight", params=[], timeout=timeout, proxy=proxy)
+
+
+def get_slot(node: str, timeout: float = 10, proxy: str | None = None) -> Result[int]:
+    return rpc_call(node=node, method="getSlot", params=[], timeout=timeout, proxy=proxy)
+
+
+def get_epoch_info(node: str, epoch: int | None = None, timeout: float = 10, proxy: str | None = None) -> Result[EpochInfo]:
+    params = [epoch] if epoch else []
+    return rpc_call(node=node, method="getEpochInfo", params=params, timeout=timeout, proxy=proxy).map(lambda r: EpochInfo(**r))
+
+
+def get_health(node: str, timeout: float = 10, proxy: str | None = None) -> Result[bool]:
+    return rpc_call(node=node, method="getHealth", params=[], timeout=timeout, proxy=proxy).map(lambda r: r == "ok")
+
+
+def get_cluster_nodes(node: str, timeout: float = 30, proxy: str | None = None) -> Result[list[ClusterNode]]:
+    return rpc_call(node=node, method="getClusterNodes", timeout=timeout, proxy=proxy, params=[]).map(
+        lambda r: [ClusterNode(**n) for n in r],
+    )
+
+
+def get_vote_accounts(node: str, timeout: float = 30, proxy: str | None = None) -> Result[list[VoteAccount]]:
+    res = rpc_call(node=node, method="getVoteAccounts", timeout=timeout, proxy=proxy, params=[])
+    if res.is_error():
+        return res
+    try:
+        data = res.unwrap()
+        result: list[VoteAccount] = []
+        for a in data["current"]:
+            result.append(  # noqa: PERF401
+                VoteAccount(
+                    validator=a["nodePubkey"],
+                    vote=a["votePubkey"],
+                    commission=a["commission"],
+                    stake=a["activatedStake"],
+                    credits=[
+                        VoteAccount.EpochCredits(epoch=c[0], credits=c[1], previous_credits=c[2]) for c in a["epochCredits"]
+                    ],
+                    delinquent=False,
+                    epoch_vote_account=a["epochVoteAccount"],
+                    root_slot=a["rootSlot"],
+                    last_vote=a["lastVote"],
+                ),
+            )
+        for a in data["delinquent"]:
+            result.append(  # noqa: PERF401
+                VoteAccount(
+                    validator=a["nodePubkey"],
+                    vote=a["votePubkey"],
+                    commission=a["commission"],
+                    stake=a["activatedStake"],
+                    credits=[
+                        VoteAccount.EpochCredits(epoch=c[0], credits=c[1], previous_credits=c[2]) for c in a["epochCredits"]
+                    ],
+                    delinquent=True,
+                    epoch_vote_account=a["epochVoteAccount"],
+                    root_slot=a["rootSlot"],
+                    last_vote=a["lastVote"],
+                ),
+            )
+        return Result.success(result, res.extra)
+    except Exception as e:
+        return Result.failure(e, res.extra)
+
+
+def get_leader_scheduler(
+    node: str,
+    slot: int | None = None,
+    timeout: float = 10,
+    proxy: str | None = None,
+) -> Result[dict[str, list[int]]]:
+    return rpc_call(
+        node=node,
+        method="getLeaderSchedule",
+        timeout=timeout,
+        proxy=proxy,
+        params=[slot],
+    )
+
+
+def get_block_production(node: str, timeout: float = 60, proxy: str | None = None) -> Result[BlockProduction]:
+    res = rpc_call(node=node, method="getBlockProduction", timeout=timeout, proxy=proxy, params=[])
+    if res.is_error():
+        return res
+    try:
+        res_ok = res.unwrap()
+        slot = res_ok["context"]["slot"]
+        first_slot = res_ok["value"]["range"]["firstSlot"]
+        last_slot = res_ok["value"]["range"]["lastSlot"]
+        leaders = []
+        for address, (leader, produced) in res.ok["value"]["byIdentity"].items():  # type: ignore[index]
+            leaders.append(BlockProduction.Leader(address=address, produced=produced, skipped=leader - produced))
+        return Result.success(BlockProduction(slot=slot, first_slot=first_slot, last_slot=last_slot, leaders=leaders), res.extra)
+    except Exception as e:
+        return Result.failure(e, res.extra)
+
+
+def get_stake_activation(node: str, address: str, timeout: float = 60, proxy: str | None = None) -> Result[StakeActivation]:
+    return rpc_call(node=node, method="getStakeActivation", timeout=timeout, proxy=proxy, params=[address]).map(
+        lambda ok: StakeActivation(**ok),
+    )
+
+
+def get_transaction(
+    node: str,
+    signature: str,
+    max_supported_transaction_version: int | None = None,
+    encoding: str = "json",
+    timeout: float = 60,
+    proxy: str | None = None,
+) -> Result[dict[str, object] | None]:
+    if max_supported_transaction_version is not None:
+        params = [signature, {"maxSupportedTransactionVersion": max_supported_transaction_version, "encoding": encoding}]
+    else:
+        params = [signature, encoding]
+    return rpc_call(node=node, method="getTransaction", timeout=timeout, proxy=proxy, params=params)

mm_sol/spl_token.py

@@ -0,0 +1,85 @@
+from mm_crypto_utils import Nodes, Proxies, retry_with_node_and_proxy
+from mm_std import Result
+from solana.exceptions import SolanaRpcException
+from solana.rpc.core import RPCException
+from solders.solders import InvalidParamsMessage, Pubkey, get_associated_token_address
+
+from mm_sol.utils import get_async_client
+
+
+async def get_balance(
+    node: str,
+    owner: str,
+    token: str,
+    token_account: str | None = None,
+    timeout: float = 5,
+    proxy: str | None = None,
+) -> Result[int]:
+    response = None
+    try:
+        client = get_async_client(node, proxy=proxy, timeout=timeout)
+        if not token_account:
+            token_account = str(get_associated_token_address(Pubkey.from_string(owner), Pubkey.from_string(token)))
+
+        res = await client.get_token_account_balance(Pubkey.from_string(token_account))
+        response = res.to_json()
+
+        # Sometimes it not raise an error, but it returns this :(
+        if isinstance(res, InvalidParamsMessage) and "could not find account" in res.message:
+            return Result.success(0, {"response": response})
+        return Result.success(int(res.value.amount), {"response": response})
+    except RPCException as e:
+        if "could not find account" in str(e):
+            return Result.success(0, {"response": response, "rpc_exception": str(e)})
+        return Result.failure(e, {"response": response})
+    except SolanaRpcException as e:
+        return Result.failure((e.error_msg, e), {"response": response})
+    except Exception as e:
+        return Result.failure(e, {"response": response})
+
+
+async def get_balance_with_retries(
+    retries: int,
+    nodes: Nodes,
+    proxies: Proxies,
+    *,
+    owner: str,
+    token: str,
+    token_account: str | None = None,
+    timeout: float = 5,
+) -> Result[int]:
+    return await retry_with_node_and_proxy(
+        retries,
+        nodes,
+        proxies,
+        lambda node, proxy: get_balance(
+            node,
+            owner=owner,
+            token=token,
+            token_account=token_account,
+            timeout=timeout,
+            proxy=proxy,
+        ),
+    )
+
+
+async def get_decimals(node: str, token: str, timeout: float = 5, proxy: str | None = None) -> Result[int]:
+    response = None
+    try:
+        client = get_async_client(node, proxy=proxy, timeout=timeout)
+        res = await client.get_token_supply(Pubkey.from_string(token))
+        response = res.to_json()
+        return Result.success(res.value.decimals, {"response": response})
+    except Exception as e:
+        return Result.failure(e, {"response": response})
+
+
+async def get_decimals_with_retries(
+    retries: int, nodes: Nodes, proxies: Proxies, *, token: str, timeout: float = 5
+) -> Result[int]:
+    return await retry_with_node_and_proxy(
+        retries,
+        nodes,
+        proxies,
+        lambda node, proxy: get_decimals(node, token=token, proxy=proxy, timeout=timeout),
+    )