Merge pull request #37 from BalancerMaxis/corepool-model

add core pools model

bal_tools/models.py

@@ -1,4 +1,4 @@
-from typing import Optional, List, Tuple
+from typing import Optional, List, Tuple, Dict, NewType
 from decimal import Decimal
 from dataclasses import dataclass, fields
 from enum import Enum
@@ -20,6 +20,30 @@ class GqlChain(Enum):
 
 DateRange = Tuple[int, int]
 
+PoolId = NewType("PoolId", str)
+Symbol = NewType("Symbol", str)
+
+class CorePools(BaseModel):
+    pools: Dict[PoolId, Symbol] = Field(default_factory=dict)
+
+    @field_validator('pools')
+    @classmethod
+    def validate_pools(cls, v):
+        return {str(k): str(v) for k, v in v.items()}
+
+    def __getitem__(self, key: str) -> str:
+        return self.pools[str(key)]
+
+    def __getattr__(self, key: str) -> str:
+        return self.pools[key]
+
+    def __iter__(self):
+        return iter(self.pools.items())
+
+    def __len__(self):
+        return len(self.pools)
+
+
 
 @dataclass
 class TWAPResult:
@@ -129,4 +153,4 @@ def validate_model(cls, values):
         for field in ["protocolFee", "swapFees", "swapVolume", "liquidity"]:
             if field in values:
                 values[field] = cls.str_to_decimal(values[field])
-        return values
+        return values

bal_tools/pools_gauges.py

@@ -1,6 +1,7 @@
 from typing import Dict
 import requests
 from .utils import to_checksum_address
+from .models import CorePools, PoolId, Symbol
 
 from gql.transport.exceptions import TransportQueryError
 from bal_tools.subgraph import Subgraph
@@ -22,11 +23,8 @@ def __init__(self, chain="mainnet", use_cached_core_pools=True):
             "apiv3", "vebal_get_voting_list"
         )["veBalGetVotingList"]
         if use_cached_core_pools:
-            self.core_pools = (
-                requests.get(f"{GITHUB_RAW_OUTPUTS}/core_pools.json")
-                .json()
-                .get(chain, {})
-            )
+            core_pools_data = requests.get(f"{GITHUB_RAW_OUTPUTS}/core_pools.json").json()
+            self.core_pools = CorePools(pools=core_pools_data.get(self.chain, {}))
         else:
             self.core_pools = self.build_core_pools()
 
@@ -89,7 +87,7 @@ def is_core_pool(self, pool_id: str) -> bool:
         returns:
         True if the pool is a core pool
         """
-        return pool_id in self.core_pools
+        return pool_id in self.core_pools.pools
 
     def query_preferential_gauges(self, skip=0, step_size=100) -> list:
         """
@@ -204,7 +202,7 @@ def has_alive_preferential_gauge(self, pool_id: str) -> bool:
                 return True
         print(f"Pool {pool_id} on {self.chain} has no alive preferential gauge")
 
-    def build_core_pools(self):
+    def build_core_pools(self) -> CorePools:
         """
         build the core pools dictionary by taking pools from `get_pools_with_rate_provider` and:
         - confirm the pool has an active gauge on the vebal voting list
@@ -213,7 +211,7 @@ def build_core_pools(self):
         - remove pools from blacklist
 
         returns:
-        dictionary of the format {pool_id: symbol}
+        CorePools object containing the core pools for the current chain
         """
         core_pools = self.get_liquid_pools_with_protocol_yield_fee()
 
@@ -254,4 +252,4 @@ def build_core_pools(self):
             # no results for this chain
             pass
 
-        return core_pools
+        return CorePools(pools={PoolId(k): Symbol(v) for k, v in core_pools.items()})

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = "0.1.5"
+VERSION = "0.1.8"
 DESCRIPTION = "Balancer Tools"
 LONG_DESCRIPTION = "Balancer Maxi helper and ecosystem tools"
 

tests/test_pools_gauges.py

@@ -1,5 +1,7 @@
 import pytest
 from gql.transport.exceptions import TransportQueryError
+from bal_tools.models import CorePools
+import json
 
 
 EXAMPLE_PREFERENTIAL_GAUGES = {
@@ -47,18 +49,18 @@ def test_has_alive_preferential_gauge(bal_pools_gauges):
     assert bal_pools_gauges.has_alive_preferential_gauge(example)
 
 
-def test_core_pools_dict(bal_pools_gauges):
+def test_core_pools_type(bal_pools_gauges):
     """
-    confirm we get a dict back with entries
+    confirm we get a CorePools type with entries
     """
     core_pools = bal_pools_gauges.core_pools
-    assert isinstance(core_pools, dict)
+    assert isinstance(core_pools, CorePools)
 
 
 @pytest.mark.skip(reason="core pool list can change; examples may not be valid")
 def test_core_pools_attr(bal_pools_gauges):
     """
-    confirm example core pool is in the dict
+    confirm example core pool is in CorePools
     """
     core_pools = bal_pools_gauges.core_pools
 
@@ -109,10 +111,10 @@ def test_is_pool_exempt_from_yield_fee(bal_pools_gauges):
 
 def test_build_core_pools(bal_pools_gauges):
     """
-    confirm core_pools can be built and is a dict
+    confirm core_pools can be built and is a CorePools type
     """
     try:
-        assert isinstance(bal_pools_gauges.build_core_pools(), dict)
+        assert isinstance(bal_pools_gauges.build_core_pools(), CorePools)
     except TransportQueryError as e:
         if "Too Many Requests" in str(e):
             pytest.skip(f"Skipping {bal_pools_gauges.chain}, too many requests")

tests/test_subgraph.py

@@ -45,7 +45,7 @@ def test_get_twap_price_token(subgraph, date_range):
         date_range=date_range,
     )
     assert isinstance(res.twap_price, Decimal)
-    assert pytest.approx(res.twap_price, rel=Decimal(1e-2)) == Decimal(3743.80)
+    assert pytest.approx(res.twap_price, rel=Decimal(0.05)) == Decimal(3743.80)
 
 
 def test_get_twap_prices(subgraph, date_range):
@@ -55,7 +55,7 @@ def test_get_twap_prices(subgraph, date_range):
         date_range=date_range,
     )
     assert isinstance(prices.bpt_price, Decimal)
-    assert pytest.approx(prices.bpt_price, rel=Decimal(1e-2)) == Decimal(4149.46)
+    assert pytest.approx(prices.bpt_price, rel=Decimal(0.05)) == Decimal(4149.46)
     assert all(isinstance(price.twap_price, Decimal) for price in prices.token_prices)
 
 
@@ -68,7 +68,7 @@ def test_get_twap_prices_custom_price_logic(subgraph, date_range, web3):
         block=20059322,
     )
     assert isinstance(prices.bpt_price, Decimal)
-    assert pytest.approx(prices.bpt_price, rel=Decimal(1e-2)) == Decimal(3707.99)
+    assert pytest.approx(prices.bpt_price, rel=Decimal(0.05)) == Decimal(3707.99)
     assert all(isinstance(price.twap_price, Decimal) for price in prices.token_prices)