dao-treasury 0.0.17__cp312-cp312-win32.whl → 0.0.61__cp312-cp312-win32.whl

dao_treasury/sorting/rule.py

@@ -31,6 +31,7 @@ See Also:
 
 from collections import defaultdict
 from dataclasses import dataclass
+from logging import getLogger
 from typing import (
     TYPE_CHECKING,
     DefaultDict,
@@ -53,6 +54,9 @@ if TYPE_CHECKING:
     from dao_treasury.db import TreasuryTx
 
 
+logger: Final = getLogger(__name__)
+_log_debug: Final = logger.debug
+
 SORT_RULES: DefaultDict[Type[SortRule], List[SortRule]] = defaultdict(list)
 """Mapping from sort rule classes to lists of instantiated rules, in creation order per class.
 
@@ -126,8 +130,6 @@ class _SortRule:
     func: Optional[SortFunction] = None
     """Custom matching function that takes a `TreasuryTx` and returns a bool or an awaitable that returns a bool."""
 
-    # __instances__: ClassVar[List[Self]] = []
-
     def __post_init__(self) -> None:
         """Validate inputs, checksum addresses, and register the rule.
 
@@ -214,6 +216,7 @@ class _SortRule:
                 getattr(tx, matcher) == getattr(self, matcher) for matcher in matchers
             )
 
+        _log_debug("checking %s for %s", tx, self.func)
         match = self.func(tx)  # type: ignore [misc]
         return match if isinstance(match, bool) else await match
 
@@ -229,8 +232,8 @@ class _InboundSortRule(_SortRule):
     async def match(self, tx: "TreasuryTx") -> bool:
         return (
             tx.to_address is not None
-            and TreasuryWallet._get_instance(tx.to_address.address) is not None
-            and await super().match(tx)
+            and TreasuryWallet.check_membership(tx.to_address.address, tx.block)
+            and await super(_InboundSortRule, self).match(tx)
         )
 
 
@@ -243,9 +246,9 @@ class _OutboundSortRule(_SortRule):
     """
 
     async def match(self, tx: "TreasuryTx") -> bool:
-        return TreasuryWallet._get_instance(
-            tx.from_address.address
-        ) is not None and await super().match(tx)
+        return TreasuryWallet.check_membership(
+            tx.from_address.address, tx.block
+        ) and await super(_OutboundSortRule, self).match(tx)
 
 
 @mypyc_attr(native_class=False)
@@ -262,7 +265,7 @@ class RevenueSortRule(_InboundSortRule):
     def __post_init__(self) -> None:
         """Prepends `self.txgroup` with 'Revenue:'."""
         object.__setattr__(self, "txgroup", f"Revenue:{self.txgroup}")
-        super().__post_init__()
+        super(RevenueSortRule, self).__post_init__()
 
 
 @mypyc_attr(native_class=False)
@@ -275,7 +278,7 @@ class CostOfRevenueSortRule(_OutboundSortRule):
     def __post_init__(self) -> None:
         """Prepends `self.txgroup` with 'Cost of Revenue:'."""
         object.__setattr__(self, "txgroup", f"Cost of Revenue:{self.txgroup}")
-        super().__post_init__()
+        super(CostOfRevenueSortRule, self).__post_init__()
 
 
 @mypyc_attr(native_class=False)
@@ -288,7 +291,7 @@ class ExpenseSortRule(_OutboundSortRule):
     def __post_init__(self) -> None:
         """Prepends `self.txgroup` with 'Expenses:'."""
         object.__setattr__(self, "txgroup", f"Expenses:{self.txgroup}")
-        super().__post_init__()
+        super(ExpenseSortRule, self).__post_init__()
 
 
 @mypyc_attr(native_class=False)
@@ -301,7 +304,7 @@ class OtherIncomeSortRule(_InboundSortRule):
     def __post_init__(self) -> None:
         """Prepends `self.txgroup` with 'Other Income:'."""
         object.__setattr__(self, "txgroup", f"Other Income:{self.txgroup}")
-        super().__post_init__()
+        super(OtherIncomeSortRule, self).__post_init__()
 
 
 @mypyc_attr(native_class=False)
@@ -314,7 +317,7 @@ class OtherExpenseSortRule(_OutboundSortRule):
     def __post_init__(self) -> None:
         """Prepends `self.txgroup` with 'Other Expenses:'."""
         object.__setattr__(self, "txgroup", f"Other Expenses:{self.txgroup}")
-        super().__post_init__()
+        super(OtherExpenseSortRule, self).__post_init__()
 
 
 @mypyc_attr(native_class=False)
@@ -327,7 +330,7 @@ class IgnoreSortRule(_SortRule):
     def __post_init__(self) -> None:
         """Prepends `self.txgroup` with 'Ignore:'."""
         object.__setattr__(self, "txgroup", f"Ignore:{self.txgroup}")
-        super().__post_init__()
+        super(IgnoreSortRule, self).__post_init__()
 
 
 TRule = TypeVar(

dao_treasury/sorting/rules/__init__.py

@@ -0,0 +1 @@
+from dao_treasury.sorting.rules.ignore import *

dao_treasury/sorting/rules/ignore/__init__.py

@@ -0,0 +1 @@
+from dao_treasury.sorting.rules.ignore.llamapay import *

dao_treasury/sorting/rules/ignore/llamapay.py

@@ -0,0 +1,20 @@
+from dao_treasury import TreasuryTx
+from dao_treasury.sorting.factory import ignore
+from dao_treasury.streams import llamapay
+
+
+@ignore("LlamaPay")
+def is_llamapay_stream_replenishment(tx: TreasuryTx) -> bool:
+    if tx.to_address.address in llamapay.factories:  # type: ignore [operator]
+        # We amortize these streams daily in the `llamapay` module, you'll sort each stream appropriately.
+        return True
+
+    # NOTE: not sure if we want this yet
+    # Puling unused funds back from vesting escrow / llamapay
+    # elif tx.from_address == "Contract: LlamaPay" and "StreamCancelled" in tx.events:
+    #    if tx.amount > 0:
+    #        tx.amount *= -1
+    #    if tx.value_usd > 0:
+    #        tx.value_usd *= -1
+    #    return True
+    return False

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,4 +1,21 @@
-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
@@ -9,11 +26,14 @@ from eth_typing import BlockNumber
 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]
@@ -56,20 +76,22 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
             TypeError: If any item in `wallets` is not a str or TreasuryWallet.
 
         Examples:
-            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
-            ... )
+            .. 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:
@@ -94,7 +116,11 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
         self.sort_rules: Final = Rules(sort_rules) if sort_rules else None
 
         self.portfolio: Final = ExportablePortfolio(
-            addresses=(wallet.address for wallet in self.wallets),
+            addresses=(
+                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,
@@ -102,6 +128,10 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
         )
         """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"""
 
@@ -114,7 +144,7 @@ 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(
+    async def _insert_txs(
         self, start_block: BlockNumber, end_block: BlockNumber
     ) -> None:
         """Populate the database with treasury transactions in a block range.
@@ -130,15 +160,32 @@ class Treasury(a_sync.ASyncGenericBase):  # type: ignore [misc]
 
         Examples:
             >>> # Insert transactions from block 0 to 10000
-            >>> await treasury.populate_db(0, 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.
         """
-        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")
+        tasks = [self._insert_txs(start_block, end_block)]
+        if self._llamapay:
+            tasks.append(self._process_streams())
+        await gather(*tasks)
+        logger.info("db connection closed")