Get wallet state from db

conftest.py

@@ -6,7 +6,7 @@
 
 import pytest
 from agent0.test_fixtures import cycle_trade_policy
-from chainsync.test_fixtures import database_engine, db_session, dummy_session, psql_docker
+from chainsync.test_fixtures import database_engine, db_api, db_session, dummy_session, psql_docker
 from ethpy.test_fixtures import local_chain, local_hyperdrive_chain
 
 # Hack to allow for vscode debugger to throw exception immediately
@@ -49,6 +49,7 @@ def pytest_internalerror(excinfo):
 # TODO this means pytest can only be ran from this directory
 __all__ = [
     "database_engine",
+    "db_api",
     "db_session",
     "dummy_session",
     "psql_docker",

lib/agent0/agent0/base/__init__.py

@@ -0,0 +1,2 @@
+"""The agent0 base module exports some types here"""
+from .types import FrozenClass, Quantity, TokenType, freezable

lib/agent0/agent0/base/agents/__init__.py

@@ -1,3 +1,2 @@
 """Base implementation of agents."""
 from .eth_agent import EthAgent
-from .eth_wallet import EthWallet

lib/agent0/agent0/base/agents/eth_agent.py

@@ -3,15 +3,15 @@
 
 from typing import Generic, TypeVar
 
+from agent0.base import Quantity, TokenType
 from agent0.base.policies import BasePolicy, NoActionPolicy
-from elfpy.types import Quantity, TokenType, Trade
+from agent0.base.state import EthWallet
+from elfpy.types import Trade
 from eth_account.signers.local import LocalAccount
 from eth_typing import ChecksumAddress
 from hexbytes import HexBytes
 from web3 import Web3
 
-from .eth_wallet import EthWallet
-
 Policy = TypeVar("Policy", bound=BasePolicy)
 Market = TypeVar("Market")
 MarketAction = TypeVar("MarketAction")

lib/agent0/agent0/base/config/environment_config.py

@@ -4,15 +4,16 @@
 import json
 from dataclasses import dataclass
 
-from elfpy import DEFAULT_LOG_LEVEL, DEFAULT_LOG_MAXBYTES, types
+from agent0.base import FrozenClass, freezable
+from elfpy import DEFAULT_LOG_LEVEL, DEFAULT_LOG_MAXBYTES
 from elfpy.utils import json as output_utils
 
 DEFAULT_USERNAME = "changeme"
 
 
-@types.freezable(frozen=False, no_new_attribs=True)
+@freezable(frozen=False, no_new_attribs=True)
 @dataclass
-class EnvironmentConfig(types.FrozenClass):
+class EnvironmentConfig(FrozenClass):
     """Parameters that can be set either locally or passed from docker."""
 
     # lots of configs!
@@ -37,8 +38,8 @@ class EnvironmentConfig(types.FrozenClass):
     max_bytes: int = DEFAULT_LOG_MAXBYTES  # int(2e6) or 2MB
     # int to be used for the random seed
     random_seed: int = 1
-    # username registration URI
-    username_register_uri: str = "http://localhost:5002"
+    # Database api for username registration and wallet queries
+    database_api_uri: str = "http://localhost:5002"
 
     def __getitem__(self, attrib) -> None:
         return getattr(self, attrib)

lib/agent0/agent0/base/policies/base.py

@@ -8,7 +8,7 @@
 from numpy.random import default_rng
 
 if TYPE_CHECKING:
-    from agent0.base.agents import EthWallet
+    from agent0.base.state import EthWallet
 
     # from agent0.base.state import BaseMarketState # TODO: don't rely on elfpy base market
     from elfpy.markets.base import BaseMarket as BaseMarketState

lib/agent0/agent0/base/state/__init__.py

@@ -1,2 +1,3 @@
 """Base classes for interface objects"""
+from .eth_wallet import EthWallet, EthWalletDeltas
 from .market_state import BaseMarketState

lib/agent0/agent0/base/agents/eth_wallet.py → lib/agent0/agent0/base/state/eth_wallet.py

@@ -6,13 +6,32 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+from agent0.base import Quantity, TokenType, freezable
 from elfpy import check_non_zero
-from elfpy.types import Quantity, TokenType
-from elfpy.wallet.wallet_deltas import WalletDeltas
 from fixedpointmath import FixedPoint
 from hexbytes import HexBytes
 
 
+@freezable()
+@dataclass()
+class EthWalletDeltas:
+    r"""Stores changes for an agent's wallet
+
+    Arguments
+    ----------
+    balance : Quantity
+        The base assets that held by the trader.
+    """
+    # fungible
+    balance: Quantity = field(default_factory=lambda: Quantity(amount=FixedPoint(0), unit=TokenType.BASE))
+
+    # TODO: Support multiple typed balances:
+    #     balance: Dict[TokenType, Quantity] = field(default_factory=dict)
+    def copy(self) -> EthWalletDeltas:
+        """Returns a new copy of self"""
+        return EthWalletDeltas(**copy.deepcopy(self.__dict__))
+
+
 @dataclass(kw_only=True)
 class EthWallet:
     r"""Stateful variable for storing what is in the agent's wallet
@@ -41,7 +60,7 @@ def copy(self) -> EthWallet:
         """Returns a new copy of self"""
         return EthWallet(**copy.deepcopy(self.__dict__))
 
-    def update(self, wallet_deltas: WalletDeltas) -> None:
+    def update(self, wallet_deltas: EthWalletDeltas) -> None:
         """Update the agent's wallet in-place
 
         Arguments

lib/agent0/agent0/base/state/market_state.py

@@ -5,10 +5,10 @@
 from dataclasses import dataclass
 
 import elfpy
-from elfpy import types
+from agent0.base import freezable
 
 
-@types.freezable(frozen=False, no_new_attribs=False)
+@freezable(frozen=False, no_new_attribs=False)
 @dataclass(kw_only=True)
 class BaseMarketState:
     r"""The state of an AMM."""

lib/agent0/agent0/base/types.py

@@ -0,0 +1,117 @@
+"""Core types used across the repo"""
+from __future__ import annotations  # types will be strings by default in 3.11
+
+from dataclasses import asdict, dataclass, is_dataclass, replace
+from enum import Enum
+from functools import wraps
+from typing import Any, Type
+
+from fixedpointmath import FixedPoint
+
+
+class FrozenClass:
+    """Config object with frozen attributes"""
+
+    def freeze(self):
+        """Disallows changing existing members"""
+        return NotImplemented
+
+    def disable_new_attribs(self):
+        """Disallows adding new members"""
+        return NotImplemented
+
+    def astype(self, _new_type):
+        """Cast all member attributes to a new type"""
+        return NotImplemented
+
+    @property
+    def dtypes(self):
+        """Return a dict listing name & type of each member variable"""
+        return NotImplemented
+
+
+def freezable(frozen: bool = False, no_new_attribs: bool = False) -> Type:
+    r"""A wrapper that allows classes to be frozen, such that existing member attributes cannot be changed"""
+
+    def decorator(cls: Type) -> Type:
+        # this decorator should only be placed atop a dataclass
+        if not is_dataclass(cls):
+            raise TypeError("The class must be a data class.")
+
+        @wraps(wrapped=cls, updated=())
+        class DecoratedFrozenClass(cls, FrozenClass):
+            """Subclass cls to enable freezing of attributes
+
+            .. todo:: resolve why pyright cannot access member "freeze" when instantiated_class.freeze() is called
+            """
+
+            def __init__(self, *args, frozen=frozen, no_new_attribs=no_new_attribs, **kwargs) -> None:
+                super().__init__(*args, **kwargs)
+                super().__setattr__("frozen", frozen)
+                super().__setattr__("no_new_attribs", no_new_attribs)
+
+            def __setattr__(self, attrib: str, value: Any) -> None:
+                if hasattr(self, attrib) and hasattr(self, "frozen") and getattr(self, "frozen"):
+                    raise AttributeError(f"{self.__class__.__name__} is frozen, cannot change attribute '{attrib}'.")
+                if not hasattr(self, attrib) and hasattr(self, "no_new_attribs") and getattr(self, "no_new_attribs"):
+                    raise AttributeError(
+                        f"{self.__class__.__name__} has no_new_attribs set, cannot add attribute '{attrib}'."
+                    )
+                super().__setattr__(attrib, value)
+
+            def freeze(self) -> None:
+                """disallows changing existing members"""
+                super().__setattr__("frozen", True)
+
+            def disable_new_attribs(self) -> None:
+                """disallows adding new members"""
+                super().__setattr__("no_new_attribs", True)
+
+            def astype(self, new_type):
+                """Cast all member attributes to a new type"""
+                new_data = {}
+                for attr_name, attr_value in asdict(self).items():
+                    try:
+                        if isinstance(attr_value, list):
+                            new_data[attr_name] = [new_type(val) for val in attr_value]
+                        else:
+                            new_data[attr_name] = new_type(attr_value)
+                        self.__annotations__[attr_name] = new_type
+                    except (ValueError, TypeError) as err:
+                        raise TypeError(
+                            f"unable to cast {attr_name=} of type {type(attr_value)=} to {new_type=}"
+                        ) from err
+                # create a new instance of the data class with the updated
+                # attributes, rather than modifying the current instance in-place
+                return replace(self, **new_data)
+
+            @property
+            def dtypes(self) -> dict[str, type]:
+                """Return a dict listing name & type of each member variable"""
+                dtypes_dict: dict[str, type] = {}
+                for attr_name, attr_value in asdict(self).items():
+                    dtypes_dict[attr_name] = type(attr_value)
+                return dtypes_dict
+
+        # Set the name of the wrapped class to the name of the input class to preserve metadata
+        DecoratedFrozenClass.__name__ = cls.__name__
+        return DecoratedFrozenClass
+
+    return decorator
+
+
+class TokenType(Enum):
+    r"""A type of token"""
+
+    BASE = "base"
+
+
+@dataclass
+class Quantity:
+    r"""An amount with a unit"""
+
+    amount: FixedPoint
+    unit: TokenType
+
+    def __neg__(self):
+        return Quantity(amount=-self.amount, unit=self.unit)

lib/agent0/agent0/hyperdrive/agents/__init__.py

@@ -1,3 +1,2 @@
 """Account and wallet with Hyperdrive specific parts"""
 from .hyperdrive_account import HyperdriveAgent
-from .hyperdrive_wallet import HyperdriveWallet

lib/agent0/agent0/hyperdrive/agents/hyperdrive_account.py

@@ -4,15 +4,15 @@
 import logging
 from typing import Generic, TypeVar
 
+from agent0.base import Quantity, TokenType
 from agent0.base.agents import EthAgent
 from agent0.base.policies import BasePolicy
-from elfpy.markets.hyperdrive import HyperdriveMarket, HyperdriveMarketAction, MarketActionType
-from elfpy.types import MarketType, Quantity, TokenType, Trade
+from agent0.hyperdrive.state import HyperdriveActionType, HyperdriveMarketAction, HyperdriveWallet
+from elfpy.markets.hyperdrive import HyperdriveMarket
+from elfpy.types import MarketType, Trade
 from eth_account.signers.local import LocalAccount
 from hexbytes import HexBytes
 
-from .hyperdrive_wallet import HyperdriveWallet
-
 Policy = TypeVar("Policy", bound=BasePolicy)
 Market = TypeVar(
     "Market", bound=HyperdriveMarket
@@ -89,7 +89,7 @@ def liquidation_trades(self) -> list[Trade[MarketAction]]:
                     Trade(
                         market_type=MarketType.HYPERDRIVE,
                         market_action=HyperdriveMarketAction(
-                            action_type=MarketActionType.CLOSE_LONG,
+                            action_type=HyperdriveActionType.CLOSE_LONG,
                             trade_amount=long.balance,
                             wallet=self.wallet,  # type: ignore
                             maturity_time=maturity_time,
@@ -104,7 +104,7 @@ def liquidation_trades(self) -> list[Trade[MarketAction]]:
                     Trade(
                         market_type=MarketType.HYPERDRIVE,
                         market_action=HyperdriveMarketAction(
-                            action_type=MarketActionType.CLOSE_SHORT,
+                            action_type=HyperdriveActionType.CLOSE_SHORT,
                             trade_amount=short.balance,
                             wallet=self.wallet,  # type: ignore
                             maturity_time=maturity_time,
@@ -118,7 +118,7 @@ def liquidation_trades(self) -> list[Trade[MarketAction]]:
                 Trade(
                     market_type=MarketType.HYPERDRIVE,
                     market_action=HyperdriveMarketAction(
-                        action_type=MarketActionType.REMOVE_LIQUIDITY,
+                        action_type=HyperdriveActionType.REMOVE_LIQUIDITY,
                         trade_amount=self.wallet.lp_tokens,
                         wallet=self.wallet,  # type: ignore
                     ),
@@ -145,7 +145,10 @@ def get_trades(self, market: Market) -> list[Trade[MarketAction]]:
         # edit each action in place
         for action in actions:
             if action.market_type == MarketType.HYPERDRIVE and action.market_action.maturity_time is None:
-                action.market_action.maturity_time = market.latest_checkpoint_time + market.position_duration.seconds
+                # TODO market latest_checkpoint_time and position_duration should be in ints
+                action.market_action.maturity_time = int(market.latest_checkpoint_time) + int(
+                    market.position_duration.seconds
+                )
                 if action.market_action.trade_amount <= 0:
                     raise ValueError("Trade amount cannot be zero or negative.")
         return actions

lib/agent0/agent0/hyperdrive/exec/__init__.py

@@ -12,5 +12,5 @@
 from .fund_agents import fund_agents
 from .get_agent_accounts import get_agent_accounts
 from .run_agents import run_agents
-from .setup_experiment import register_username, setup_experiment
+from .setup_experiment import setup_experiment
 from .trade_loop import get_wait_for_new_block, trade_if_new_block