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

tristero/data.py

@@ -6,20 +6,41 @@ 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 _chain_key(chain: str | int) -> str:
+    """Resolve a chain key from either a chain key (e.g. 'arbitrum') or a chain id (e.g. 42161/'42161')."""
+    s = str(chain)
+    if s in CHAINS:
+        return s
 
-def get_permit2_addr(chain: str):
+    try:
+        chain_id = int(s)
+    except ValueError as exc:
+        raise KeyError(f"Unknown chain: {chain}") from exc
+
+    for key, data in CHAINS.items():
+        if data.get("chainId") == chain_id:
+            return key
+
+    raise KeyError(f"Unknown chain id: {chain}")
+
+
+def get_address(chain: str | int, id: str) -> str | None:
+    key = _chain_key(chain)
+    addresses = CHAINS[key].get("addresses") or {}
+    return addresses.get(id)
+
+def get_permit2_addr(chain: str | int):
     return get_address(chain, 'permit2')
 
-def get_wrapped_gas_addr(chain: str):
+def get_wrapped_gas_addr(chain: str | int):
     return get_address(chain, 'wrappedGasToken')
 
-def get_gas_addr(chain: str):
+def get_gas_addr(chain: str | int):
     return get_address(chain, 'gasToken')
 
-def chain(id: str) -> str | None:
-    return CHAINS[str(id)].get('chainId')
+def chain(id: str | int) -> str | None:
+    """Return the chain key for a given chain id or chain key."""
+    return _chain_key(id)
 
 if TYPE_CHECKING:
     from enum import IntEnum

tristero/eip712/__init__.py

@@ -0,0 +1,48 @@
+from .simple_types import (
+    EIP712OrderParameters,
+    EIP712TokenPermissions,
+    EIP712Domain,
+    EIP712Permit,
+    EIP712MMWitness,
+    EIP712MarginWitness,
+    EIP712CloseWithSwap,
+)
+
+from .nested_types import (
+    EIP712SignedOrder,
+    EIP712PermitWitnessTransferFromOrder,
+    EIP712PermitWitnessTransferFromMM,
+    EIP712PermitWitnessTransferFromMargin,
+)
+
+from .eip712_struct import EIP712Struct
+from .eip712_auto import EIP712AutoStruct
+from .escrow_utils import get_escrow_domain, get_close_with_swap_types, ESCROW_NAME, ESCROW_VERSION
+from eth_utils import keccak
+
+
+def eip712_digest(domain: "EIP712Struct", message: "EIP712Struct") -> bytes:
+    prefix = b"\x19\x01"
+    return keccak(prefix + domain.hash() + message.hash())
+
+
+__all__ = [
+    "EIP712OrderParameters",
+    "EIP712TokenPermissions",
+    "EIP712SignedOrder",
+    "EIP712PermitWitnessTransferFromOrder",
+    "EIP712PermitWitnessTransferFromMM",
+    "EIP712PermitWitnessTransferFromMargin",
+    "EIP712MMWitness",
+    "EIP712MarginWitness",
+    "EIP712CloseWithSwap",
+    "EIP712Struct",
+    "EIP712Domain",
+    "EIP712Permit",
+    "EIP712AutoStruct",
+    "eip712_digest",
+    "get_escrow_domain",
+    "get_close_with_swap_types",
+    "ESCROW_NAME",
+    "ESCROW_VERSION",
+]

tristero/eip712/eip712_auto.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+from web3 import Web3
+
+import typing as _t
+from typing import get_args, get_origin, ClassVar
+from typing import Any, Dict, Type
+
+from eth_abi import encode as abi_encode
+from eth_utils import keccak
+from pydantic import BaseModel
+
+from .eip712_struct import EIP712Struct
+
+
+def _struct_string(cls: Type["EIP712Struct"]) -> str:
+    sol_types = getattr(cls, "SOL_TYPES", {})
+    if not sol_types:
+        return cls.TYPE_STRING
+
+    hints = _t.get_type_hints(cls)
+    parts: list[str] = []
+
+    for field, t in hints.items():
+        if get_origin(t) is ClassVar:
+            continue
+        if isinstance(t, type) and issubclass(t, EIP712Struct):
+            parts.append(f"{sol_types.get(field, t.SOL_TYPE_NAME)} {field}")
+        else:
+            parts.append(f"{sol_types[field]} {field}")
+
+    return f"{cls.SOL_TYPE_NAME}(" + ",".join(parts) + ")"
+
+
+def _gather_deps(
+    cls: Type[EIP712Struct], seen: set[Type[EIP712Struct]]
+) -> list[Type[EIP712Struct]]:
+    """Depth-first collect all user-defined nested structs (no duplicates)."""
+    deps: list[Type[EIP712Struct]] = []
+    for t in _t.get_type_hints(cls).values():
+        if get_origin(t) is ClassVar:
+            continue
+        if isinstance(t, type) and issubclass(t, EIP712Struct) and t not in seen:
+            seen.add(t)
+            deps.extend(_gather_deps(t, seen))  # recurse first
+            deps.append(t)  # then this struct itself
+    return deps
+
+
+def _canonical_type_string(cls: Type[EIP712Struct]) -> str:
+    """Root encodeType + ASCII-sorted dependencies (EIP-712 protocol definition)."""
+    deps = sorted(_gather_deps(cls, set()), key=lambda c: c.SOL_TYPE_NAME)
+    return _struct_string(cls) + "".join(_struct_string(d) for d in deps)
+
+
+class _AutoMeta(type(BaseModel)):
+    """Populate TYPE_STRING / TYPE_STRUCT at **class-definition** time."""
+
+    def __init__(cls, name, bases, ns, **kw):
+        super().__init__(name, bases, ns, **kw)
+
+        if cls.__name__ == "EIP712AutoStruct":
+            return
+
+        cls.TYPE_STRING = _canonical_type_string(cls)  # type: ignore[attr-defined]
+
+        hints = _t.get_type_hints(cls)
+        struct: list[Dict[str, str]] = []
+        for field, t in hints.items():
+            if get_origin(t) is ClassVar:
+                continue
+            if isinstance(t, type) and issubclass(t, EIP712Struct):
+                struct.append(
+                    {
+                        "name": field,
+                        "type": getattr(cls, "SOL_TYPES", {}).get(field, t.SOL_TYPE_NAME),
+                    }
+                )
+            else:
+                struct.append({"name": field, "type": cls.SOL_TYPES[field]})
+        cls.TYPE_STRUCT = struct  # type: ignore[attr-defined]
+
+
+def _coerce_bytes(val: Any) -> bytes:
+    """Accept raw bytes *or* '0x…' str/bytes and return binary bytes."""
+    if isinstance(val, (bytes, bytearray)):
+        b = bytes(val)
+        if b.startswith(b"0x"):
+            return Web3.to_bytes(hexstr=b.decode())
+        return b
+    if isinstance(val, str):
+        if val.startswith("0x"):
+            return Web3.to_bytes(hexstr=val)
+        return val.encode()
+    raise TypeError("Expected bytes-like or 0x-string")
+
+
+def _hash_fixed_bytes_array(n: int, seq: list[bytes]) -> bytes:
+    """Hash a ``bytesN[]`` value-type array (N ≤ 32)."""
+    packed = bytearray(32 * len(seq))
+    for i, raw in enumerate(seq):
+        b = _coerce_bytes(raw)
+        if len(b) != n:
+            raise ValueError(
+                f"bytes{n} element length mismatch: Expected {n} bytes, got {len(b)}, value: {b.hex()}"
+            )
+        packed[i * 32 : i * 32 + n] = b
+    return keccak(bytes(packed))
+
+
+def _hash_dynamic_bytes_array(seq: list[bytes]) -> bytes:
+    """Hash a ``bytes[]`` (or ``string[]``) array."""
+    packed = bytearray(32 * len(seq))
+    for i, raw in enumerate(seq):
+        packed[i * 32 : (i + 1) * 32] = keccak(_coerce_bytes(raw))
+    return keccak(bytes(packed))
+
+
+def _hash_value_type_array(base: str, seq: list[Any]) -> bytes:
+    packed = b"".join(abi_encode([base], [x]) for x in seq)
+    return keccak(packed)
+
+
+def _hash_array(sol_array_type: str, seq: list[Any]) -> bytes:
+    base = sol_array_type[:-2]
+
+    if base in ("bytes", "string"):
+        return _hash_dynamic_bytes_array(seq)
+
+    if base.startswith("bytes") and base != "bytes":
+        n = int(base[5:])
+        return _hash_fixed_bytes_array(n, seq)
+
+    return _hash_value_type_array(base, seq)
+
+
+class EIP712AutoStruct(EIP712Struct, metaclass=_AutoMeta):
+    SOL_TYPES: ClassVar[Dict[str, str]] = {}
+    model_config = {"arbitrary_types_allowed": True}
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "EIP712AutoStruct":
+        field_types = _t.get_type_hints(cls)
+        kwargs = {}
+
+        for field in cls.__annotations__.keys():
+            if get_origin(field_types.get(field)) is ClassVar:
+                continue
+
+            if field not in data:
+                continue
+
+            field_value = data[field]
+            field_type = field_types.get(field)
+
+            if isinstance(field_type, type) and issubclass(field_type, EIP712Struct):
+                kwargs[field] = field_type.from_dict(field_value)
+            else:
+                if (
+                    field_type is bytes
+                    and isinstance(field_value, str)
+                    and field_value.startswith("0x")
+                ):
+                    kwargs[field] = Web3.to_bytes(hexstr=field_value)
+                elif (
+                    field_type is str
+                    and cls.SOL_TYPES.get(field) == "address"
+                    and isinstance(field_value, str)
+                ):
+                    kwargs[field] = Web3.to_checksum_address(field_value)
+                elif (
+                    get_origin(field_type) in (list, _t.List)
+                    and get_args(field_type)[0] is bytes
+                ):
+                    cleaned = []
+                    for elem in field_value:
+                        if isinstance(elem, str) and elem.startswith("0x"):
+                            cleaned.append(Web3.to_bytes(hexstr=elem))
+                        else:
+                            cleaned.append(elem)
+                    kwargs[field] = cleaned
+                else:
+                    kwargs[field] = field_value
+
+        return cls(**kwargs)
+
+    def hash(self) -> bytes:
+        abi_types: list[str] = ["bytes32"]
+        abi_vals: list[Any] = [self.type_hash]
+
+        for field in self.__class__.__annotations__.keys():
+            sol_type = self.SOL_TYPES[field]
+            val = getattr(self, field)
+
+            if isinstance(val, EIP712Struct):
+                abi_types.append("bytes32")
+                abi_vals.append(val.hash())
+                continue
+
+            if sol_type in ("bytes", "string"):
+                abi_types.append("bytes32")
+                abi_vals.append(keccak(_coerce_bytes(val)))
+                continue
+
+            if sol_type.endswith("[]"):
+                abi_types.append("bytes32")
+                abi_vals.append(_hash_array(sol_type, val))
+                continue
+
+            abi_types.append(sol_type)
+            abi_vals.append(val)
+
+        return keccak(abi_encode(abi_types, abi_vals))

tristero/eip712/eip712_struct.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, ClassVar
+from pydantic import BaseModel
+from web3 import Web3
+from eth_utils import keccak
+
+
+class EIP712Struct(ABC, BaseModel):
+    SOL_TYPE_NAME: ClassVar[str]
+    TYPE_STRING: ClassVar[str]
+    TYPE_STRUCT: ClassVar[List[Dict]]
+
+    @property
+    def sol_type_name(self) -> str:
+        return self.__class__.SOL_TYPE_NAME
+
+    @property
+    def type_string(self) -> str:
+        return self.__class__.TYPE_STRING
+
+    @property
+    def type_struct(self) -> List[Dict]:
+        return self.__class__.TYPE_STRUCT
+
+    @property
+    def type_hash(self) -> bytes:
+        return keccak(text=self.__class__.TYPE_STRING)
+
+    @abstractmethod
+    def hash(self) -> bytes:
+        pass
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "EIP712Struct":
+        field_names = cls.__annotations__.keys()
+        kwargs = {field: data.get(field) for field in field_names if field in data}
+
+        for field, value in kwargs.items():
+            field_type = cls.__annotations__.get(field, None)
+
+            if (
+                isinstance(value, dict)
+                and hasattr(field_type, "__origin__")
+                and issubclass(field_type.__origin__, EIP712Struct)
+            ):
+                kwargs[field] = field_type.from_dict(value)
+
+            elif (
+                isinstance(value, dict)
+                and isinstance(field_type, type)
+                and issubclass(field_type, EIP712Struct)
+            ):
+                kwargs[field] = field_type.from_dict(value)
+
+        return cls(**kwargs)
+
+    @staticmethod
+    def _clean(value: Any):
+        if isinstance(value, (bytes, bytearray)):
+            return Web3.to_hex(value)
+        if isinstance(value, EIP712Struct):
+            return value.eip_signable_struct
+        if isinstance(value, (list, tuple)):
+            return [EIP712Struct._clean(v) for v in value]
+        if isinstance(value, dict):
+            return {k: EIP712Struct._clean(v) for k, v in value.items()}
+        return value
+
+    @property
+    def eip_signable_struct(self) -> Dict[str, Any]:
+        raw = self.model_dump(mode="python", exclude_none=True)
+        return {k: self._clean(v) for k, v in raw.items()}

tristero/eip712/escrow_utils.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from eth_utils import to_checksum_address
+
+from .simple_types import EIP712CloseWithSwap
+
+
+ESCROW_NAME = "Tristero Margin Position"
+ESCROW_VERSION = "3"
+
+
+def get_escrow_domain(chain_id: int, verifying_contract: str) -> Dict[str, Any]:
+    return {
+        "name": ESCROW_NAME,
+        "version": ESCROW_VERSION,
+        "chainId": int(chain_id),
+        "verifyingContract": to_checksum_address(verifying_contract),
+    }
+
+
+def get_close_with_swap_types() -> Dict[str, Any]:
+    return {
+        "EIP712Domain": [
+            {"name": "name", "type": "string"},
+            {"name": "version", "type": "string"},
+            {"name": "chainId", "type": "uint256"},
+            {"name": "verifyingContract", "type": "address"},
+        ],
+        "CloseWithSwap": EIP712CloseWithSwap.TYPE_STRUCT,
+    }

tristero/eip712/nested_types.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+from typing import List
+from .eip712_auto import EIP712AutoStruct
+from .simple_types import (
+    EIP712OrderParameters,
+    EIP712TokenPermissions,
+    EIP712MMWitness,
+    EIP712MarginWitness,
+)
+
+
+class EIP712SignedOrder(EIP712AutoStruct):
+    SOL_TYPE_NAME = "SignedOrder"
+
+    sender: str
+    parameters: EIP712OrderParameters
+    deadline: int
+    target: str
+    filler: str
+    orderType: str
+    customData: List[bytes]
+
+    SOL_TYPES = {
+        "sender": "address",
+        "parameters": "OrderParameters",
+        "deadline": "uint256",
+        "target": "address",
+        "filler": "address",
+        "orderType": "string",
+        "customData": "bytes[]",
+    }
+
+
+class EIP712PermitWitnessTransferFromOrder(EIP712AutoStruct):
+    SOL_TYPE_NAME = "PermitWitnessTransferFrom"
+
+    permitted: EIP712TokenPermissions
+    spender: str
+    nonce: int
+    deadline: int
+    witness: EIP712SignedOrder
+
+    SOL_TYPES = {
+        "permitted": "TokenPermissions",
+        "spender": "address",
+        "nonce": "uint256",
+        "deadline": "uint256",
+        "witness": "SignedOrder",
+    }
+
+
+class EIP712PermitWitnessTransferFromMM(EIP712AutoStruct):
+    SOL_TYPE_NAME = "PermitWitnessTransferFrom"
+
+    permitted: EIP712TokenPermissions
+    spender: str
+    nonce: int
+    deadline: int
+    witness: EIP712MMWitness
+
+    SOL_TYPES = {
+        "permitted": "TokenPermissions",
+        "spender": "address",
+        "nonce": "uint256",
+        "deadline": "uint256",
+        "witness": "MMWitness",
+    }
+
+
+class EIP712PermitWitnessTransferFromMargin(EIP712AutoStruct):
+    SOL_TYPE_NAME = "PermitWitnessTransferFrom"
+
+    permitted: EIP712TokenPermissions
+    spender: str
+    nonce: int
+    deadline: int
+    witness: EIP712MarginWitness
+
+    SOL_TYPES = {
+        "permitted": "TokenPermissions",
+        "spender": "address",
+        "nonce": "uint256",
+        "deadline": "uint256",
+        "witness": "MarginWitness",
+    }