Source code for tradeexecutor.backtest.simulated_wallet

"""On-chain balance simulator for backtesting.

- We simulate ERC-20 token balances as they would appear in live trading

- The simulation will catch bugs like incorrect accounting or rounding
  errors for the trade execution

"""

import logging

from decimal import Decimal
from typing import Dict, Tuple

import pandas as pd

from tradeexecutor.state.identifier import AssetIdentifier
from tradeexecutor.utils.accuracy import QUANTITY_EPSILON
from tradeexecutor.state.types import JSONHexAddress


logger = logging.getLogger(__name__)


class OutOfSimulatedBalance(Exception):
    """Backtest detected a situation that we run out of a token balance.

    We simulate both internal ledger and on-chain balances in backtesting.
    This exception usually means that there is a bug an internal accounting of trade execution.
    """


[docs]class SimulatedWallet: """A wallet that keeps token balances by ERC-20 address. - Simulates different incoming and outgoing tokens from a wallet includive, aToken and vToken interest amounts with :py:meth:`rebalance`. - If a backtest tries to transfer a token it does not have, or does not have enough of it, raise an error - Will catch bugs in internal accounting """
[docs] def __init__(self): #: Raw balances token address -> balance self.balances: Dict[str, Decimal] = {} #: token address -> asset info for debug logging self.tokens: Dict[str, AssetIdentifier] = {} #: Start with zero nonce like Ethereum acconts self.nonce = 0
[docs] def update_balance( self, token: JSONHexAddress | AssetIdentifier, delta: Decimal, reason: str = None, epsilon=QUANTITY_EPSILON, ): """Change the token balance of some delta. Check that balance does not go zero. :param token: Token we receive or send. Give either raw address or asset definition. Any asset definion is automatically added to our internal tracking list for diagnostics. :param delta: The amount of token, human units :param reason: Reason for this change. Only used for backtesting diagnostics. :param epsilon: If the balance goes below this dust threshold, go all the way to zero """ if isinstance(token, AssetIdentifier): self.update_token_info(token) if abs(delta) <= epsilon: # No changes, don't make noise for logging return token_symbol = self.get_token_symbol(token) assert isinstance(delta, Decimal), f"Expected decimal got: {delta.__class__}: {delta}" old_balance = self.get_balance(token) new_balance = old_balance + delta if abs(new_balance) < epsilon: # Decimal rounding errors new_balance = Decimal(0) reason = reason or "unhinted reason" logger.info("Wallet balance for %s: %f -> %f (%+f), %s", token_symbol, old_balance, new_balance, delta, reason) if new_balance < 0: raise OutOfSimulatedBalance(f"Simulated wallet balance went negative {new_balance} for token {token_symbol}, because of {reason}.\nOld balance was {old_balance}") if new_balance != 0 and new_balance <= epsilon: logger.info("Fixing dust balance to zero, calculated balance %f, epsilon %s", new_balance, epsilon) new_balance = Decimal(0) self.set_balance(token, new_balance)
[docs] def set_balance(self, token: JSONHexAddress | AssetIdentifier, amount: Decimal): """Directly set balance. :param token: Token we receive or send. Give either raw address or asset definition. Any asset definion is automatically added to our internal tracking list for diagnostics. :param amount: New absolute balance. """ if isinstance(token, AssetIdentifier): self.update_token_info(token) token_address = token.address else: token_address = token assert token_address.lower() == token_address, f"No checksummed addresses: {token_address}" assert token_address.startswith("0x") assert isinstance(amount, Decimal), f"Expected decimal got: {amount.__class__}: {amount}" self.balances[token_address] = amount
[docs] def rebase(self, token: JSONHexAddress | AssetIdentifier, new_amount: Decimal): """Set rebase token amount. aToken / vToken accrues interest or debt. :param new_amount: Abs token amount on the chain """ if isinstance(token, AssetIdentifier): token_address = token.address elif isinstance(token, str): token_address = token else: raise AssertionError(f"Does not understand: {token.__class__}: {token}") assert token_address in self.balances, f"Cannot rebase token we do not have: {token}" self.set_balance(token, new_amount)
[docs] def get_balance(self, token: JSONHexAddress | AssetIdentifier) -> Decimal: """Get on-chain balance of one token. :return: Human-readable token balance """ if isinstance(token, AssetIdentifier): token_address = token.address elif isinstance(token, str): token_address = token else: raise AssertionError(f"Does not understand: {token.__class__}: {token}") assert token_address.lower() == token_address, f"Non-checksummed address required: {token}" return self.balances.get(token_address, Decimal(0))
[docs] def fetch_nonce_and_tx_hash(self) -> Tuple[int, str]: """Allocates a dummy nonce for a transaction. :return: Tuple (nonce, tx_hash) """ nonce = self.nonce tx_hash = hex(nonce) self.nonce += 1 return nonce, tx_hash
[docs] def get_token_symbol(self, token: JSONHexAddress | AssetIdentifier) -> str: """Get the human readable name of token for diagnostics output.""" if isinstance(token, AssetIdentifier): return token.token_symbol # Cached symbol asset = self.tokens.get(token, {}) if asset: return asset.token_symbol return f"<token {token}>"
[docs] def update_token_info(self, asset: AssetIdentifier): """Set the token info for a particular ERC-20. This way the wallet has metadata on what token it has and can produce better diagnostics output. Automatically called by :py:meth:`update_balance`. """ assert isinstance(asset, AssetIdentifier) self.tokens[asset.address] = asset
[docs] def get_all_balances(self) -> pd.DataFrame: """Show the status of the wallet as a printable DataFrame. Example: .. code-block:: python print(wallet.get_all_balances()) Output: .. code-block:: text Balance Token USDC 9500 aUSDC 998.4999999999999999687749774 vWETH 0.3003021039165400376391259260 """ tokens = sorted([a for a in self.balances.keys()]) data = [(self.get_token_symbol(a), self.get_balance(a)) for a in tokens] df = pd.DataFrame(data, columns=["Token", "Balance"]) df = df.set_index("Token") return df
[docs] def verify_balances( self, expected: Dict[AssetIdentifier, Decimal], epsilon=0.0001, ) -> Tuple[bool, pd.DataFrame]: """Check that our simulated balances are what we expect. :return: Clean or not, all assets table. """ data = [] clean = True for asset, amount in expected.items(): on_chain_balance = self.get_balance(asset.address) mismatch = False diff = abs(on_chain_balance - amount) if diff / amount >= epsilon: clean = False mismatch = True data.append((asset.token_symbol, amount, on_chain_balance, diff, mismatch)) df = pd.DataFrame(data, columns=["Asset", "Expected", "Actual", "Diff", "Mismatch"]) df.set_index("Asset") return clean, df