From 27c5a7580bed29fd8c53662453b49cbc1e3f0992 Mon Sep 17 00:00:00 2001
From: Thales Lima <thales@datarevenue.com>
Date: Thu, 11 Jul 2024 19:14:32 +0200
Subject: [PATCH] Load tokens from Tycho RPC

---
 tycho/tycho/exceptions.py    |  8 ++++
 tycho/tycho/models.py        |  9 +++++
 tycho/tycho/tycho_adapter.py | 76 ++++++++++++++++++++++++++++++++++--
 tycho/tycho/utils.py         |  3 ++
 4 files changed, 93 insertions(+), 3 deletions(-)
 create mode 100644 tycho/tycho/utils.py

diff --git a/tycho/tycho/exceptions.py b/tycho/tycho/exceptions.py
index e69de29..8d88057 100644
--- a/tycho/tycho/exceptions.py
+++ b/tycho/tycho/exceptions.py
@@ -0,0 +1,8 @@
+class TychoDecodeError(Exception):
+    def __init__(self, msg: str, pool_id: str):
+        super().__init__(msg)
+        self.pool_id = pool_id
+
+
+class APIRequestError(Exception):
+    pass
diff --git a/tycho/tycho/models.py b/tycho/tycho/models.py
index a71deb8..4b7529d 100644
--- a/tycho/tycho/models.py
+++ b/tycho/tycho/models.py
@@ -1,5 +1,7 @@
 import datetime
 from enum import Enum
+from typing import Union
+
 from pydantic import BaseModel, Field
 
 
@@ -18,3 +20,10 @@ class EVMBlock(BaseModel):
 
 class ThirdPartyPool:
     pass
+
+
+class EthereumToken(BaseModel):
+    symbol: str
+    address: str
+    decimals: int
+    gas: Union[int, list[int]] = 29000
diff --git a/tycho/tycho/tycho_adapter.py b/tycho/tycho/tycho_adapter.py
index 61d62ac..b90f337 100644
--- a/tycho/tycho/tycho_adapter.py
+++ b/tycho/tycho/tycho_adapter.py
@@ -1,13 +1,16 @@
 import asyncio
 import json
 import platform
+import requests
 import time
+
 from asyncio.subprocess import STDOUT, PIPE
 from collections import defaultdict
 from datetime import datetime
 from decimal import Decimal
+from http.client import HTTPException
 from logging import getLogger
-from typing import Tuple, Any, Optional
+from typing import Tuple, Any, Optional, Dict
 
 from eth_utils import to_checksum_address
 from protosim_py import (
@@ -25,7 +28,8 @@ from tycho.tycho.constants import (
     MAX_BALANCE,
 )
 from tycho.tycho.decoders import ThirdPartyPoolTychoDecoder
-from tycho.tycho.models import Blockchain, EVMBlock, ThirdPartyPool
+from tycho.tycho.exceptions import APIRequestError
+from tycho.tycho.models import Blockchain, EVMBlock, ThirdPartyPool, EthereumToken
 
 log = getLogger(__name__)
 
@@ -34,6 +38,65 @@ class TychoClientException(Exception):
     pass
 
 
+class TokenLoader:
+    def __init__(
+        self,
+        tycho_url: str,
+        blockchain: Blockchain,
+        min_token_quality: Optional[int] = 51,
+    ):
+        self.tycho_url = tycho_url
+        self.blockchain = blockchain
+        self.min_token_quality = min_token_quality
+        self.endpoint = "/v1/{}/tokens"
+        self._token_limit = 10000
+
+    def get_tokens(self) -> dict[str, EthereumToken]:
+        """Loads all tokens from Tycho RPC"""
+        url = self.tycho_url + self.endpoint.format(self.blockchain.value)
+        page = 0
+
+        start = time.monotonic()
+        all_tokens = []
+        while data := self._get_all_with_pagination(
+            url=url,
+            page=page,
+            limit=self._token_limit,
+            params={"min_quality": self.min_token_quality},
+        ):
+            all_tokens.extend(data)
+            page += 1
+            if len(data) < self._token_limit:
+                break
+
+        log.info(f"Loaded {len(all_tokens)} tokens in {time.monotonic() - start:.2f}s")
+
+        formatted_tokens = dict()
+
+        for token in all_tokens:
+            token["address"] = to_checksum_address(token["address"])
+            formatted = EthereumToken(**token)
+            formatted_tokens[formatted.address] = formatted
+
+        return formatted_tokens
+
+    @staticmethod
+    def _get_all_with_pagination(
+        url: str, params: Optional[Dict] = None, page: int = 0, limit: int = 50
+    ) -> Dict:
+        if params is None:
+            params = {}
+
+        params["pagination"] = {"page": page, "page_size": limit}
+        r = requests.post(url, json=params)
+        try:
+            r.raise_for_status()
+        except HTTPException as e:
+            log.error(f"Request status {r.status_code} with content {r.json()}")
+            raise APIRequestError("Failed to load token configurations")
+        return r.json()["tokens"]
+
+
 class TychoPoolStateStreamAdapter:
     def __init__(
         self,
@@ -56,7 +119,7 @@ class TychoPoolStateStreamAdapter:
         self.tycho_url = tycho_url
         self.min_tvl = min_tvl
         self.tycho_client = None
-        self.protocol = "vm:protocol"
+        self.protocol = f"vm:{protocol}"
         self._include_state = include_state
         self._blockchain = blockchain
 
@@ -70,6 +133,13 @@ class TychoPoolStateStreamAdapter:
             permanent_storage=None,
         )
 
+        # Loads tokens from Tycho
+        self._tokens: dict[str, EthereumToken] = TokenLoader(
+            tycho_url=self.tycho_url,
+            blockchain=self._blockchain,
+            min_token_quality=self.min_token_quality,
+        ).get_tokens()
+
         # TODO: Check if it's necessary
         self.ignored_pools = []
         self.vm_contracts = defaultdict(list)
diff --git a/tycho/tycho/utils.py b/tycho/tycho/utils.py
new file mode 100644
index 0000000..72b4ff5
--- /dev/null
+++ b/tycho/tycho/utils.py
@@ -0,0 +1,3 @@
+def decode_tycho_exchange(exchange: str) -> (str, bool):
+    # removes vm prefix if present, returns True if vm prefix was present (vm protocol) or False if native protocol
+    return (exchange.split(":")[1], False) if "vm:" in exchange else (exchange, True)