Add redeemer stuff (#184)

* Add memoize

* Add async_batched

* Add os token converter

* Add META_VAULT_ID

* Fix lint

* Fix func metadata

Author: evgeny-stakewise

sw_utils/__init__.py

@@ -1,4 +1,4 @@
-from .common import InterruptHandler
+from .common import InterruptHandler, async_batched
 from .consensus import (
     PENDING_STATUSES,
     ExtendedAsyncBeacon,
@@ -9,7 +9,7 @@
     get_chain_latest_head,
     get_consensus_client,
 )
-from .decorators import retry_aiohttp_errors, retry_ipfs_exception, safe
+from .decorators import memoize, retry_aiohttp_errors, retry_ipfs_exception, safe
 from .event_scanner import EventProcessor, EventScanner
 from .exceptions import IpfsException
 from .execution import GasManager, get_execution_client
@@ -30,6 +30,7 @@
     NETWORKS,
     BaseNetworkConfig,
 )
+from .os_token_converter import OsTokenConverter
 from .password import generate_password
 from .protocol_config import build_protocol_config
 from .signing import (

sw_utils/common.py

@@ -1,9 +1,11 @@
 import asyncio
 import logging
 import signal
-from typing import Any
+from typing import Any, AsyncGenerator, TypeVar
 from urllib.parse import urlparse, urlunparse
 
+T = TypeVar('T')
+
 logger = logging.getLogger(__name__)
 
 
@@ -65,3 +67,17 @@ def urljoin(base: str, *args: str) -> str:
 
 def _join_paths(*args: str) -> str:
     return '/'.join(str(x).strip('/') for x in args)
+
+
+async def async_batched(
+    async_gen: AsyncGenerator[T, None], batch_size: int
+) -> AsyncGenerator[list[T], None]:
+    """Batch items from an async iterator. Replacement for itertools.batched."""
+    batch = []
+    async for item in async_gen:
+        batch.append(item)
+        if len(batch) == batch_size:
+            yield batch
+            batch = []
+    if batch:
+        yield batch

sw_utils/decorators.py

@@ -18,6 +18,36 @@
 default_logger = logging.getLogger(__name__)
 
 
+def memoize(func: Callable) -> Callable:
+    """
+    Helper to memoize both sync and async functions.
+    Main usage is for async functions because `functools.cache` won't work with them.
+    """
+    cache: dict = {}
+
+    @wraps(func)
+    async def memoized_async_func(*args, **kwargs):  # type: ignore
+        key = (args, frozenset(sorted(kwargs.items())))
+        if key in cache:
+            return cache[key]
+        result = await func(*args, **kwargs)
+        cache[key] = result
+        return result
+
+    @wraps(func)
+    def memoized_sync_func(*args, **kwargs):  # type: ignore
+        key = (args, frozenset(sorted(kwargs.items())))
+        if key in cache:
+            return cache[key]
+        result = func(*args, **kwargs)
+        cache[key] = result
+        return result
+
+    if asyncio.iscoroutinefunction(func):
+        return memoized_async_func
+    return memoized_sync_func
+
+
 def safe(func: Callable) -> Callable:
     if asyncio.iscoroutinefunction(func):
 

sw_utils/networks.py

@@ -50,6 +50,7 @@ class BaseNetworkConfig:
     PECTRA_VAULT_VERSION: int
     OS_TOKEN_VAULT_CONTROLLER_CONTRACT_ADDRESS: ChecksumAddress
     MIN_EFFECTIVE_PRIORITY_FEE_PER_GAS: Wei
+    META_VAULT_ID: HexStr
 
     @property
     def SECONDS_PER_BLOCK(self) -> int:
@@ -128,6 +129,7 @@ def PECTRA_SLOT(self) -> int:
             '0x2A261e60FB14586B474C208b1B7AC6D0f5000306'
         ),
         MIN_EFFECTIVE_PRIORITY_FEE_PER_GAS=Web3.to_wei(0, 'gwei'),
+        META_VAULT_ID=HexStr('0xcfece609e9557b5f0b085dadb8b7c99d43a9052d28ab3d68c9e3c1a9c3ab85c0'),
     ),
     HOODI: BaseNetworkConfig(
         SLOTS_PER_EPOCH=32,
@@ -179,6 +181,7 @@ def PECTRA_SLOT(self) -> int:
             '0x140Fc69Eabd77fFF91d9852B612B2323256f7Ac1'
         ),
         MIN_EFFECTIVE_PRIORITY_FEE_PER_GAS=Web3.to_wei(0, 'gwei'),
+        META_VAULT_ID=HexStr('0xcfece609e9557b5f0b085dadb8b7c99d43a9052d28ab3d68c9e3c1a9c3ab85c0'),
     ),
     GNOSIS: BaseNetworkConfig(
         SLOTS_PER_EPOCH=16,
@@ -234,5 +237,6 @@ def PECTRA_SLOT(self) -> int:
             '0x60B2053d7f2a0bBa70fe6CDd88FB47b579B9179a'
         ),
         MIN_EFFECTIVE_PRIORITY_FEE_PER_GAS=Web3.to_wei(1, 'gwei'),
+        META_VAULT_ID=HexStr('0xfb5cee5ecc2ff8d1a7a5ad55f89156551c040a926bec689720ab06063922454d'),
     ),
 }

sw_utils/os_token_converter.py

@@ -0,0 +1,22 @@
+from web3.types import Wei
+
+
+class OsTokenConverter:
+    """
+    Convert between shares and assets based on total assets and total shares.
+    Helps to avoid repeating calls to the contract.
+    """
+
+    def __init__(self, total_assets: Wei, total_shares: Wei):
+        self.total_assets = total_assets
+        self.total_shares = total_shares
+
+    def to_shares(self, assets: Wei) -> Wei:
+        if self.total_assets == 0:
+            return Wei(0)
+        return Wei((assets * self.total_shares) // self.total_assets)
+
+    def to_assets(self, shares: Wei) -> Wei:
+        if self.total_shares == 0:
+            return Wei(0)
+        return Wei((shares * self.total_assets) // self.total_shares)

sw_utils/tests/test_os_token_converter.py

@@ -0,0 +1,34 @@
+import pytest
+from web3.types import Wei
+
+from sw_utils.os_token_converter import OsTokenConverter
+
+
+class TestOsTokenConverter:
+    @pytest.mark.parametrize(
+        'total_assets,total_shares,assets,expected_shares',
+        [
+            (Wei(1000), Wei(100), Wei(500), Wei(50)),
+            (Wei(2000), Wei(200), Wei(1000), Wei(100)),
+            (Wei(0), Wei(100), Wei(500), Wei(0)),  # Edge case: total_assets is 0
+            (Wei(1000), Wei(0), Wei(500), Wei(0)),  # Edge case: total_shares is 0
+        ],
+    )
+    def test_to_shares(self, total_assets, total_shares, assets, expected_shares):
+        converter = OsTokenConverter(total_assets=total_assets, total_shares=total_shares)
+        shares = converter.to_shares(assets)
+        assert shares == expected_shares
+
+    @pytest.mark.parametrize(
+        'total_assets,total_shares,shares,expected_assets',
+        [
+            (Wei(1000), Wei(100), Wei(50), Wei(500)),
+            (Wei(2000), Wei(200), Wei(100), Wei(1000)),
+            (Wei(1000), Wei(0), Wei(50), Wei(0)),  # Edge case: total_shares is 0
+            (Wei(0), Wei(100), Wei(50), Wei(0)),  # Edge case: total_assets is 0
+        ],
+    )
+    def test_to_assets(self, total_assets, total_shares, shares, expected_assets):
+        converter = OsTokenConverter(total_assets=total_assets, total_shares=total_shares)
+        assets = converter.to_assets(shares)
+        assert assets == expected_assets