""" Universe Service — Historical stock universe construction for backtesting Builds monthly market_cap snapshots by combining: - SEC EDGAR companyfacts → shares_outstanding (quarterly, carried forward) - yfinance monthly close prices Data flow: 1. discover_tickers() → universe_ticker_registry (current US stocks via yf.screen) 2. build_snapshots() → universe_snapshot (monthly market_cap per ticker) 3. screen_historical() → filtered results for a given historical date Survivorship bias note: v1 universe is based on currently-listed stocks only. Delisted companies are not included. """ import asyncio import logging import math from datetime import date, datetime, timezone from typing import Dict, List, Optional, Tuple from sqlalchemy import and_, desc, func, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession from app.models.universe_snapshot import UniverseSnapshot, UniverseTickerRegistry from app.services.sec_http_client import SECHttpClient logger = logging.getLogger(__name__) _CHUNK = 3000 # asyncpg 32767-param limit (10 cols × 3000 = 30000) _PRICE_BATCH = 30 # tickers per yfinance bulk download _SEC_CONCURRENCY = 3 # concurrent SEC EDGAR companyfacts requests _SCREEN_PAGE = 250 # max per yf.screen() call class UniverseService: def __init__(self): self._http = SECHttpClient("Stock Oracle Universe Service") # ------------------------------------------------------------------ # Step 1: Discover and register tickers # ------------------------------------------------------------------ def _screen_page_sync(self, query, offset: int) -> dict: """Synchronous yfinance screen() — runs in executor.""" import yfinance_plus as yf return yf.screen( query, offset=offset, size=_SCREEN_PAGE, sortField="intradaymarketcap", sortAsc=False, ) async def _discover_via_yfinance(self, market_cap_min: float) -> List[dict]: """Try yfinance screener. Returns list of quote dicts or raises on 401/error.""" from yfinance import EquityQuery query = EquityQuery("and", [ EquityQuery("eq", ["region", "us"]), EquityQuery("gt", ["intradaymarketcap", market_cap_min]), ]) loop = asyncio.get_event_loop() all_quotes: List[dict] = [] offset = 0 while True: raw = await loop.run_in_executor( None, self._screen_page_sync, query, offset ) quotes = raw.get("quotes", []) if not quotes: break all_quotes.extend(quotes) total = raw.get("count") or raw.get("total") or len(all_quotes) offset += len(quotes) if offset >= total: break await asyncio.sleep(0.3) return all_quotes async def _discover_via_sec_edgar(self) -> List[dict]: """ Fallback: use SEC EDGAR company_tickers_exchange.json. Returns list of dicts with ticker/name/cik/exchange. No sector/industry (those remain NULL in registry). No market_cap filter (filter happens at screening time). """ # Case-insensitive match; normalize to canonical uppercase form _US_EXCHANGE_MAP = { "nyse": "NYSE", "nasdaq": "NASDAQ", "amex": "AMEX", "nysemkt": "AMEX", "arca": "ARCA", "nysearca": "ARCA", "bats": "BATS", "otc": "OTC", } try: data = await self._http.fetch_json( "https://www.sec.gov/files/company_tickers_exchange.json" ) except Exception as e: raise RuntimeError(f"SEC EDGAR company_tickers_exchange.json fetch failed: {e}") fields = data.get("fields", []) rows_raw = data.get("data", []) try: cik_idx = fields.index("cik") name_idx = fields.index("name") ticker_idx = fields.index("ticker") exchange_idx = fields.index("exchange") except ValueError as e: raise RuntimeError(f"Unexpected company_tickers_exchange.json schema: {e}") quotes = [] for row in rows_raw: try: ticker = str(row[ticker_idx]).upper().strip() exchange_raw = str(row[exchange_idx]).strip() canonical_exchange = _US_EXCHANGE_MAP.get(exchange_raw.lower()) if not ticker or len(ticker) > 10: continue # Keep only major US exchanges (case-insensitive) if not canonical_exchange: continue # Skip preferred stocks / warrants / rights / units (contain - or end in W/R/U/Z) if "-" in ticker: continue if len(ticker) > 1 and ticker[-1] in ("W", "R", "Z") and ticker[:-1].isalpha(): continue quotes.append({ "symbol": ticker, "shortName": str(row[name_idx]) if row[name_idx] else None, "cik_override": str(row[cik_idx]).zfill(10), "exchange": canonical_exchange, "sector": None, "industry": None, "quoteType": "EQUITY", }) except (IndexError, TypeError): continue logger.info(f"Universe: SEC EDGAR fallback returned {len(quotes)} tickers") return quotes async def discover_tickers(self, db: AsyncSession, market_cap_min: float = 1e8) -> Dict: """ Discover US-listed stocks and store in registry. Primary: yfinance screener (includes market_cap filter + sector/industry) Fallback: SEC EDGAR company_tickers_exchange.json (no market_cap filter, no sector/industry — those remain NULL) Returns: { tickers_found, tickers_registered, source } """ source = "yfinance" all_quotes: List[dict] = [] try: all_quotes = await self._discover_via_yfinance(market_cap_min) if not all_quotes: raise RuntimeError("yfinance screener returned 0 results") logger.info(f"Universe: yfinance screener returned {len(all_quotes)} quotes") except Exception as e: logger.warning( f"Universe: yfinance screener failed ({e}), " "falling back to SEC EDGAR company_tickers_exchange.json" ) source = "sec_edgar" all_quotes = await self._discover_via_sec_edgar() # Fetch CIK map from SEC once (disk-cached after first call) cik_map = await self._fetch_cik_map() _REVERSE_EXCHANGE = { "NYQ": "NYSE", "NMS": "NASDAQ", "NGM": "NASDAQ", "NCM": "NASDAQ", "ASE": "AMEX", "PCX": "NYSE_ARCA", } _SKIP_TYPES = {"ETF", "MUTUALFUND", "INDEX", "CURRENCY", "FUTURE", "OPTION"} rows = [] for q in all_quotes: ticker = (q.get("symbol") or "").upper() if not ticker or len(ticker) > 10: continue if (q.get("quoteType") or "").upper() in _SKIP_TYPES: continue exchange_code = q.get("exchange", "") rows.append({ "ticker": ticker, "name": q.get("shortName") or q.get("longName"), "cik": q.get("cik_override") or cik_map.get(ticker), "sector": q.get("sector"), "industry": q.get("industry"), "exchange": _REVERSE_EXCHANGE.get(exchange_code, exchange_code) or None, "is_active": True, }) if not rows: return {"tickers_found": len(all_quotes), "tickers_registered": 0, "source": source} for i in range(0, len(rows), _CHUNK): chunk = rows[i:i + _CHUNK] stmt = pg_insert(UniverseTickerRegistry).values(chunk) stmt = stmt.on_conflict_do_update( constraint="uq_universe_ticker_registry", set_={ "name": stmt.excluded.name, "cik": stmt.excluded.cik, "sector": stmt.excluded.sector, "industry": stmt.excluded.industry, "exchange": stmt.excluded.exchange, "is_active": stmt.excluded.is_active, "updated_at": func.now(), }, ) await db.execute(stmt) await db.commit() logger.info(f"Universe: upserted {len(rows)} tickers into registry (source={source})") return {"tickers_found": len(all_quotes), "tickers_registered": len(rows), "source": source} async def _fetch_cik_map(self) -> Dict[str, str]: """Fetch SEC company_tickers.json once → {TICKER: padded_cik}.""" try: data = await self._http.fetch_json( "https://www.sec.gov/files/company_tickers.json" ) return { entry["ticker"].upper(): str(entry["cik_str"]).zfill(10) for entry in data.values() if "ticker" in entry and "cik_str" in entry } except Exception as e: logger.warning(f"Universe: CIK map fetch failed: {e}") return {} # ------------------------------------------------------------------ # Step 2: Build monthly snapshots # ------------------------------------------------------------------ async def build_snapshots( self, db_or_factory, tickers: Optional[List[str]], start_date: str, end_date: str, force_rebuild: bool = False, ) -> Dict: """ Build monthly market_cap snapshots for given tickers (or all registry). db_or_factory: Either an AsyncSession (for small sync calls ≤20 tickers) or an AsyncSessionLocal factory (for large background jobs). Using a factory means each batch opens/closes its own DB connection, preventing the long-running job from monopolising the connection pool and blocking the API. Returns: { tickers_processed, tickers_failed, snapshots_created } """ from app.core.database import AsyncSessionLocal # Normalise: wrap a plain session in a no-op factory so the rest of the # code can always call `async with _db_ctx() as db:` from contextlib import asynccontextmanager if callable(db_or_factory) and not isinstance(db_or_factory, AsyncSession): # It's a factory (e.g. AsyncSessionLocal) _db_ctx = db_or_factory else: # It's an existing session — wrap it so we never close it here _session = db_or_factory @asynccontextmanager async def _db_ctx(): yield _session # ---- Resolve ticker list ---- async with _db_ctx() as db: if tickers: ticker_list = [t.upper() for t in tickers] else: result = await db.execute( select(UniverseTickerRegistry.ticker) .where(UniverseTickerRegistry.is_active == True) .order_by(UniverseTickerRegistry.ticker) ) ticker_list = [r[0] for r in result.fetchall()] if not ticker_list: logger.warning("Universe: no tickers to process") return {"tickers_processed": 0, "tickers_failed": 0, "snapshots_created": 0} logger.info( f"Universe: building snapshots for {len(ticker_list)} tickers " f"({start_date} → {end_date})" ) if force_rebuild: if not tickers: # Full rebuild: wipe ALL snapshots (also removes stale tickers # that were de-registered from the registry) await db.execute(UniverseSnapshot.__table__.delete()) else: # Partial rebuild: delete only the requested tickers for i in range(0, len(ticker_list), _CHUNK): chunk = ticker_list[i:i + _CHUNK] await db.execute( UniverseSnapshot.__table__.delete().where( UniverseSnapshot.ticker.in_(chunk) ) ) await db.commit() # Load registry metadata (sector/industry/exchange/CIK) registry_map: Dict[str, UniverseTickerRegistry] = {} for i in range(0, len(ticker_list), _CHUNK): chunk = ticker_list[i:i + _CHUNK] reg_result = await db.execute( select(UniverseTickerRegistry).where( UniverseTickerRegistry.ticker.in_(chunk) ) ) for r in reg_result.scalars().all(): registry_map[r.ticker] = r total_snapshots = 0 total_failed = 0 total_batches = (len(ticker_list) + _PRICE_BATCH - 1) // _PRICE_BATCH _CACHE_CLEAR_EVERY = 20 # clear SEC in-memory cache every N batches to prevent OOM # ---- Process in batches — each batch uses its own DB session ---- for batch_start in range(0, len(ticker_list), _PRICE_BATCH): batch = ticker_list[batch_start:batch_start + _PRICE_BATCH] batch_num = batch_start // _PRICE_BATCH + 1 # Fetch monthly prices (synchronous yfinance, runs in thread) price_data = await asyncio.to_thread( self._fetch_bulk_monthly_prices, batch, start_date, end_date ) # Fetch shares_outstanding from SEC EDGAR (async, rate-limited) sem = asyncio.Semaphore(_SEC_CONCURRENCY) async def _fetch_one(tkr: str) -> Tuple[str, List]: async with sem: history = await self._fetch_shares_history( tkr, registry_map.get(tkr) ) return tkr, history results = await asyncio.gather( *[_fetch_one(t) for t in batch], return_exceptions=True ) shares_map: Dict[str, List[Tuple[date, float]]] = {} for item in results: if isinstance(item, Exception): continue tkr, history = item shares_map[tkr] = history # Build snapshot rows batch_rows = [] for ticker in batch: shares_history = shares_map.get(ticker, []) ticker_prices = price_data.get(ticker, {}) reg = registry_map.get(ticker) if not ticker_prices: total_failed += 1 continue # Use the LATEST available shares_outstanding (most recent SEC filing). # yfinance returns split-adjusted prices retroactively, so using the # post-split shares count gives correct market_cap across all periods. latest_shares = shares_history[-1][1] if shares_history else None for snap_date_key, close in ticker_prices.items(): if isinstance(snap_date_key, datetime): snap_date = snap_date_key.date() else: snap_date = snap_date_key shares = latest_shares if shares is None or close is None: continue # Sanity checks: skip absurd values (data quality) if close > 1_000_000 or close <= 0: # max BRK-A ~$600K continue if shares < 100_000: # too few shares for a public co continue market_cap = shares * close if math.isnan(market_cap) or market_cap <= 0 or market_cap > 5e12: continue snapshot_dt = datetime( snap_date.year, snap_date.month, 1, tzinfo=timezone.utc ) batch_rows.append({ "ticker": ticker, "snapshot_date": snapshot_dt, "close_price": round(close, 4), "shares_outstanding": shares, "market_cap": round(market_cap, 0), "sector": reg.sector if reg else None, "industry": reg.industry if reg else None, "exchange": reg.exchange if reg else None, }) # Upsert with a fresh short-lived DB session (releases connection immediately after) if batch_rows: async with _db_ctx() as db: for i in range(0, len(batch_rows), _CHUNK): chunk = batch_rows[i:i + _CHUNK] stmt = pg_insert(UniverseSnapshot).values(chunk) stmt = stmt.on_conflict_do_update( constraint="uq_universe_snapshot", set_={ "close_price": stmt.excluded.close_price, "shares_outstanding": stmt.excluded.shares_outstanding, "market_cap": stmt.excluded.market_cap, "sector": stmt.excluded.sector, "industry": stmt.excluded.industry, "exchange": stmt.excluded.exchange, }, ) result = await db.execute(stmt) total_snapshots += result.rowcount await db.commit() logger.info( f"Universe: batch {batch_num}/{total_batches} — " f"{len(batch)} tickers, {len(batch_rows)} snapshot rows" ) # Periodically clear SEC in-memory JSON cache to prevent OOM. # Disk cache is retained — subsequent lookups re-load from disk. if batch_num % _CACHE_CLEAR_EVERY == 0: import gc self._http._json_cache.clear() self._http._text_cache.clear() gc.collect() logger.info(f"Universe: cleared SEC JSON cache at batch {batch_num}") # Yield to event loop between batches to keep API responsive await asyncio.sleep(0) logger.info( f"Universe: build complete — {len(ticker_list)} tickers, " f"{total_snapshots} snapshots created, {total_failed} failed" ) return { "tickers_processed": len(ticker_list), "tickers_failed": total_failed, "snapshots_created": total_snapshots, } async def _fetch_shares_history( self, ticker: str, registry_entry: Optional[UniverseTickerRegistry], ) -> List[Tuple[date, float]]: """Fetch shares_outstanding history from SEC EDGAR companyfacts.""" cik = None if registry_entry and registry_entry.cik: cik = registry_entry.cik else: cik = await self._http.get_company_cik(ticker) if not cik: return [] try: url = ( f"{self._http.sec_base_data}" f"/api/xbrl/companyfacts/CIK{str(cik).zfill(10)}.json" ) facts = await self._http.fetch_json(url) return _extract_shares_history(facts) except Exception as e: logger.debug(f"Universe: companyfacts failed for {ticker}: {e}") return [] def _fetch_bulk_monthly_prices( self, tickers: List[str], start: str, end: str, ) -> Dict[str, Dict[date, float]]: """ Synchronous bulk monthly price fetch via yfinance. Returns {ticker: {date: close_price}}. """ try: import math as _math import yfinance_plus as yf except ImportError: logger.error("yfinance not available for price download") return {} if not tickers: return {} # Always use group_by='ticker' for consistent MultiIndex column structure: # data[ticker] -> DataFrame with ['Close', 'Adj Close', ...] try: data = yf.download( tickers=tickers, start=start, end=end, interval="1mo", auto_adjust=False, progress=False, threads=False, # disable internal thread pool — prevents OS thread explosion group_by="ticker", ) except Exception as e: logger.warning(f"Universe: bulk price download failed: {e}") return {} if data is None or data.empty: return {} result: Dict[str, Dict[date, float]] = {} def _safe_close(val) -> Optional[float]: if val is None: return None try: f = float(val) return None if _math.isnan(f) or f <= 0 else f except (TypeError, ValueError): return None def _to_date(dt) -> Optional[date]: if hasattr(dt, "date"): return dt.date() if isinstance(dt, date): return dt return None for ticker in tickers: result[ticker] = {} try: # With group_by='ticker', columns are MultiIndex (ticker, price_type) # data[ticker] gives a flat DataFrame with price columns lvl0 = data.columns.get_level_values(0) if ticker not in lvl0: continue ticker_df = data[ticker] close_col = ticker_df.get("Close") if close_col is None: close_col = ticker_df.get("Adj Close") if close_col is None: continue for dt, val in close_col.items(): d = _to_date(dt) c = _safe_close(val) if d and c: result[ticker][d] = c except Exception as e: logger.debug(f"Universe: price parse error for {ticker}: {e}") return result # ------------------------------------------------------------------ # Step 3: Historical screening # ------------------------------------------------------------------ async def screen_historical( self, db: AsyncSession, date_str: str, market_cap_min: Optional[float], market_cap_max: Optional[float], sector: Optional[str], exchange: Optional[str], page: int, page_size: int, sort_by: str, sort_ascending: bool, ) -> Tuple[List[Dict], int, str]: """ Screen stocks at a historical date based on market_cap and other criteria. Returns: (items, total_count, actual_snapshot_date_str) """ try: target = date.fromisoformat(date_str) except ValueError: raise ValueError(f"Invalid date format: {date_str!r} — use YYYY-MM-DD") snapshot_dt = datetime(target.year, target.month, 1, tzinfo=timezone.utc) # Always exclude clearly bad data (sanity cap: $5T max, historical record is ~$3.7T) _MAX_MARKET_CAP = 5e12 conditions = [ UniverseSnapshot.snapshot_date == snapshot_dt, UniverseSnapshot.market_cap <= _MAX_MARKET_CAP, ] if market_cap_min is not None: conditions.append(UniverseSnapshot.market_cap >= market_cap_min) if market_cap_max is not None: conditions.append(UniverseSnapshot.market_cap <= market_cap_max) if sector: conditions.append(UniverseSnapshot.sector == sector) if exchange: conditions.append(UniverseSnapshot.exchange == exchange.upper()) filter_clause = and_(*conditions) # Total count count_q = await db.execute( select(func.count(UniverseSnapshot.id)).where(filter_clause) ) total = count_q.scalar() or 0 # Sort sort_col = { "market_cap": UniverseSnapshot.market_cap, "ticker": UniverseSnapshot.ticker, }.get(sort_by, UniverseSnapshot.market_cap) order = sort_col.asc() if sort_ascending else desc(sort_col) # Paginated fetch rows_result = await db.execute( select(UniverseSnapshot) .where(filter_clause) .order_by(order) .limit(page_size) .offset((page - 1) * page_size) ) rows = rows_result.scalars().all() # Enrich with names from registry tickers_in_page = [r.ticker for r in rows] name_map: Dict[str, Optional[str]] = {} if tickers_in_page: name_result = await db.execute( select(UniverseTickerRegistry.ticker, UniverseTickerRegistry.name) .where(UniverseTickerRegistry.ticker.in_(tickers_in_page)) ) name_map = {r[0]: r[1] for r in name_result.fetchall()} items = [ { "ticker": r.ticker, "name": name_map.get(r.ticker), "market_cap": r.market_cap, "close_price": r.close_price, "shares_outstanding": r.shares_outstanding, "sector": r.sector, "industry": r.industry, "exchange": r.exchange, "snapshot_date": r.snapshot_date, } for r in rows ] return items, total, snapshot_dt.strftime("%Y-%m-%d") async def get_registry( self, db: AsyncSession, sector: Optional[str], exchange: Optional[str], is_active: Optional[bool], page: int, page_size: int, ) -> Tuple[List, int]: """Browse the ticker registry with optional filters.""" conditions = [] if sector: conditions.append(UniverseTickerRegistry.sector == sector) if exchange: conditions.append(UniverseTickerRegistry.exchange == exchange.upper()) if is_active is not None: conditions.append(UniverseTickerRegistry.is_active == is_active) where = and_(*conditions) if conditions else True count_q = await db.execute( select(func.count(UniverseTickerRegistry.id)).where(where) ) total = count_q.scalar() or 0 result = await db.execute( select(UniverseTickerRegistry) .where(where) .order_by(UniverseTickerRegistry.ticker) .limit(page_size) .offset((page - 1) * page_size) ) return result.scalars().all(), total # ------------------------------------------------------------------ # Module-level helpers # ------------------------------------------------------------------ def _extract_shares_history( facts_data: dict, ) -> List[Tuple[date, float]]: """ Extract sorted (period_end_date, shares_outstanding) from SEC companyfacts JSON. Prefers CommonStockSharesOutstanding over weighted average concepts. """ if not facts_data or "facts" not in facts_data: return [] us_gaap = facts_data["facts"].get("us-gaap", {}) _CONCEPTS = [ "CommonStockSharesOutstanding", "WeightedAverageNumberOfSharesOutstandingBasic", "WeightedAverageNumberOfDilutedSharesOutstanding", ] _FORMS = {"10-K", "10-Q", "10-K/A", "10-Q/A"} all_points: Dict[date, float] = {} for concept in _CONCEPTS: if concept not in us_gaap: continue for unit_key, entries in us_gaap[concept].get("units", {}).items(): if unit_key != "shares": continue for entry in entries: if entry.get("form") not in _FORMS: continue end_str = entry.get("end") val = entry.get("val") if not end_str or val is None: continue try: end_d = date.fromisoformat(end_str) all_points[end_d] = float(val) except (ValueError, TypeError): continue return sorted(all_points.items()) def _get_shares_at_date( shares_history: List[Tuple[date, float]], target: date ) -> Optional[float]: """Carry-forward: most recent shares_outstanding on or before target date.""" result = None for entry_date, shares in shares_history: if entry_date <= target: result = shares else: break return result