Merge pull request #21 from cowprotocol/marshmallow-proposal

Implemented data serialization/deserialization with marshmallow

.pylintrc

@@ -1,2 +1,2 @@
 [MASTER]
-disable=fixme,logging-fstring-interpolation,missing-module-docstring,invalid-name
+disable=fixme,logging-fstring-interpolation,missing-module-docstring,invalid-name,too-many-instance-attributes

requirements.txt

@@ -15,3 +15,4 @@ psycopg2-binary>=2.9.3
 SQLAlchemy==1.4.41
 pandas==1.5.0
 click==8.1.3
+marshmallow==3.18.0

src/missing_prices.py

@@ -8,18 +8,19 @@
 from dotenv import load_dotenv
 from dune_client.client import DuneClient
 from dune_client.query import Query
-from dune_client.types import DuneRecord
+from marshmallow import fields
 
 from duneapi.api import DuneAPI
 from duneapi.types import Address, DuneQuery, Network
 from duneapi.util import open_query
 
-DuneTokenPriceRow = tuple[str, str, str, str, int]
+from src.utils import TokenSchema, CoinSchema, CoinsSchema, Token, EthereumAddress, Coin
 
+DuneTokenPriceRow = tuple[str, str, str, Address, int]
 
-# TODO - remove the Anys here: https://github.com/cowprotocol/data-misc/issues/20
-def load_coins() -> dict[str, dict[str, Any]]:
-    """ "
+
+def load_coins() -> dict[Address, Coin]:
+    """
     Loads and returns coin dictionaries from Coin Paprika via their API.
     Excludes, inactive, new and non "token" types
     """
@@ -42,13 +43,19 @@ def load_coins() -> dict[str, dict[str, Any]]:
             # only include ethereum tokens
             try:
                 entry["address"] = contract_dict[entry["id"]].lower()
-                coin_dict[entry["address"]] = entry
+                try:
+                    # coin_dict[Address(entry["address"])] = entry
+                    coin_dict[entry["address"]] = entry
+                except ValueError:
+                    print(f"{entry['address']} is not a valid ethereum address")
+                    continue
             except KeyError:
                 missed += 1
-                # print(f"Error with {err}, excluding entry {entry}")
 
     print(f"Excluded address for {missed} entries out of {len(entries)}")
-    return coin_dict
+    return CoinsSchema(keys=EthereumAddress, values=fields.Nested(CoinSchema)).load(
+        coin_dict
+    )
 
 
 def write_results(results: list[DuneTokenPriceRow], path: str, filename: str) -> None:
@@ -57,7 +64,8 @@ def write_results(results: list[DuneTokenPriceRow], path: str, filename: str) ->
         os.makedirs(path)
     with open(os.path.join(path, filename), "w", encoding="utf-8") as file:
         for row in results:
-            file.write(str(row) + ",\n")
+            # TODO: [duneapi#68] Fix __repr__ of duneapi.types.Address
+            file.write(f"('{row[0]}', '{row[1]}', '{row[2]}', '{row[3]}', {row[4]}),\n")
         print(f"Results written to {filename}")
 
 
@@ -101,9 +109,10 @@ def as_dune_repr(self, coin_id: str) -> dict[str, Any]:
         }
 
 
-def load_tokens(dune: DuneClient) -> list[DuneRecord]:
+def load_tokens(dune: DuneClient) -> list[Token]:
     """Loads Tokens with missing prices from Dune"""
-    return dune.refresh(Query(query_id=1317238, name="Tokens with Missing Prices"))
+    results = dune.refresh(Query(query_id=1317238, name="Tokens with Missing Prices"))
+    return [TokenSchema().load(r) for r in results]
 
 
 def fetch_tokens_without_prices(dune: DuneAPI) -> list[CoinPaprikaToken]:
@@ -128,14 +137,14 @@ def run_missing_prices() -> None:
     print(f"Fetched {len(tokens)} traded tokens from Dune without prices")
     found, res = 0, []
     for token in tokens:
-        if token["address"].lower() in coins:
-            paprika_data = coins[token["address"].lower()]
+        if token.address in coins:
+            paprika_data = coins[token.address]
             dune_row = (
-                str(paprika_data["id"]),
+                paprika_data.id,
                 "ethereum",
-                str(paprika_data["symbol"]),
-                str(paprika_data["address"].lower()),
-                int(token["decimals"]),
+                paprika_data.symbol,
+                paprika_data.address,
+                token.decimals,
             )
             res.append(dune_row)
             found += 1

src/utils.py

@@ -4,8 +4,11 @@
 from datetime import datetime
 from enum import Enum
 from typing import Any
+from dataclasses import dataclass
 
 from duneapi.types import Network as LegacyDuneNetwork
+from duneapi.types import Address
+from marshmallow import fields, Schema, post_load, ValidationError
 
 
 def partition_array(arr: list[Any], size: int) -> list[list[Any]]:
@@ -79,3 +82,105 @@ def chain_id(self) -> int:
         Aligned with https://chainlist.org/
         """
         return {Network.MAINNET: 1, Network.GNOSIS: 100}[self]
+
+
+class EthereumAddress(fields.Field):
+    """Field that serializes to a string of numbers and deserializes
+    to a list of numbers.
+    """
+
+    def _serialize(self, value, attr, obj, **kwargs):
+        if value is None:
+            return ""
+        return f"{value}".lower()
+
+    def _deserialize(self, value, attr, data, **kwargs):
+        try:
+            return Address(value)
+        except ValueError as error:
+            raise ValidationError("Not a valid address") from error
+
+
+@dataclass
+class Token:
+    """Dataclass for holding Token data"""
+
+    address: Address
+    decimals: int
+    symbol: str
+    popularity: int
+
+
+@dataclass
+class Coin:
+    """Dataclass for holding Coin data"""
+
+    id: str
+    name: str
+    symbol: str
+    rank: int
+    is_new: bool
+    is_active: bool
+    type: str
+    address: Address
+
+
+class TokenSchema(Schema):
+    """TokenSchema CoinSchema for serializing/deserializing token data"""
+
+    address = EthereumAddress(required=True)
+    decimals = fields.Int(required=True)
+    popularity = fields.Int()
+    symbol = fields.String()
+
+    @post_load
+    def make_token(self, data, **_kwargs):
+        """Turns Token data into Token instance"""
+        return Token(**data)
+
+
+class CoinSchema(Schema):
+    """CoinSchema for serializing/deserializing coin data"""
+
+    # pylint: disable=too-many-instance-attributes
+    # Eight are passed from the API
+
+    id = fields.String(required=True)
+    name = fields.String()
+    symbol = fields.String(required=True)
+    rank = fields.Int()
+    is_new = fields.Bool()
+    is_active = fields.Bool()
+    type = fields.String()
+    address = EthereumAddress(required=True)
+
+    def load(self, *args, **kwargs):
+        try:
+            return super().load(*args, **kwargs)
+        except ValidationError as e:
+            return e.valid_data
+
+    @post_load
+    def make_coin(self, data, **_kwargs):
+        """Turns Coin data into Coin instance"""
+        return Coin(**data)
+
+
+class CoinsSchema(fields.Dict):
+    """CoinsSchema for containing multiple Coinschema-s"""
+
+    @staticmethod
+    def _get_obj(obj, _attr, _default):
+        """Accessor for the dump method"""
+        return obj
+
+    def dump(self, obj: Any):
+        """Serializes data"""
+        return self.serialize("", obj, accessor=self._get_obj)
+
+    def load(self, data: dict[Address, Any]):
+        """Loads data into mapping"""
+        try:
+            return self.deserialize(data)
+        except ValidationError as e:
+            return e.valid_data