Source code for trader.persistence.index_membership

"""Repository for `index_membership` (issue #27).

A set, not a history: `trader scan-universe` upserts on every scrape, and
`discover_symbols`'s analyst-scan candidate source reads the current members
of one named index.
"""

from collections.abc import Iterable
from datetime import datetime

from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker

from trader.persistence.models import IndexMembership

__all__ = ["IndexMembershipRepository"]




class IndexMembershipRepository:
    """Writes and reads `index_membership`."""

    def __init__(self, session_factory: sessionmaker[Session]) -> None:
        self._session_factory = session_factory



    def upsert(self, index_name: str, symbols: Iterable[str], *, now: datetime) -> None:
        """Record `symbols` as the current membership of `index_name`.

        Existing members not in `symbols` are left as they were — this only
        ever adds. A departed constituent lingering in `index_membership`
        after the next scrape costs nothing: it still has to clear
        `is_strong_buy_consensus` and the full `_evaluate` chain like any
        other candidate, and `asset_gate` rejects it if it is no longer
        tradable.
        """
        normalized = {s.strip().upper() for s in symbols if s.strip()}
        if not normalized:
            return
        with self._session_factory() as session:
            existing = set(
                session.scalars(
                    select(IndexMembership.symbol).where(
                        IndexMembership.index_name == index_name,
                        IndexMembership.symbol.in_(normalized),
                    )
                ).all()
            )
            new_rows = [
                IndexMembership(index_name=index_name, symbol=symbol, discovered_at=now)
                for symbol in normalized - existing
            ]
            session.add_all(new_rows)
            session.commit()




    def members(self, index_name: str) -> list[str]:
        """Every symbol currently recorded under `index_name`."""
        statement = select(IndexMembership.symbol).where(
            IndexMembership.index_name == index_name
        )
        with self._session_factory() as session:
            return list(session.scalars(statement).all())