diff --git a/bal_tools/models.py b/bal_tools/models.py index 5673e67..3349b53 100644 --- a/bal_tools/models.py +++ b/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 \ No newline at end of file + return values diff --git a/bal_tools/pools_gauges.py b/bal_tools/pools_gauges.py index 61b03b1..95effa9 100644 --- a/bal_tools/pools_gauges.py +++ b/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()}) diff --git a/setup.py b/setup.py index f2b6636..12bed54 100644 --- a/setup.py +++ b/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" diff --git a/tests/test_pools_gauges.py b/tests/test_pools_gauges.py index cfb44dc..d84d2b7 100644 --- a/tests/test_pools_gauges.py +++ b/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") diff --git a/tests/test_subgraph.py b/tests/test_subgraph.py index 6286ffb..d5f69e1 100644 --- a/tests/test_subgraph.py +++ b/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)