Add relative imports, small bugfixes

testing/docker-compose.yaml

@@ -1,12 +1,14 @@
 version: '3.1'
 services:
   db:
-    image: ghcr.io/dbsystel/postgresql-partman:15-5
+    build:
+      dockerfile: postgres.Dockerfile
     restart: "always"
     environment:
       POSTGRESQL_PASSWORD: mypassword
       POSTGRESQL_DATABASE: tycho_indexer_0
       POSTGRESQL_USERNAME: postgres
+      POSTGRESQL_SHARED_PRELOAD_LIBRARIES: pg_cron
     ports:
       - "5431:5432"
     volumes:

testing/postgres.Dockerfile

@@ -0,0 +1,16 @@
+# This Dockerfile creates a custom postgres image used for CI and local deployment.
+# This is required because we use some postgres extensions that aren't in the generic Postgres image such as pg_partman or pg_cron.
+
+# As an image with pg_partman already exist, we start from this one an add pg_cron and possibly other extensions on top of that.
+FROM ghcr.io/dbsystel/postgresql-partman:15-5
+ARG PGCRON_VERSION="1.6.2"
+USER root
+RUN cd /tmp \
+    && wget "https://github.com/citusdata/pg_cron/archive/refs/tags/v${PGCRON_VERSION}.tar.gz" \
+    && tar zxf v${PGCRON_VERSION}.tar.gz \
+    && cd pg_cron-${PGCRON_VERSION} \
+    && make \
+    && make install \
+    && cd .. && rm -r pg_cron-${PGCRON_VERSION} v${PGCRON_VERSION}.tar.gz
+RUN echo "cron.database_name = 'tycho_indexer_0'" >> /opt/bitnami/postgresql/conf/postgresql.conf
+USER 1001

testing/tycho.py

@@ -7,7 +7,9 @@ import os
 import psycopg2
 from psycopg2 import sql
 
-binary_path = "./testing/tycho-indexer"
+from pathlib import Path
+
+binary_path = Path(__file__).parent / "tycho-indexer"
 
 
 class TychoRunner:

tycho/tycho/models.py

@@ -1,68 +0,0 @@
-import datetime
-from dataclasses import dataclass
-from enum import Enum, IntEnum, auto
-from typing import Union
-
-from pydantic import BaseModel, Field
-
-from tycho.tycho.pool_state import ThirdPartyPool
-
-Address = str
-
-
-class Blockchain(Enum):
-    ethereum = "ethereum"
-    arbitrum = "arbitrum"
-    polygon = "polygon"
-    zksync = "zksync"
-
-
-class EVMBlock(BaseModel):
-    id: int
-    ts: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
-    hash_: str
-
-
-class EthereumToken(BaseModel):
-    symbol: str
-    address: str
-    decimals: int
-    gas: Union[int, list[int]] = 29000
-
-
-class DatabaseType(Enum):
-    # Make call to the node each time it needs a storage (unless cached from a previous call).
-    rpc_reader = "rpc_reader"
-    # Connect to Tycho and cache the whole state of a target contract, the state is continuously updated by Tycho.
-    # To use this we need Tycho to be configured to index the target contract state.
-    tycho = "tycho"
-
-
-class Capability(IntEnum):
-    SellSide = auto()
-    BuySide = auto()
-    PriceFunction = auto()
-    FeeOnTransfer = auto()
-    ConstantPrice = auto()
-    TokenBalanceIndependent = auto()
-    ScaledPrice = auto()
-
-
-class SynchronizerState(Enum):
-    started = "started"
-    ready = "ready"
-    stale = "stale"
-    delayed = "delayed"
-    advanced = "advanced"
-    ended = "ended"
-
-
-@dataclass(repr=False)
-class BlockProtocolChanges:
-    block: EVMBlock
-    pool_states: dict[Address, ThirdPartyPool]
-    """All updated pools"""
-    removed_pools: set[Address]
-    sync_states: dict[str, SynchronizerState]
-    deserialization_time: float
-    """The time it took to deserialize the pool states from the tycho feed message"""

tycho_client/tycho/adapter_contract.py

@@ -17,9 +17,9 @@ from protosim_py import (
     StateUpdate,
 )
 
-from tycho.tycho.constants import EXTERNAL_ACCOUNT
-from tycho.tycho.models import Address, EthereumToken, EVMBlock, Capability
-from tycho.tycho.utils import load_abi, maybe_coerce_error
+from .constants import EXTERNAL_ACCOUNT
+from .models import Address, EthereumToken, EVMBlock, Capability
+from .utils import load_abi, maybe_coerce_error
 
 log = logging.getLogger(__name__)
 

tycho_client/tycho/assets/IERC20.sol

@@ -0,0 +1,78 @@
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts (last updated v4.9.0) (token/ERC20/IERC20.sol)
+
+pragma solidity ^0.8.19;
+
+/**
+ * @dev Interface of the ERC20 standard as defined in the EIP.
+ */
+interface IERC20 {
+    /**
+     * @dev Emitted when `value` tokens are moved from one account (`from`) to
+     * another (`to`).
+     *
+     * Note that `value` may be zero.
+     */
+    event Transfer(address indexed from, address indexed to, uint256 value);
+
+    /**
+     * @dev Emitted when the allowance of a `spender` for an `owner` is set by
+     * a call to {approve}. `value` is the new allowance.
+     */
+    event Approval(address indexed owner, address indexed spender, uint256 value);
+
+    /**
+     * @dev Returns the amount of tokens in existence.
+     */
+    function totalSupply() external view returns (uint256);
+
+    /**
+     * @dev Returns the amount of tokens owned by `account`.
+     */
+    function balanceOf(address account) external view returns (uint256);
+
+    /**
+     * @dev Moves `amount` tokens from the caller's account to `to`.
+     *
+     * Returns a boolean value indicating whether the operation succeeded.
+     *
+     * Emits a {Transfer} event.
+     */
+    function transfer(address to, uint256 amount) external returns (bool);
+
+    /**
+     * @dev Returns the remaining number of tokens that `spender` will be
+     * allowed to spend on behalf of `owner` through {transferFrom}. This is
+     * zero by default.
+     *
+     * This value changes when {approve} or {transferFrom} are called.
+     */
+    function allowance(address owner, address spender) external view returns (uint256);
+
+    /**
+     * @dev Sets `amount` as the allowance of `spender` over the caller's tokens.
+     *
+     * Returns a boolean value indicating whether the operation succeeded.
+     *
+     * IMPORTANT: Beware that changing an allowance with this method brings the risk
+     * that someone may use both the old and the new allowance by unfortunate
+     * transaction ordering. One possible solution to mitigate this race
+     * condition is to first reduce the spender's allowance to 0 and set the
+     * desired value afterwards:
+     * https://github.com/ethereum/EIPs/issues/20#issuecomment-263524729
+     *
+     * Emits an {Approval} event.
+     */
+    function approve(address spender, uint256 amount) external returns (bool);
+
+    /**
+     * @dev Moves `amount` tokens from `from` to `to` using the
+     * allowance mechanism. `amount` is then deducted from the caller's
+     * allowance.
+     *
+     * Returns a boolean value indicating whether the operation succeeded.
+     *
+     * Emits a {Transfer} event.
+     */
+    function transferFrom(address from, address to, uint256 amount) external returns (bool);
+}

tycho_client/tycho/assets/mocked_ERC20.sol

@@ -0,0 +1,363 @@
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts (last updated v4.9.0) (token/ERC20/ERC20.sol)
+
+pragma solidity ^0.8.19;
+
+import "./IERC20.sol";
+
+
+/**
+ * @dev Provides information about the current execution context, including the
+ * sender of the transaction and its data. While these are generally available
+ * via msg.sender and msg.data, they should not be accessed in such a direct
+ * manner, since when dealing with meta-transactions the account sending and
+ * paying for execution may not be the actual sender (as far as an application
+ * is concerned).
+ *
+ * This contract is only required for intermediate, library-like contracts.
+ */
+abstract contract Context {
+    function _msgSender() internal view virtual returns (address) {
+        return msg.sender;
+    }
+
+    function _msgData() internal view virtual returns (bytes calldata) {
+        return msg.data;
+    }
+}
+
+/**
+ * @dev Interface for the optional metadata functions from the ERC20 standard.
+ *
+ * _Available since v4.1._
+ */
+interface IERC20Metadata is IERC20 {
+    /**
+     * @dev Returns the name of the token.
+     */
+    function name() external view returns (string memory);
+
+    /**
+     * @dev Returns the symbol of the token.
+     */
+    function symbol() external view returns (string memory);
+
+    /**
+     * @dev Returns the decimals places of the token.
+     */
+    function decimals() external view returns (uint8);
+}
+
+/**
+ * @dev Implementation of the {IERC20} interface.
+ *
+ * This implementation is agnostic to the way tokens are created. This means
+ * that a supply mechanism has to be added in a derived contract using {_mint}.
+ *
+ * TIP: For a detailed writeup see our guide
+ * https://forum.openzeppelin.com/t/how-to-implement-erc20-supply-mechanisms/226[How
+ * to implement supply mechanisms].
+ *
+ * The default value of {decimals} is 18. To change this, you should override
+ * this function so it returns a different value.
+ *
+ * We have followed general OpenZeppelin Contracts guidelines: functions revert
+ * instead returning `false` on failure. This behavior is nonetheless
+ * conventional and does not conflict with the expectations of ERC20
+ * applications.
+ *
+ * Additionally, an {Approval} event is emitted on calls to {transferFrom}.
+ * This allows applications to reconstruct the allowance for all accounts just
+ * by listening to said events. Other implementations of the EIP may not emit
+ * these events, as it isn't required by the specification.
+ *
+ * Finally, the non-standard {decreaseAllowance} and {increaseAllowance}
+ * functions have been added to mitigate the well-known issues around setting
+ * allowances. See {IERC20-approve}.
+ */
+contract ERC20 is Context, IERC20, IERC20Metadata {
+    mapping(address => uint256) private _balances;
+
+    mapping(address => mapping(address => uint256)) private _allowances;
+
+    uint256 private _totalSupply;
+
+    string private _name;
+    string private _symbol;
+    uint8 private _decimals;
+
+    /**
+     * @dev Sets the values for {name}, {symbol} and {decimals}.
+     *
+     * All three of these values are immutable: they can only be set once during
+     * construction.
+     */
+    constructor(string memory name_, string memory symbol_, uint8 decimals_) {
+        _name = name_;
+        _symbol = symbol_;
+        _decimals = decimals_;
+    }
+
+    /**
+     * @dev Returns the name of the token.
+     */
+    function name() public view virtual returns (string memory) {
+        return _name;
+    }
+
+    /**
+     * @dev Returns the symbol of the token, usually a shorter version of the
+     * name.
+     */
+    function symbol() public view virtual returns (string memory) {
+        return _symbol;
+    }
+
+    /**
+     * @dev Returns the number of decimals used to get its user representation.
+     * For example, if `decimals` equals `2`, a balance of `505` tokens should
+     * be displayed to a user as `5.05` (`505 / 10 ** 2`).
+     *
+     * Tokens usually opt for a value of 18, imitating the relationship between
+     * Ether and Wei. This is the default value returned by this function, unless
+     * it's overridden.
+     *
+     * NOTE: This information is only used for _display_ purposes: it in
+     * no way affects any of the arithmetic of the contract, including
+     * {IERC20-balanceOf} and {IERC20-transfer}.
+     */
+    function decimals() public view virtual returns (uint8) {
+        return _decimals;
+    }
+
+    /**
+     * @dev See {IERC20-totalSupply}.
+     */
+    function totalSupply() public view virtual returns (uint256) {
+        return _totalSupply;
+    }
+
+    /**
+     * @dev See {IERC20-balanceOf}.
+     */
+    function balanceOf(address account) public view virtual returns (uint256) {
+        return _balances[account];
+    }
+
+    /**
+     * @dev See {IERC20-transfer}.
+     *
+     * Requirements:
+     *
+     * - `to` cannot be the zero address.
+     * - the caller must have a balance of at least `amount`.
+     */
+    function transfer(address to, uint256 amount) public virtual returns (bool) {
+        address owner = _msgSender();
+        _transfer(owner, to, amount);
+        return true;
+    }
+
+    /**
+     * @dev See {IERC20-allowance}.
+     */
+    function allowance(address owner, address spender) public view virtual returns (uint256) {
+        return _allowances[owner][spender];
+    }
+
+    /**
+     * @dev See {IERC20-approve}.
+     *
+     * NOTE: If `amount` is the maximum `uint256`, the allowance is not updated on
+     * `transferFrom`. This is semantically equivalent to an infinite approval.
+     *
+     * Requirements:
+     *
+     * - `spender` cannot be the zero address.
+     */
+    function approve(address spender, uint256 amount) public virtual returns (bool) {
+        address owner = _msgSender();
+        _approve(owner, spender, amount);
+        return true;
+    }
+
+    /**
+     * @dev See {IERC20-transferFrom}.
+     *
+     * Emits an {Approval} event indicating the updated allowance. This is not
+     * required by the EIP. See the note at the beginning of {ERC20}.
+     *
+     * NOTE: Does not update the allowance if the current allowance
+     * is the maximum `uint256`.
+     *
+     * Requirements:
+     *
+     * - `from` and `to` cannot be the zero address.
+     * - `from` must have a balance of at least `amount`.
+     * - the caller must have allowance for ``from``'s tokens of at least
+     * `amount`.
+     */
+    function transferFrom(address from, address to, uint256 amount) public virtual returns (bool) {
+        address spender = _msgSender();
+        _spendAllowance(from, spender, amount);
+        _transfer(from, to, amount);
+        return true;
+    }
+
+    /**
+     * @dev Atomically increases the allowance granted to `spender` by the caller.
+     *
+     * This is an alternative to {approve} that can be used as a mitigation for
+     * problems described in {IERC20-approve}.
+     *
+     * Emits an {Approval} event indicating the updated allowance.
+     *
+     * Requirements:
+     *
+     * - `spender` cannot be the zero address.
+     */
+    function increaseAllowance(address spender, uint256 addedValue) public virtual returns (bool) {
+        address owner = _msgSender();
+        _approve(owner, spender, allowance(owner, spender) + addedValue);
+        return true;
+    }
+
+    /**
+     * @dev Atomically decreases the allowance granted to `spender` by the caller.
+     *
+     * This is an alternative to {approve} that can be used as a mitigation for
+     * problems described in {IERC20-approve}.
+     *
+     * Emits an {Approval} event indicating the updated allowance.
+     *
+     * Requirements:
+     *
+     * - `spender` cannot be the zero address.
+     * - `spender` must have allowance for the caller of at least
+     * `subtractedValue`.
+     */
+    function decreaseAllowance(address spender, uint256 subtractedValue) public virtual returns (bool) {
+        address owner = _msgSender();
+        uint256 currentAllowance = allowance(owner, spender);
+        require(currentAllowance >= subtractedValue, "ERC20: decreased allowance below zero");
+        unchecked {
+            _approve(owner, spender, currentAllowance - subtractedValue);
+        }
+
+        return true;
+    }
+
+    /**
+     * @dev Moves `amount` of tokens from `from` to `to`.
+     *
+     * This internal function is equivalent to {transfer}, and can be used to
+     * e.g. implement automatic token fees, slashing mechanisms, etc.
+     *
+     * Emits a {Transfer} event.
+     *
+     * NOTE: This function is not virtual, {_update} should be overridden instead.
+     */
+    function _transfer(address from, address to, uint256 amount) internal {
+        require(from != address(0), "ERC20: transfer from the zero address");
+        require(to != address(0), "ERC20: transfer to the zero address");
+        _update(from, to, amount);
+    }
+
+    /**
+     * @dev Transfers `amount` of tokens from `from` to `to`, or alternatively mints (or burns) if `from` (or `to`) is
+     * the zero address. All customizations to transfers, mints, and burns should be done by overriding this function.
+     *
+     * Emits a {Transfer} event.
+     */
+    function _update(address from, address to, uint256 amount) internal virtual {
+        if (from == address(0)) {
+            _totalSupply += amount;
+        } else {
+            uint256 fromBalance = _balances[from];
+            require(fromBalance >= amount, "ERC20: transfer amount exceeds balance");
+            unchecked {
+                // Overflow not possible: amount <= fromBalance <= totalSupply.
+                _balances[from] = fromBalance - amount;
+            }
+        }
+
+        if (to == address(0)) {
+            unchecked {
+                // Overflow not possible: amount <= totalSupply or amount <= fromBalance <= totalSupply.
+                _totalSupply -= amount;
+            }
+        } else {
+            unchecked {
+                // Overflow not possible: balance + amount is at most totalSupply, which we know fits into a uint256.
+                _balances[to] += amount;
+            }
+        }
+
+        emit Transfer(from, to, amount);
+    }
+
+    /**
+     * @dev Creates `amount` tokens and assigns them to `account`, by transferring it from address(0).
+     * Relies on the `_update` mechanism
+     *
+     * Emits a {Transfer} event with `from` set to the zero address.
+     *
+     * NOTE: This function is not virtual, {_update} should be overridden instead.
+     */
+    function _mint(address account, uint256 amount) internal {
+        require(account != address(0), "ERC20: mint to the zero address");
+        _update(address(0), account, amount);
+    }
+
+    /**
+     * @dev Destroys `amount` tokens from `account`, by transferring it to address(0).
+     * Relies on the `_update` mechanism.
+     *
+     * Emits a {Transfer} event with `to` set to the zero address.
+     *
+     * NOTE: This function is not virtual, {_update} should be overridden instead
+     */
+    function _burn(address account, uint256 amount) internal {
+        require(account != address(0), "ERC20: burn from the zero address");
+        _update(account, address(0), amount);
+    }
+
+    /**
+     * @dev Sets `amount` as the allowance of `spender` over the `owner` s tokens.
+     *
+     * This internal function is equivalent to `approve`, and can be used to
+     * e.g. set automatic allowances for certain subsystems, etc.
+     *
+     * Emits an {Approval} event.
+     *
+     * Requirements:
+     *
+     * - `owner` cannot be the zero address.
+     * - `spender` cannot be the zero address.
+     */
+    function _approve(address owner, address spender, uint256 amount) internal virtual {
+        require(owner != address(0), "ERC20: approve from the zero address");
+        require(spender != address(0), "ERC20: approve to the zero address");
+
+        _allowances[owner][spender] = amount;
+        emit Approval(owner, spender, amount);
+    }
+
+    /**
+     * @dev Updates `owner` s allowance for `spender` based on spent `amount`.
+     *
+     * Does not update the allowance amount in case of infinite allowance.
+     * Revert if not enough allowance is available.
+     *
+     * Might emit an {Approval} event.
+     */
+    function _spendAllowance(address owner, address spender, uint256 amount) internal virtual {
+        uint256 currentAllowance = allowance(owner, spender);
+        if (currentAllowance != type(uint256).max) {
+            require(currentAllowance >= amount, "ERC20: insufficient allowance");
+            unchecked {
+                _approve(owner, spender, currentAllowance - amount);
+            }
+        }
+    }
+}

tycho_client/tycho/constants.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Final
 
-TYCHO_CLIENT_FOLDER = Path(__file__) / "bins"
+TYCHO_CLIENT_FOLDER = Path(__file__).parent / "bins"
 TYCHO_CLIENT_LOG_FOLDER = TYCHO_CLIENT_FOLDER / "logs"
 
 EXTERNAL_ACCOUNT: Final[str] = "0xf847a638E44186F3287ee9F8cAF73FF4d4B80784"

tycho_client/tycho/decoders.py

@@ -2,10 +2,10 @@ from decimal import Decimal
 from logging import getLogger
 from typing import Any
 
-from tycho.tycho.exceptions import TychoDecodeError
-from tycho.tycho.models import EVMBlock, EthereumToken
-from tycho.tycho.pool_state import ThirdPartyPool
-from tycho.tycho.utils import decode_tycho_exchange
+from .exceptions import TychoDecodeError
+from .models import EVMBlock, EthereumToken
+from .pool_state import ThirdPartyPool
+from .utils import decode_tycho_exchange
 
 log = getLogger(__name__)
 
@@ -98,7 +98,7 @@ class ThirdPartyPoolTychoDecoder:
 
     @staticmethod
     def apply_update(
-            pool: ThirdPartyPool,
+        pool: ThirdPartyPool,
         pool_update: dict[str, Any],
         balance_updates: dict[str, Any],
         block: EVMBlock,

tycho_client/tycho/models.py

@@ -0,0 +1,108 @@
+import datetime
+from decimal import Decimal, localcontext, Context, ROUND_FLOOR, InvalidOperation
+from enum import Enum, IntEnum, auto
+from fractions import Fraction
+from logging import getLogger
+from typing import Union
+
+from pydantic import BaseModel, Field
+
+Address = str
+
+log = getLogger(__name__)
+
+
+class Blockchain(Enum):
+    ethereum = "ethereum"
+    arbitrum = "arbitrum"
+    polygon = "polygon"
+    zksync = "zksync"
+
+
+class EVMBlock(BaseModel):
+    id: int
+    ts: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
+    hash_: str
+
+
+class EthereumToken(BaseModel):
+    symbol: str
+    address: str
+    decimals: int
+    gas: Union[int, list[int]] = 29000
+
+    def to_onchain_amount(self, amount: Union[float, Decimal, str]) -> int:
+        """Converts floating-point numerals to an integer, by shifting right by the
+        token's maximum amount of decimals (e.g.: 1.000000 becomes 1000000).
+        For the reverse operation please see self.from_onchain_amount
+        """
+        if not isinstance(amount, Decimal):
+            log.warning(f"Expected variable of type Decimal. Got {type(amount)}.")
+
+        with localcontext(Context(rounding=ROUND_FLOOR, prec=256)):
+            amount = Decimal(str(amount)) * (10 ** self.decimals)
+            try:
+                amount = amount.quantize(Decimal("1.0"))
+            except InvalidOperation:
+                log.error(
+                    f"Quantize failed for {self.symbol}, {amount}, {self.decimals}"
+                )
+            return int(amount)
+
+    def from_onchain_amount(
+        self, onchain_amount: Union[int, Fraction], quantize: bool = True
+    ) -> Decimal:
+        """Converts an Integer to a quantized decimal, by shifting left by the token's
+        maximum amount of decimals (e.g.: 1000000 becomes 1.000000 for a 6-decimal token
+        For the reverse operation please see self.to_onchain_amount
+
+        If the onchain_amount is too low, then using quantize can underflow without
+        raising and the offchain amount returned is 0.
+        See _decimal.Decimal.quantize docstrings for details.
+
+        Quantize is needed for UniswapV2.
+        """
+        with localcontext(self._dec_context):
+            if isinstance(onchain_amount, Fraction):
+                return (
+                    Decimal(onchain_amount.numerator)
+                    / Decimal(onchain_amount.denominator)
+                    / Decimal(10 ** self.decimals)
+                ).quantize(Decimal(f"{1 / 10 ** self.decimals}"))
+            if quantize is True:
+                try:
+                    amount = (
+                        Decimal(str(onchain_amount)) / 10 ** self.decimals
+                    ).quantize(Decimal(f"{1 / 10 ** self.decimals}"))
+                except InvalidOperation:
+                    amount = Decimal(str(onchain_amount)) / Decimal(10 ** self.decimals)
+            else:
+                amount = Decimal(str(onchain_amount)) / Decimal(10 ** self.decimals)
+            return amount
+
+
+class DatabaseType(Enum):
+    # Make call to the node each time it needs a storage (unless cached from a previous call).
+    rpc_reader = "rpc_reader"
+    # Connect to Tycho and cache the whole state of a target contract, the state is continuously updated by Tycho.
+    # To use this we need Tycho to be configured to index the target contract state.
+    tycho = "tycho"
+
+
+class Capability(IntEnum):
+    SellSide = auto()
+    BuySide = auto()
+    PriceFunction = auto()
+    FeeOnTransfer = auto()
+    ConstantPrice = auto()
+    TokenBalanceIndependent = auto()
+    ScaledPrice = auto()
+
+
+class SynchronizerState(Enum):
+    started = "started"
+    ready = "ready"
+    stale = "stale"
+    delayed = "delayed"
+    advanced = "advanced"
+    ended = "ended"

tycho_client/tycho/pool_state.py

@@ -5,17 +5,17 @@ from copy import deepcopy
 from decimal import Decimal
 from fractions import Fraction
 from logging import getLogger
-from typing import Optional, cast, TypeVar
+from typing import Optional, cast, TypeVar, Annotated, DefaultDict
 
 from eth_typing import HexStr
 from protosim_py import SimulationEngine, AccountInfo
 from pydantic import BaseModel, PrivateAttr, Field
 
-from tycho.tycho.adapter_contract import AdapterContract
-from tycho.tycho.constants import MAX_BALANCE, EXTERNAL_ACCOUNT
-from tycho.tycho.exceptions import RecoverableSimulationException
-from tycho.tycho.models import EVMBlock, Capability, Address, EthereumToken
-from tycho.tycho.utils import (
+from .adapter_contract import AdapterContract
+from .constants import MAX_BALANCE, EXTERNAL_ACCOUNT
+from .exceptions import RecoverableSimulationException
+from .models import EVMBlock, Capability, Address, EthereumToken
+from .utils import (
     create_engine,
     get_contract_bytecode,
     frac_to_decimal,
@@ -54,9 +54,11 @@ class ThirdPartyPool(BaseModel):
     """The contract address for where protocol balances are stored (i.e. a vault contract).
     If given, balances will be overwritten here instead of on the pool contract during simulations."""
 
-    block_lasting_overwrites: defaultdict[Address, dict[int, int]] = Field(
-        default_factory=lambda: defaultdict(dict)
-    )
+    block_lasting_overwrites: DefaultDict[
+        Address,
+        Annotated[dict[int, int], Field(default_factory=lambda: defaultdict[dict])],
+    ]
+
     """Storage overwrites that will be applied to all simulations. They will be cleared
     when ``clear_all_cache`` is called, i.e. usually at each block. Hence the name."""
 

tycho_client/tycho/tycho_adapter.py

@@ -4,6 +4,7 @@ import platform
 import time
 from asyncio.subprocess import STDOUT, PIPE
 from collections import defaultdict
+from dataclasses import dataclass
 from datetime import datetime
 from decimal import Decimal
 from http.client import HTTPException
@@ -13,18 +14,13 @@ from typing import Any, Optional, Dict
 import requests
 from protosim_py import AccountUpdate, AccountInfo, BlockHeader
 
-from tycho.tycho.constants import TYCHO_CLIENT_LOG_FOLDER, TYCHO_CLIENT_FOLDER
-from tycho.tycho.decoders import ThirdPartyPoolTychoDecoder
-from tycho.tycho.exceptions import APIRequestError, TychoClientException
-from tycho.tycho.models import (
-    Blockchain,
-    EVMBlock,
-    EthereumToken,
-    BlockProtocolChanges,
-    SynchronizerState,
-)
-from tycho.tycho.tycho_db import TychoDBSingleton
-from tycho.tycho.utils import create_engine
+from .pool_state import ThirdPartyPool
+from .constants import TYCHO_CLIENT_LOG_FOLDER, TYCHO_CLIENT_FOLDER
+from .decoders import ThirdPartyPoolTychoDecoder
+from .exceptions import APIRequestError, TychoClientException
+from .models import Blockchain, EVMBlock, EthereumToken, SynchronizerState, Address
+from .tycho_db import TychoDBSingleton
+from .utils import create_engine
 
 log = getLogger(__name__)
 
@@ -34,7 +30,7 @@ class TokenLoader:
         self,
         tycho_url: str,
         blockchain: Blockchain,
-        min_token_quality: Optional[int] = 51,
+        min_token_quality: Optional[int] = 0,
     ):
         self.tycho_url = tycho_url
         self.blockchain = blockchain
@@ -70,6 +66,34 @@ class TokenLoader:
 
         return formatted_tokens
 
+    def get_token_subset(self, addresses: list[str]) -> dict[str, EthereumToken]:
+        """Loads a subset of tokens from Tycho RPC"""
+        url = self.tycho_url + self.endpoint.format(self.blockchain.value)
+        page = 0
+
+        start = time.monotonic()
+        all_tokens = []
+        while data := self._get_all_with_pagination(
+            url=url,
+            page=page,
+            limit=self._token_limit,
+            params={"min_quality": self.min_token_quality, "addresses": addresses},
+        ):
+            all_tokens.extend(data)
+            page += 1
+            if len(data) < self._token_limit:
+                break
+
+        log.info(f"Loaded {len(all_tokens)} tokens in {time.monotonic() - start:.2f}s")
+
+        formatted_tokens = dict()
+
+        for token in all_tokens:
+            formatted = EthereumToken(**token)
+            formatted_tokens[formatted.address] = formatted
+
+        return formatted_tokens
+
     @staticmethod
     def _get_all_with_pagination(
         url: str, params: Optional[Dict] = None, page: int = 0, limit: int = 50
@@ -87,6 +111,17 @@ class TokenLoader:
         return r.json()["tokens"]
 
 
+@dataclass(repr=False)
+class BlockProtocolChanges:
+    block: EVMBlock
+    pool_states: dict[Address, ThirdPartyPool]
+    """All updated pools"""
+    removed_pools: set[Address]
+    sync_states: dict[str, SynchronizerState]
+    deserialization_time: float
+    """The time it took to deserialize the pool states from the tycho feed message"""
+
+
 class TychoPoolStateStreamAdapter:
     def __init__(
         self,
@@ -95,7 +130,7 @@ class TychoPoolStateStreamAdapter:
         decoder: ThirdPartyPoolTychoDecoder,
         blockchain: Blockchain,
         min_tvl: Optional[Decimal] = 10,
-        min_token_quality: Optional[int] = 51,
+        min_token_quality: Optional[int] = 0,
         include_state=True,
     ):
         """
@@ -122,7 +157,7 @@ class TychoPoolStateStreamAdapter:
 
         # Loads tokens from Tycho
         self._tokens: dict[str, EthereumToken] = TokenLoader(
-            tycho_url=self.tycho_url,
+            tycho_url=f"http://{self.tycho_url}",
             blockchain=self._blockchain,
             min_token_quality=self.min_token_quality,
         ).get_tokens()
@@ -139,11 +174,11 @@ class TychoPoolStateStreamAdapter:
 
         cmd = [
             "--log-folder",
-            TYCHO_CLIENT_LOG_FOLDER,
+            str(TYCHO_CLIENT_LOG_FOLDER),
             "--tycho-url",
             self.tycho_url,
             "--min-tvl",
-            self.min_tvl,
+            str(self.min_tvl),
         ]
         if not self._include_state:
             cmd.append("--no-state")
@@ -152,7 +187,7 @@ class TychoPoolStateStreamAdapter:
 
         log.debug(f"Starting tycho-client binary at {bin_path}. CMD: {cmd}")
         self.tycho_client = await asyncio.create_subprocess_exec(
-            bin_path, *cmd, stdout=PIPE, stderr=STDOUT, limit=2 ** 64
+            str(bin_path), *cmd, stdout=PIPE, stderr=STDOUT, limit=2 ** 64
         )
 
     @staticmethod

tycho_client/tycho/utils.py

@@ -13,10 +13,10 @@ from hexbytes import HexBytes
 from protosim_py import SimulationEngine, AccountInfo
 from web3 import Web3
 
-from tycho.tycho.constants import EXTERNAL_ACCOUNT, MAX_BALANCE
-from tycho.tycho.exceptions import OutOfGas
-from tycho.tycho.models import Address, EthereumToken
-from tycho.tycho.tycho_db import TychoDBSingleton
+from .constants import EXTERNAL_ACCOUNT, MAX_BALANCE
+from .exceptions import OutOfGas
+from .models import Address, EthereumToken
+from .tycho_db import TychoDBSingleton
 
 log = getLogger(__name__)
 
@@ -182,6 +182,8 @@ def get_storage_slot_at_key(key: Address, mapping_slot: int) -> int:
 
 @lru_cache
 def get_contract_bytecode(name: str) -> bytes:
+    """Load contract bytecode from a file in the assets directory"""
+    # TODO: Check if this locaation is correct
     with open(Path(__file__).parent / "assets" / name, "rb") as fh:
         code = fh.read()
     return code