Implement a scaffold for running arbitrage "strategies."

src/contracts/auction.py

@@ -1,6 +1,7 @@
 from src.util import NEUTRON_NETWORK_CONFIG
 from cosmpy.aerial.contract import LedgerContract
 from cosmpy.aerial.client import LedgerClient
+from typing import Any
 
 
 class AuctionProvider:
@@ -13,7 +14,7 @@ def __init__(self, contract: LedgerContract, asset_a: str, asset_b: str):
         self.asset_a_denom = asset_a
         self.asset_b_denom = asset_b
 
-    def simulate_swap_asset_a(self, amount: int) -> int:
+    def simulate_swap_asset_a(self, amount: int) -> float:
         auction_info = self.contract.query("get_auction")
 
         print(self.contract.address)
@@ -38,7 +39,7 @@ class AuctionDirectory:
     - AuctionProviders for each auction
     """
 
-    def __init__(self, deployments: dict[str, any]):
+    def __init__(self, deployments: dict[str, Any]):
         self.client = LedgerClient(NEUTRON_NETWORK_CONFIG)
         self.deployment_info = deployments["auctions"]["neutron"]
 
@@ -55,7 +56,7 @@ def auctions(self) -> dict[str, dict[str, AuctionProvider]]:
         auction_infos = self.directory_contract.query(
             {"get_pairs": {"start_after": None, "limit": None}}
         )
-        auctions = {}
+        auctions: dict[str, dict[str, AuctionProvider]] = {}
 
         for auction in auction_infos:
             pair, addr = auction

src/contracts/pool/astroport.py

@@ -2,6 +2,7 @@
 from src.util import NEUTRON_NETWORK_CONFIG
 from cosmpy.aerial.contract import LedgerContract
 from cosmpy.aerial.client import LedgerClient
+from typing import Any, cast
 
 
 class Token:
@@ -14,16 +15,16 @@ def __init__(self, denom: str):
         self.denom = denom
 
 
-def asset_info_to_token(info: dict[str, any]) -> NativeToken | Token:
+def asset_info_to_token(info: dict[str, Any]) -> NativeToken | Token:
     """
     Converts an Astroport Pairs {} query message response list member to a
     representation as Token or NativeToken.
     """
 
     if "token" in info:
-        return Token(info["token"]["contract_addr"])
+        return Token(cast(dict[str, Any], info["token"])["contract_addr"])
 
-    return NativeToken(info["native_token"]["denom"])
+    return NativeToken(cast(dict[str, Any], info["native_token"])["denom"])
 
 
 def token_to_addr(token: NativeToken | Token) -> str:
@@ -37,7 +38,7 @@ def token_to_addr(token: NativeToken | Token) -> str:
     return token.contract_addr
 
 
-def token_to_asset_info(token: NativeToken | Token) -> dict[str, any]:
+def token_to_asset_info(token: NativeToken | Token) -> dict[str, Any]:
     """
     Gets the JSON astroport AssetInfo representation of a token representation.
     """
@@ -100,7 +101,7 @@ class AstroportPoolDirectory:
     - AstroportPoolProviders for each pair
     """
 
-    def __init__(self, deployments: dict[str, any]):
+    def __init__(self, deployments: dict[str, Any]):
         self.client = LedgerClient(NEUTRON_NETWORK_CONFIG)
         self.deployment_info = deployments["pools"]["astroport"]["neutron"]
 
@@ -119,7 +120,7 @@ def pools(self) -> dict[str, dict[str, AstroportPoolProvider]]:
         )["pairs"]
 
         # All denom symbols and token contract addresses
-        asset_pools = {}
+        asset_pools: dict[str, dict[str, AstroportPoolProvider]] = {}
 
         # Pool wrappers for each asset
         for pool in pools:

src/contracts/pool/osmosis.py

@@ -1,6 +1,7 @@
 import urllib3
 import json
 from src.contracts.pool.provider import PoolProvider
+from typing import Any
 
 
 class OsmosisPoolProvider(PoolProvider):
@@ -15,8 +16,8 @@ def __init__(self, pool_id: int, asset_a: str, asset_b: str):
         self.asset_b_denom = asset_b
         self.pool_id = pool_id
 
-    def __exchange_rate(self, asset_a: str, asset_b: str, amount: int) -> int:
-        return int(
+    def __exchange_rate(self, asset_a: str, asset_b: str, amount: int) -> float:
+        return float(
             json.loads(
                 self.client.request(
                     "GET",
@@ -25,10 +26,10 @@ def __exchange_rate(self, asset_a: str, asset_b: str, amount: int) -> int:
             )["token_out_amount"]
         )
 
-    def simulate_swap_asset_a(self, amount: int) -> int:
+    def simulate_swap_asset_a(self, amount: int) -> float:
         return self.__exchange_rate(self.asset_a_denom, self.asset_b_denom, amount)
 
-    def simulate_swap_asset_b(self, amount: int) -> int:
+    def simulate_swap_asset_b(self, amount: int) -> float:
         return self.__exchange_rate(self.asset_b_denom, self.asset_a_denom, amount)
 
     def asset_a(self) -> str:
@@ -54,7 +55,7 @@ def pools(self) -> dict[str, dict[str, OsmosisPoolProvider]]:
         Gets an OsmosisPoolProvider for every pair on Osmosis.
         """
 
-        def denoms(pool: dict[str, any]) -> list[str]:
+        def denoms(pool: dict[str, Any]) -> list[str]:
             if "pool_liquidity" in pool:
                 return [asset["denom"] for asset in pool["pool_liquidity"]][:2]
 
@@ -74,7 +75,7 @@ def denoms(pool: dict[str, any]) -> list[str]:
         )["pools"]
 
         # Match each symbol with multiple trading pairs
-        asset_pools = {}
+        asset_pools: dict[str, dict[str, OsmosisPoolProvider]] = {}
 
         for pool_id, pool in enumerate(pools, 1):
             denom_addrs = denoms(pool)

src/contracts/pool/provider.py

@@ -5,30 +5,30 @@ class PoolProvider:
     exchange and chain-specific via extension of this base class.
     """
 
-    def simulate_swap_asset_a(self, amount: int) -> int:
+    def simulate_swap_asset_a(self, amount: int) -> float:
         """
         Gets the current exchange rate per quantity of asset a in the base denomination.
         """
 
-        pass
+        raise NotImplementedError
 
-    def simulate_swap_asset_b(self, amount: int) -> int:
+    def simulate_swap_asset_b(self, amount: int) -> float:
         """
         Gets the current exchange rate per quantity of asset b in the base denomination.
         """
 
-        pass
+        raise NotImplementedError
 
     def asset_a(self) -> str:
         """
         Gets the contract address or ticker (if a native asset) of the first denomination in the pair.
         """
 
-        pass
+        raise NotImplementedError
 
     def asset_b(self) -> str:
         """
         Gets the contract address or ticker (if a native asset) of the second denomination in the pair.
         """
 
-        pass
+        raise NotImplementedError

src/scheduler.py

@@ -0,0 +1,57 @@
+from src.contracts.auction import AuctionDirectory, AuctionProvider
+from src.contracts.pool.provider import PoolProvider
+from src.contracts.pool.astroport import AstroportPoolDirectory
+from src.contracts.pool.osmosis import OsmosisPoolDirectory
+from src.util import deployments
+from typing import Callable, List
+
+
+class Scheduler:
+    """
+    A registry of pricing providers for different assets,
+    which can be polled alongside a strategy function which
+    may interact with registered providers.
+    """
+
+    def __init__(
+        self,
+        strategy: Callable[
+            [
+                dict[str, dict[str, List[PoolProvider]]],
+                dict[str, dict[str, AuctionProvider]],
+            ],
+            None,
+        ],
+    ):
+        self.strategy = strategy
+        self.providers: dict[str, dict[str, List[PoolProvider]]] = {}
+
+        auction_manager = AuctionDirectory(deployments())
+        self.auctions = auction_manager.auctions()
+
+    def register_provider(self, provider: PoolProvider):
+        """
+        Registers a pool provider, enqueing it to future strategy function polls.
+        """
+
+        if provider.asset_a() not in self.providers:
+            self.providers[provider.asset_a()] = {}
+
+        if provider.asset_b() not in self.providers:
+            self.providers[provider.asset_b()] = {}
+
+        if provider.asset_b() not in self.providers[provider.asset_a()]:
+            self.providers[provider.asset_a()][provider.asset_b()] = []
+
+        if provider.asset_a() not in self.providers[provider.asset_b()]:
+            self.providers[provider.asset_b()][provider.asset_a()] = []
+
+        self.providers[provider.asset_a()][provider.asset_b()].append(provider)
+        self.providers[provider.asset_b()][provider.asset_a()].append(provider)
+
+    def poll(self):
+        """
+        Polls the strategy functionw with all registered providers.
+        """
+
+        self.strategy(self.providers, self.auctions)

src/util.py

@@ -1,4 +1,6 @@
 from cosmpy.aerial.client import NetworkConfig
+import json
+from typing import Any
 
 """
 Network configuration to be used for all neutron clients.
@@ -10,3 +12,8 @@
     fee_denomination="untrn",
     staking_denomination="untrn",
 )
+
+
+def deployments() -> dict[str, Any]:
+    with open("contracts/deployments.json") as f:
+        return json.load(f)

tests/test_auction.py

@@ -2,7 +2,7 @@
 from src.contracts.auction import (
     AuctionDirectory,
 )
-from tests.util import deployments
+from src.util import deployments
 import json
 import pytest
 

tests/test_pool.py

@@ -5,7 +5,7 @@
     token_to_addr,
 )
 from src.contracts.pool.osmosis import OsmosisPoolDirectory, OsmosisPoolProvider
-from tests.util import deployments
+from src.util import deployments
 import pytest
 
 
@@ -16,7 +16,13 @@ def test_astroport_pools():
     """
 
     astroport = AstroportPoolDirectory(deployments())
-    assert len(astroport.pools()) != 0
+    pools = astroport.pools()
+
+    assert len(pools) != 0
+
+    for base in pools.values():
+        for pool in base.values():
+            assert isinstance(pool, AstroportPoolProvider)
 
 
 def test_osmosis_pools():
@@ -26,7 +32,13 @@ def test_osmosis_pools():
     """
 
     osmosis = OsmosisPoolDirectory()
-    assert len(osmosis.pools()) != 0
+    pools = osmosis.pools()
+
+    assert len(pools) != 0
+
+    for base in pools.values():
+        for pool in base.values():
+            assert isinstance(pool, OsmosisPoolProvider)
 
 
 def test_astroport_provider():

tests/test_scheduler.py

@@ -0,0 +1,77 @@
+from src.scheduler import Scheduler
+from src.util import deployments
+from src.contracts.pool.osmosis import OsmosisPoolDirectory
+from src.contracts.pool.astroport import AstroportPoolDirectory
+from src.contracts.pool.provider import PoolProvider
+from src.contracts.auction import AuctionProvider
+import json
+import pytest
+from typing import List
+
+
+# Noop strategy
+def strategy(
+    pools: dict[str, dict[str, List[PoolProvider]]],
+    auctions: dict[str, dict[str, AuctionProvider]],
+):
+    return None
+
+
+def test_init():
+    """
+    Test that a scheduler can be instantiated.
+    """
+
+    sched = Scheduler(strategy)
+    assert sched is not None
+
+
+def test_register_provider():
+    """
+    Test that a provider can be registered to a scheduler.
+    """
+
+    osmosis = OsmosisPoolDirectory()
+    pool = list(list(osmosis.pools().values())[0].values())[0]
+
+    sched = Scheduler(strategy)
+
+    directory = OsmosisPoolDirectory()
+    pools = directory.pools()
+
+    for base in pools.values():
+        for pool in base.values():
+            sched.register_provider(pool)
+
+    assert len(sched.providers) > 0
+
+
+def test_poll():
+    """
+    Test that a strategy function can be run.
+    """
+
+    osmosis = OsmosisPoolDirectory()
+    astroport = AstroportPoolDirectory(deployments())
+
+    def simple_strategy(
+        pools: dict[str, dict[str, List[PoolProvider]]],
+        auctions: dict[str, dict[str, AuctionProvider]],
+    ):
+        assert len(pools) > 0
+        assert len(auctions) > 0
+
+    sched = Scheduler(simple_strategy)
+
+    osmos_pools = osmosis.pools()
+    astro_pools = astroport.pools()
+
+    for base in osmos_pools.values():
+        for pool in base.values():
+            sched.register_provider(pool)
+
+    for base in astro_pools.values():
+        for pool in base.values():
+            sched.register_provider(pool)
+
+    sched.poll()