tristero 0.1.6__py3-none-any.whl → 0.2.1__py3-none-any.whl

tristero/__init__.py

@@ -1,12 +1,18 @@
-from tristero.client import OrderFailedException, StuckException, SwapException, TokenSpec, execute_swap, start_swap
+from tristero.client import OrderFailedException, StuckException, SwapException, TokenSpec, execute_permit2_swap, start_permit2_swap, start_feather_swap, wait_for_feather_completion, wait_for_completion, wait_for_completion_with_retry
 from tristero.config import set_config
+from tristero.data import ChainID
 
 __all__ = [
+    "ChainID",
     "TokenSpec",
     "StuckException",
     "OrderFailedException",
     "SwapException",
-    "execute_swap",
-    "start_swap",
+    "start_permit2_swap",
+    "start_feather_swap",
+    "execute_permit2_swap",
+    "wait_for_feather_completion",
+    "wait_for_completion",
+    "wait_for_completion_with_retry",
     "set_config",
 ]

tristero/api.py

@@ -1,73 +1,8 @@
-from enum import Enum
 from typing import Any, Optional
 import httpx
 from websockets.asyncio.client import connect
 from tristero.config import get_config
-
-class ChainID(str, Enum):
-    arbitrum = "42161"
-    avalanche = "43114"
-    base = "8453"
-    blast = "81457"
-    bsc = "56"
-    celo = "42220"
-    ethereum = "1"
-    mantle = "5000"
-    mode = "34443"
-    opbnb = "204"
-    optimism = "10"
-    polygon = "137"
-    scroll = "534352"
-    solana = "1151111081099710"
-    tron = "728126428"
-    unichain = "130"
-    sei = "1329"
-    sonic = "146"
-    linea = "59144"
-    worldchain = "480"
-    codex = "81224"
-    plume = "98866"
-    hyperevm = "999"
-
-_WRAPPED_GAS_ADDRESSES = {
-    ChainID.ethereum: "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
-    ChainID.unichain: "0x4200000000000000000000000000000000000006",
-    ChainID.optimism: "0x4200000000000000000000000000000000000006",
-    ChainID.arbitrum: "0x82aF49447D8a07e3bd95BD0d56f35241523fBab1",
-    ChainID.avalanche: "0xB31f66AA3C1e785363F0875A1B74E27b85FD66c7",
-    ChainID.polygon: "0x0d500B1d8E8eF31E21C99d1Db9A6444d3ADf1270",
-    ChainID.base: "0x4200000000000000000000000000000000000006",
-    ChainID.bsc: "0xbb4CdB9CBd36B01bD1cBaEBF2De08d9173bc095c",
-    ChainID.blast: "0x4300000000000000000000000000000000000004",
-    ChainID.mode: "0x4200000000000000000000000000000000000006",
-    ChainID.solana: "So11111111111111111111111111111111111111112",
-    ChainID.sonic: "0x039e2fB66102314Ce7b64Ce5Ce3E5183bc94aD38",
-    ChainID.linea: "0xe5D7C2a44FfDDf6b295A15c148167daaAf5Cf34f",
-    ChainID.sei: "0xE30feDd158A2e3b13e9badaeABaFc5516e95e8C7",
-    ChainID.worldchain: "0x4200000000000000000000000000000000000006",
-    ChainID.plume: "0xEa237441c92CAe6FC17Caaf9a7acB3f953be4bd1",
-    ChainID.hyperevm: "0x5555555555555555555555555555555555555555",
-}
-
-_PERMIT2_CONTRACT_ADDRESSES = {
-    ChainID.ethereum: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.avalanche: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.arbitrum: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.base: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.bsc: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.optimism: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.polygon: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.unichain: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.linea: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.worldchain: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.sonic: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.sei: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.blast: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.mode: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.scroll: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.plume: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-    ChainID.hyperevm: "0x000000000022D473030F116dDEE9F6B43aC78BA3",
-}
+from tristero.data import get_gas_addr
 
 class APIException(Exception):
     pass
@@ -99,21 +34,17 @@ async def t_post(client: httpx.AsyncClient, url: str, body: Optional[dict[str, A
     resp = await client.post(url, headers=get_config().headers, json=body, timeout=20)
     return handle_resp(resp)
 
+def or_native(chain_id: str, address: str):
+    return get_gas_addr(chain_id) if address == "native" else address
 
-def native_address(chain_id: ChainID) -> str:
-    return "0x0000000000000000000000000000000000000000"
-
-
-def or_native(chain_id: ChainID, address: str):
-    return native_address(chain_id) if address == "native" else address
-
+c = httpx.AsyncClient()
 
 async def get_quote(
     from_wallet: str,
     to_wallet: str,
-    from_chain_id: ChainID,
+    from_chain_id: str,
     from_address: str,
-    to_chain_id: ChainID,
+    to_chain_id: str,
     to_address: str,
     amount: int,
 ):
@@ -123,36 +54,51 @@ async def get_quote(
     from_token_address = or_native(from_chain_id, from_address)
     to_token_address = or_native(to_chain_id, to_address)
 
-    async with httpx.AsyncClient() as c:
-        data = {
-            "src_chain_id": from_chain_id.value,
-            "src_token_address": from_token_address,
-            "src_token_quantity": str(int(amount)),
-            "src_wallet_address": from_wallet,
-            "dst_chain_id": to_chain_id.value,
-            "dst_token_address": to_token_address,
-            "dst_wallet_address": to_wallet,
-        }
-        resp = await c.post(
-            get_config().quoter_url,
-            json=data,
-        )
-        # print(data, resp)
-        try:
-            return handle_resp(resp)
-        except Exception as e:
-            raise QuoteException(e, data) from e
-
+    data = {
+        "src_chain_id": from_chain_id,
+        "src_token_address": from_token_address,
+        "src_token_quantity": str(int(amount)),
+        "src_wallet_address": from_wallet,
+        "dst_chain_id": to_chain_id,
+        "dst_token_address": to_token_address,
+        "dst_wallet_address": to_wallet,
+    }
+    resp = await c.post(
+        get_config().quoter_url,
+        json=data,
+    )
+    # print(data, resp)
+    try:
+        return handle_resp(resp)
+    except Exception as e:
+        raise QuoteException(e, data) from e
 
 async def fill_order(signature: str, domain: dict[str, Any], message: dict[str, Any]):
-    async with httpx.AsyncClient() as c:
-        data = {"signature": signature, "domain": domain, "message": message}
-        resp = await c.post(
-            get_config().filler_url,
-            json=data,
-        )
-        return handle_resp(resp)
+    data = {"signature": signature, "domain": domain, "message": message}
+    resp = await c.post(
+        get_config().filler_url,
+        json=data,
+    )
+    return handle_resp(resp)
+
+async def fill_order_feather(src_chain: str, dst_chain: str, dst_address: str, amount: int, client_id: str = ''):
+    data = {
+        "client_id": client_id,
+        "src_chain": src_chain,
+        "dst_chain": dst_chain,
+        "dst_address": dst_address,
+        "amount": amount
+    }
+    resp = await c.post(
+        get_config().filler_url,
+        json=data,
+    )
+    return handle_resp(resp)
 
 async def poll_updates(order_id: str):
     ws = await connect(f"{get_config().ws_url}/{order_id}")
     return ws
+
+async def poll_updates_feather(order_id: str):
+    ws = await connect(f"{get_config().ws_url}/feather/{order_id}")
+    return ws

tristero/client.py

@@ -1,4 +1,5 @@
 import asyncio
+from dataclasses import dataclass
 import json
 import logging
 from typing import Any, Optional, TypeVar
@@ -6,10 +7,10 @@ from typing import Any, Optional, TypeVar
 from eth_account.signers.local import LocalAccount
 from pydantic import BaseModel
 
-from tristero.api import ChainID, fill_order, poll_updates
-from .permit2 import create_order
+from tristero.api import fill_order, fill_order_feather, get_quote, poll_updates, poll_updates_feather
+from .permit2 import Permit2Order, create_permit2_order
+from .data import ChainID
 from web3 import AsyncBaseProvider, AsyncWeb3
-import logging
 from web3 import AsyncWeb3
 from tenacity import (
     retry,
@@ -45,7 +46,36 @@ class TokenSpec(BaseModel, frozen=True):
     chain_id: ChainID
     token_address: str
 
-async def wait_for_completion(order_id: str):
+async def wait_for_feather_completion(order_id: str):
+    ws = await poll_updates_feather(order_id)
+    try:
+        async for msg in ws:
+            if not msg:
+                continue
+            msg = json.loads(msg)
+            status = msg['status']
+            logger.info(
+                {
+                    "message": f"status={status}",
+                    "id": "order_update",
+                    "payload": msg,
+                }
+            )
+            if status in ['Expired']:
+                await ws.close()
+                raise Exception(f"Swap failed: {ws.close_reason} {msg}")
+            elif status in ['Finalized']:
+                await ws.close()
+                return msg
+        # If we exit the loop without completed/failed, raise to retry
+        raise WebSocketClosedError("WebSocket closed without completion status")
+    except Exception as e:
+        # Close cleanly if still open
+        if not ws.close_code:
+            await ws.close()
+        raise
+
+async def wait_for_permit2_completion(order_id: str):
     ws = await poll_updates(order_id)
     try:
         async for msg in ws:
@@ -72,49 +102,69 @@ async def wait_for_completion(order_id: str):
             await ws.close()
         raise
 
+async def wait_for_completion(order_id: str, feather: bool):
+    if feather:
+        return await wait_for_feather_completion(order_id)
+    else:
+        return await wait_for_permit2_completion(order_id)
+
 @retry(
     stop=stop_after_attempt(3),
     wait=wait_exponential(multiplier=1, min=4, max=10),
     retry=retry_if_exception_type(WebSocketClosedError)
 )
-async def wait_for_completion_with_retry(order_id: str):
-    return await wait_for_completion(order_id)
-
-async def start_swap(w3: AsyncWeb3[P], account: LocalAccount, from_t: TokenSpec, to_t: TokenSpec, raw_amount: int) -> str:
-    """
-    Execute a token swap operation.
-
-    Args:
-        w3: Web3 provider instance
-        account: Account to execute swap from
-        from_t: Source token specification
-        to_t: Target token specification
-        raw_amount: Amount in smallest unit (e.g., wei)
-
-    Returns:
-        Order ID for tracking the swap
-
-    Raises:
-        Exception: If order creation or submission fails
-    """
-    data, sig = await create_order(
+async def wait_for_completion_with_retry(order_id: str, feather: bool = False):
+    return await wait_for_completion(order_id, feather)
+
+@dataclass
+class FeatherSwapResult:
+    deposit_address: str
+    data: Any
+
+class FeatherException(Exception):
+    pass
+
+async def start_feather_swap(
+    src_t: TokenSpec,
+    dst_t: TokenSpec,
+    dst_addr: str,
+    raw_amount: int,
+    client_id: str = ''
+):
+    resp = await fill_order_feather(client_id, str(src_t.chain_id.value), str(dst_t.chain_id), dst_addr, raw_amount)
+    if resp['detail']:
+        raise FeatherException(resp)
+    else:
+        return FeatherSwapResult(
+            deposit_address = resp['deposit_address'],
+            data = resp
+        )
+
+async def start_permit2_swap(
+    w3: AsyncWeb3[P],
+    account: LocalAccount,
+    src_t: TokenSpec,
+    dst_t: TokenSpec,
+    raw_amount: int,
+):
+    order = await create_permit2_order(
         w3,
         account,
-        from_t.chain_id,
-        from_t.token_address,
-        to_t.chain_id,
-        to_t.token_address,
+        str(src_t.chain_id.value),
+        src_t.token_address,
+        str(dst_t.chain_id.value),
+        dst_t.token_address,
         raw_amount,
     )
     response = await fill_order(
-        str(sig.signature.to_0x_hex()),
-        data.domain.model_dump(by_alias=True, mode="json"),
-        data.message.model_dump(by_alias=True, mode="json"),
+        str(order.sig.signature.to_0x_hex()),
+        order.msg.domain.model_dump(by_alias=True, mode="json"),
+        order.msg.message.model_dump(by_alias=True, mode="json"),
     )
+    order_id = response['id']
+    return order_id
 
-    return response['id']
-
-async def execute_swap(
+async def execute_permit2_swap(
     w3: AsyncWeb3[P],
     account: LocalAccount,
     src_t: TokenSpec,
@@ -124,20 +174,14 @@ async def execute_swap(
     timeout: Optional[float] = None
 ) -> dict[str, Any]:
     """Execute and wait for swap completion."""
-    order_id = await start_swap(
-        w3,
-        account,
-        src_t,
-        dst_t,
-        raw_amount
-    )
+    order_id = await start_permit2_swap(w3, account, src_t, dst_t, raw_amount)
     logger.info(f"Swap order placed: {order_id}")
 
     waiter = wait_for_completion_with_retry if retry else wait_for_completion
 
     try:
         if timeout is None:
-            return await waiter(order_id)
+            return await waiter(order_id, False)
 
         return await asyncio.wait_for(
             waiter(order_id),
@@ -145,3 +189,4 @@ async def execute_swap(
         )
     except asyncio.TimeoutError as exc:
         raise StuckException(f"Swap {order_id} timed out after {timeout}s") from exc
+

tristero/data.py

@@ -0,0 +1,35 @@
+from importlib import resources as impresources
+import json
+from typing import TYPE_CHECKING
+import enum
+
+CHAINS_FILE = impresources.files("tristero.files") / "chains.json"
+CHAINS = json.loads(CHAINS_FILE.read_text())
+
+def get_address(chain: str, id: str) -> str | None:
+    return CHAINS[chain].get('addresses').get(id)
+
+def get_permit2_addr(chain: str):
+    return get_address(chain, 'permit2')
+
+def get_wrapped_gas_addr(chain: str):
+    return get_address(chain, 'wrappedGasToken')
+
+def get_gas_addr(chain: str):
+    return get_address(chain, 'gasToken')
+
+def chain(id: str) -> str | None:
+    return CHAINS[str(id)].get('chainId')
+
+if TYPE_CHECKING:
+    from enum import IntEnum
+    class ChainID(IntEnum):
+        _value_: int
+else:
+    ChainID = enum.Enum(
+        'ChainID',
+        {
+            name: data['chainId']
+            for name, data in CHAINS.items()
+        }
+    )