dao-treasury 0.0.10__cp310-cp310-win32.whl → 0.0.70__cp310-cp310-win32.whl

dao_treasury/streams/llamapay.py

@@ -0,0 +1,388 @@
+import asyncio
+import datetime as dt
+import decimal
+from logging import getLogger
+from typing import (
+    Awaitable,
+    Callable,
+    Dict,
+    Final,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    final,
+)
+
+import dank_mids
+import pony.orm
+from a_sync import AsyncThreadPoolExecutor, igather
+from brownie.network.event import _EventItem
+from eth_typing import BlockNumber, ChecksumAddress, HexAddress, HexStr
+from tqdm.asyncio import tqdm_asyncio
+
+import y
+from y.time import NoBlockFound, UnixTimestamp
+from y.utils.events import decode_logs, get_logs_asap
+
+from dao_treasury import constants
+from dao_treasury.db import (
+    Stream,
+    StreamedFunds,
+    Address,
+    Token,
+    must_sort_outbound_txgroup_dbid,
+)
+from dao_treasury._wallet import TreasuryWallet
+
+
+logger: Final = getLogger(__name__)
+
+_UTC: Final = dt.timezone.utc
+
+_ONE_DAY: Final = 60 * 60 * 24
+
+_STREAMS_THREAD: Final = AsyncThreadPoolExecutor(1)
+
+create_task: Final = asyncio.create_task
+sleep: Final = asyncio.sleep
+
+datetime: Final = dt.datetime
+timedelta: Final = dt.timedelta
+fromtimestamp: Final = datetime.fromtimestamp
+now: Final = datetime.now
+
+Decimal: Final = decimal.Decimal
+
+ObjectNotFound: Final = pony.orm.ObjectNotFound
+commit: Final = pony.orm.commit
+db_session: Final = pony.orm.db_session
+
+Contract: Final = y.Contract
+Network: Final = y.Network
+get_block_at_timestamp: Final = y.get_block_at_timestamp
+get_price: Final = y.get_price
+
+
+networks: Final = [Network.Mainnet]
+
+factories: List[HexAddress] = []
+
+if dai_stream_factory := {
+    Network.Mainnet: "0x60c7B0c5B3a4Dc8C690b074727a17fF7aA287Ff2",
+}.get(constants.CHAINID):
+    factories.append(dai_stream_factory)
+
+if yfi_stream_factory := {
+    Network.Mainnet: "0xf3764eC89B1ad20A31ed633b1466363FAc1741c4",
+}.get(constants.CHAINID):
+    factories.append(yfi_stream_factory)
+
+
+def _generate_dates(
+    start: dt.datetime, end: dt.datetime, stop_at_today: bool = True
+) -> Iterator[dt.datetime]:
+    current = start
+    while current < end:
+        yield current
+        current += timedelta(days=1)
+        if stop_at_today and current.date() > now(_UTC).date():
+            break
+
+
+_StreamToStart = Callable[[HexStr, Optional[BlockNumber]], Awaitable[int]]
+
+_streamToStart_cache: Final[Dict[HexStr, _StreamToStart]] = {}
+
+
+def _get_streamToStart(stream_id: HexStr) -> _StreamToStart:
+    if streamToStart := _streamToStart_cache.get(stream_id):
+        return streamToStart
+    with db_session:
+        contract: y.Contract = Stream[stream_id].contract.contract  # type: ignore [misc]
+    streamToStart = contract.streamToStart.coroutine
+    _streamToStart_cache[stream_id] = streamToStart
+    return streamToStart
+
+
+async def _get_start_timestamp(
+    stream_id: HexStr, block: Optional[BlockNumber] = None
+) -> int:
+    streamToStart = _streamToStart_cache.get(stream_id)
+    if streamToStart is None:
+        streamToStart = await _STREAMS_THREAD.run(_get_streamToStart, stream_id)
+    # try:
+    return int(await streamToStart(f"0x{stream_id}", block_identifier=block))  # type: ignore [call-arg]
+    # except Exception:
+    #    return 0
+
+
+def _pause_stream(stream_id: HexStr) -> None:
+    with db_session:
+        Stream[stream_id].pause()  # type: ignore [misc]
+
+
+def _stop_stream(stream_id: str, block: BlockNumber) -> None:
+    with db_session:
+        Stream[stream_id].stop_stream(block)  # type: ignore [misc]
+
+
+_block_timestamps: Final[Dict[BlockNumber, UnixTimestamp]] = {}
+
+
+async def _get_block_timestamp(block: BlockNumber) -> UnixTimestamp:
+    if timestamp := _block_timestamps.get(block):
+        return timestamp
+    timestamp = await dank_mids.eth.get_block_timestamp(block)
+    _block_timestamps[block] = timestamp
+    return timestamp
+
+
+"""
+class _StreamProcessor(ABC):
+    @abstractmethod
+    async def _load_streams(self) -> None:
+        ...
+"""
+
+
+@final
+class LlamaPayProcessor:
+    """
+    Generalized async processor for DAO stream contracts.
+    Args are passed in at construction time.
+    Supports time-bounded admin periods for filtering.
+    """
+
+    handled_events: Final = (
+        "StreamCreated",
+        "StreamCreatedWithReason",
+        "StreamModified",
+        "StreamPaused",
+        "StreamCancelled",
+    )
+    skipped_events: Final = (
+        "PayerDeposit",
+        "PayerWithdraw",
+        "Withdraw",
+    )
+
+    def __init__(self) -> None:
+        self.stream_contracts: Final = {Contract(addr) for addr in factories}
+
+    async def _get_streams(self) -> None:
+        await igather(
+            self._load_contract_events(stream_contract)
+            for stream_contract in self.stream_contracts
+        )
+
+    async def _load_contract_events(self, stream_contract: y.Contract) -> None:
+        events = decode_logs(
+            await get_logs_asap(stream_contract.address, None, sync=False)
+        )
+        keys: Set[str] = set(events.keys())
+        for k in keys:
+            if k not in self.handled_events and k not in self.skipped_events:
+                raise NotImplementedError(f"Need to handle event: {k}")
+
+        if "StreamCreated" in keys:
+            for event in events["StreamCreated"]:
+                from_address, *_ = event.values()
+                from_address = Address.get_or_insert(from_address).address
+                if not TreasuryWallet.check_membership(
+                    from_address, event.block_number
+                ):
+                    continue
+                await _STREAMS_THREAD.run(self._get_stream, event)
+
+        if "StreamCreatedWithReason" in keys:
+            for event in events["StreamCreatedWithReason"]:
+                from_address, *_ = event.values()
+                from_address = Address.get_or_insert(from_address).address
+                if not TreasuryWallet.check_membership(
+                    from_address, event.block_number
+                ):
+                    continue
+                await _STREAMS_THREAD.run(self._get_stream, event)
+
+        if "StreamModified" in keys:
+            for event in events["StreamModified"]:
+                from_address, _, _, old_stream_id, *_ = event.values()
+                if not TreasuryWallet.check_membership(
+                    from_address, event.block_number
+                ):
+                    continue
+                await _STREAMS_THREAD.run(
+                    _stop_stream, old_stream_id.hex(), event.block_number
+                )
+                await _STREAMS_THREAD.run(self._get_stream, event)
+
+        if "StreamPaused" in keys:
+            for event in events["StreamPaused"]:
+                from_address, *_, stream_id = event.values()
+                if not TreasuryWallet.check_membership(
+                    from_address, event.block_number
+                ):
+                    continue
+                await _STREAMS_THREAD.run(_pause_stream, stream_id.hex())
+
+        if "StreamCancelled" in keys:
+            for event in events["StreamCancelled"]:
+                from_address, *_, stream_id = event.values()
+                if not TreasuryWallet.check_membership(
+                    from_address, event.block_number
+                ):
+                    continue
+                await _STREAMS_THREAD.run(
+                    _stop_stream, stream_id.hex(), event.block_number
+                )
+
+    def _get_stream(self, log: _EventItem) -> Stream:
+        with db_session:
+            if log.name == "StreamCreated":
+                from_address, to_address, amount_per_second, stream_id = log.values()
+                reason = None
+            elif log.name == "StreamCreatedWithReason":
+                from_address, to_address, amount_per_second, stream_id, reason = (
+                    log.values()
+                )
+            elif log.name == "StreamModified":
+                (
+                    from_address,
+                    _,
+                    _,
+                    old_stream_id,
+                    to_address,
+                    amount_per_second,
+                    stream_id,
+                ) = log.values()
+                reason = Stream[old_stream_id.hex()].reason  # type: ignore [misc]
+            else:
+                raise NotImplementedError("This is not an appropriate event log.")
+
+            stream_id_hex = stream_id.hex()
+            try:
+                return Stream[stream_id_hex]  # type: ignore [misc]
+            except ObjectNotFound:
+                entity = Stream(
+                    stream_id=stream_id_hex,
+                    contract=Address.get_dbid(log.address),
+                    start_block=log.block_number,
+                    token=Token.get_dbid(Contract(log.address).token()),
+                    from_address=Address.get_dbid(from_address),
+                    to_address=Address.get_dbid(to_address),
+                    amount_per_second=amount_per_second,
+                    txgroup=must_sort_outbound_txgroup_dbid,
+                )
+                if reason is not None:
+                    entity.reason = reason
+                commit()
+                return entity
+
+    def streams_for_recipient(
+        self, recipient: ChecksumAddress, at_block: Optional[BlockNumber] = None
+    ) -> List[Stream]:
+        with db_session:
+            streams = Stream.select(lambda s: s.to_address.address == recipient)
+            if at_block is None:
+                return list(streams)
+            return [
+                s for s in streams if (s.end_block is None or at_block <= s.end_block)
+            ]
+
+    def streams_for_token(
+        self, token: ChecksumAddress, include_inactive: bool = False
+    ) -> List[Stream]:
+        with db_session:
+            streams = Stream.select(lambda s: s.token.address.address == token)
+            return (
+                list(streams)
+                if include_inactive
+                else [s for s in streams if s.is_alive]
+            )
+
+    async def process_streams(self, run_forever: bool = False) -> None:
+        logger.info("Processing stream events and streamed funds...")
+        # Always sync events before processing
+        await self._get_streams()
+        with db_session:
+            streams = [s.stream_id for s in Stream.select()]
+        await tqdm_asyncio.gather(
+            *(
+                self.process_stream(stream_id, run_forever=run_forever)
+                for stream_id in streams
+            ),
+            desc="LlamaPay Streams",
+        )
+
+    async def process_stream(
+        self, stream_id: HexStr, run_forever: bool = False
+    ) -> None:
+        start, end = await _STREAMS_THREAD.run(Stream._get_start_and_end, stream_id)
+        for date_obj in _generate_dates(start, end, stop_at_today=not run_forever):
+            if await self.process_stream_for_date(stream_id, date_obj) is None:
+                return
+
+    async def process_stream_for_date(
+        self, stream_id: HexStr, date_obj: dt.datetime
+    ) -> Optional[StreamedFunds]:
+        entity = await _STREAMS_THREAD.run(
+            StreamedFunds.get_entity, stream_id, date_obj
+        )
+        if entity:
+            return entity
+
+        stream_token, start_date = await _STREAMS_THREAD.run(
+            Stream._get_token_and_start_date, stream_id
+        )
+        check_at = date_obj + timedelta(days=1) - timedelta(seconds=1)
+        if check_at > now(tz=_UTC):
+            await sleep((check_at - now(tz=_UTC)).total_seconds())
+
+        while True:
+            try:
+                block = await get_block_at_timestamp(check_at, sync=False)
+            except NoBlockFound:
+                sleep_time = (check_at - now(tz=_UTC)).total_seconds()
+                logger.debug(
+                    "no block found for %s, sleeping %ss", check_at, sleep_time
+                )
+                await sleep(sleep_time)
+            else:
+                break
+
+        price_fut = create_task(get_price(stream_token, block, sync=False))
+        start_timestamp = await _get_start_timestamp(stream_id, block)
+        if start_timestamp == 0:
+            if await _STREAMS_THREAD.run(Stream.check_closed, stream_id):
+                price_fut.cancel()
+                return None
+
+            while start_timestamp == 0:
+                block -= 1
+                start_timestamp = await _get_start_timestamp(stream_id, block)
+
+            block_datetime = fromtimestamp(await _get_block_timestamp(block), tz=_UTC)
+            assert block_datetime.date() == date_obj.date()
+            seconds_active = (check_at - block_datetime).seconds
+            is_last_day = True
+        else:
+            seconds_active = int(check_at.timestamp()) - start_timestamp
+            is_last_day = False
+
+        seconds_active_today = min(seconds_active, _ONE_DAY)
+        if seconds_active_today < _ONE_DAY and not is_last_day:
+            if date_obj.date() != start_date:
+                seconds_active_today = _ONE_DAY
+
+        with db_session:
+            price = Decimal(await price_fut)
+            entity = await _STREAMS_THREAD.run(
+                StreamedFunds.create_entity,
+                stream_id,
+                date_obj,
+                price,
+                seconds_active_today,
+                is_last_day,
+            )
+            return entity

dao_treasury/treasury.py

@@ -1,19 +1,39 @@
-from asyncio import create_task
+"""Treasury orchestration and analytics interface.
+
+This module defines the Treasury class, which aggregates DAO wallets, sets up
+sorting rules, and manages transaction ingestion and streaming analytics.
+It coordinates the end-to-end flow from wallet configuration to database
+population and dashboard analytics.
+
+Key Responsibilities:
+    - Aggregate and manage DAO-controlled wallets.
+    - Ingest and process on-chain transactions.
+    - Apply sorting/categorization rules.
+    - Integrate with streaming protocols (e.g., LlamaPay).
+    - Populate the database for analytics and dashboards.
+
+This is the main entry point for orchestrating DAO treasury analytics.
+"""
+
+from asyncio import create_task, gather
 from logging import getLogger
 from pathlib import Path
-from typing import Final, Iterable, List, Optional, Union
+from typing import Dict, Final, Iterable, List, Optional, Union
 
 import a_sync
 from a_sync.a_sync.abstract import ASyncABC
-from eth_typing import BlockNumber
+from eth_typing import BlockNumber, HexAddress
 from eth_portfolio.structs import LedgerEntry
 from eth_portfolio.typing import PortfolioBalances
 from eth_portfolio_scripts._portfolio import ExportablePortfolio
+from pony.orm import db_session
 from tqdm.asyncio import tqdm_asyncio
 
 from dao_treasury._wallet import TreasuryWallet
+from dao_treasury.constants import CHAINID
 from dao_treasury.db import TreasuryTx
 from dao_treasury.sorting._rules import Rules
+from dao_treasury.streams import llamapay
 
 
 Wallet = Union[TreasuryWallet, str]
@@ -32,19 +52,58 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
         sort_rules: Optional[Path] = None,
         start_block: int = 0,
         label: str = "your org's treasury",
+        custom_buckets: Optional[Dict[HexAddress, str]] = None,
         asynchronous: bool = False,
     ) -> None:
-        """
+        """Initialize the Treasury singleton for managing DAO funds.
+
+        This class aggregates multiple treasury wallets, sets up sorting rules,
+        and constructs an :class:`.ExportablePortfolio` for fetching balance and
+        transaction history.
+
         Args:
-            wallets: Iterable[Union[TreasuryWallet, str]]
+            wallets: Iterable of wallet
+                addresses or :class:`.TreasuryWallet` instances representing
+                DAO-controlled wallets.
+            sort_rules: Directory path containing YAML rule files
+                for sorting transactions. See :class:`dao_treasury.sorting._rules.Rules`.
+            start_block: Block number from which to start loading portfolio
+                history.
+            label: Descriptive label for the portfolio, used in exported data.
+            asynchronous: Whether methods default to asynchronous mode.
+
+        Raises:
+            RuntimeError: If a second Treasury instance is initialized.
+            TypeError: If any item in `wallets` is not a str or TreasuryWallet.
+
+        Examples:
+            .. code-block:: python
+
+                # Create a synchronous Treasury
+                treasury = Treasury(
+                    wallets=["0xAbc123...", TreasuryWallet("0xDef456...", start_block=1000)],
+                    sort_rules=Path("/path/to/rules"),
+                    start_block=500,
+                    label="DAO Treasury",
+                    asynchronous=False
+                )
 
+                # Create an asynchronous Treasury
+                treasury_async = Treasury(
+                    wallets=["0xAbc123..."],
+                    asynchronous=True
+                )
         """
         global TREASURY
         if TREASURY is not None:
-            raise RuntimeError(f"You can only initialize one {type(self).__name__} object")
+            raise RuntimeError(
+                f"You can only initialize one {type(self).__name__} object"
+            )
         ASyncABC.__init__(self)
+
         self.wallets: Final[List[TreasuryWallet]] = []
         """The collection of wallets owned or controlled by the on-chain org"""
+
         for wallet in wallets:
             if isinstance(wallet, str):
                 self.wallets.append(TreasuryWallet(wallet))  # type: ignore [type-arg]
@@ -54,21 +113,27 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
                 raise TypeError(
                     f"`wallets` can only contain: {wallet_types}  You passed {wallet}"
                 )
-        
+
         self.sort_rules: Final = Rules(sort_rules) if sort_rules else None
 
         self.portfolio: Final = ExportablePortfolio(
             addresses=(
-                wallet.address if isinstance(wallet, TreasuryWallet) else wallet
+                wallet.address
                 for wallet in self.wallets
+                if wallet.networks is None or CHAINID in wallet.networks
             ),
             start_block=start_block,
             label=label,
             load_prices=True,
+            custom_buckets=custom_buckets,
             asynchronous=asynchronous,
         )
         """An eth_portfolio.Portfolio object used for exporting tx and balance history"""
 
+        self._llamapay: Final = (
+            llamapay.LlamaPayProcessor() if CHAINID in llamapay.networks else None
+        )
+
         self.asynchronous: Final = asynchronous
         """A boolean flag indicating whether the API for this `Treasury` object is sync or async by default"""
 
@@ -81,20 +146,48 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
     def txs(self) -> a_sync.ASyncIterator[LedgerEntry]:
         return self.portfolio.ledger.all_entries
 
-    async def populate_db(self, start_block: BlockNumber, end_block: BlockNumber) -> None:
-        """returns: number of new txs"""
-        # TODO: implement this
-        # NOTE: ensure stream loader task has been started
-        #global _streams_task
-        #if _streams_task is None:
-        #    _streams_task = create_task(streams._get_coro())
-        futs = []
-        async for entry in self.portfolio.ledger[start_block:end_block]:
-            if not entry.value:
-                # TODO: add an arg in eth-port to skip 0 value
-                logger.debug("zero value transfer, skipping %s", entry)
-                continue
-            futs.append(create_task(TreasuryTx.insert(entry)))
-        
-        if futs:
-            await tqdm_asyncio.gather(*futs, desc="Insert Txs to Postgres")
+    async def _insert_txs(
+        self, start_block: BlockNumber, end_block: BlockNumber
+    ) -> None:
+        """Populate the database with treasury transactions in a block range.
+
+        Streams ledger entries from `start_block` up to (but not including)
+        `end_block`, skips zero-value transfers, and inserts each remaining entry
+        into the DB via :meth:`dao_treasury.db.TreasuryTx.insert`. Uses
+        :class:`tqdm.asyncio.tqdm_asyncio` to display progress.
+
+        Args:
+            start_block: First block number to include (inclusive).
+            end_block: Last block number to include (exclusive).
+
+        Examples:
+            >>> # Insert transactions from block 0 to 10000
+            >>> await treasury._insert_txs(0, 10000)
+        """
+        with db_session:
+            futs = []
+            async for entry in self.portfolio.ledger[start_block:end_block]:
+                if not entry.value:
+                    # TODO: add an arg in eth-port to skip 0 value
+                    logger.debug("zero value transfer, skipping %s", entry)
+                    continue
+                futs.append(create_task(TreasuryTx.insert(entry)))
+            if futs:
+                await tqdm_asyncio.gather(*futs, desc="Insert Txs to Postgres")
+                logger.info(f"{len(futs)} transfers exported")
+
+    async def _process_streams(self) -> None:
+        if self._llamapay is not None:
+            await self._llamapay.process_streams(run_forever=True)
+
+    async def populate_db(
+        self, start_block: BlockNumber, end_block: BlockNumber
+    ) -> None:
+        """
+        Populate the database with treasury transactions and streams in parallel.
+        """
+        tasks = [self._insert_txs(start_block, end_block)]
+        if self._llamapay:
+            tasks.append(self._process_streams())
+        await gather(*tasks)
+        logger.info("db connection closed")