You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

621 lines
22 KiB
Python

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

"""
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 = 100 # 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 as yf
return yf.screen(
query, offset=offset, size=_SCREEN_PAGE,
sortField="intradaymarketcap", sortAsc=False,
)
async def discover_tickers(self, db: AsyncSession, market_cap_min: float = 1e8) -> Dict:
"""
Discover US-listed stocks via yfinance screener and store in registry.
Filters out ETFs, mutual funds, and foreign-listed stocks.
Returns: { tickers_found, tickers_registered }
"""
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)
logger.info(f"Universe: screener returned {len(all_quotes)} quotes")
# 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": 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}
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")
return {"tickers_found": len(all_quotes), "tickers_registered": len(rows)}
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: AsyncSession,
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).
For each ticker:
1. Fetch shares_outstanding history from SEC EDGAR companyfacts
2. Fetch monthly close prices via yfinance bulk download
3. Compute market_cap = close × shares per month
4. Upsert into universe_snapshot
Returns: { tickers_processed, tickers_failed, snapshots_created }
"""
# Resolve ticker list
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:
await db.execute(
UniverseSnapshot.__table__.delete().where(
UniverseSnapshot.ticker.in_(ticker_list)
)
)
await db.flush()
# Load registry metadata (sector/industry/exchange/CIK)
reg_result = await db.execute(
select(UniverseTickerRegistry).where(
UniverseTickerRegistry.ticker.in_(ticker_list)
)
)
registry_map: Dict[str, UniverseTickerRegistry] = {
r.ticker: r for r in reg_result.scalars().all()
}
total_snapshots = 0
total_failed = 0
# Process in batches of _PRICE_BATCH for yfinance bulk download
for batch_start in range(0, len(ticker_list), _PRICE_BATCH):
batch = ticker_list[batch_start:batch_start + _PRICE_BATCH]
# ---- Fetch monthly prices (synchronous, in thread) ----
price_data = await asyncio.to_thread(
self._fetch_bulk_monthly_prices, batch, start_date, end_date
)
# ---- Fetch shares_outstanding from SEC EDGAR (concurrent) ----
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.
# Small error (~10-20%) may occur for companies with large buyback programs,
# but this is acceptable for screening purposes.
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
market_cap = shares * close
if math.isnan(market_cap) or market_cap <= 0:
continue
# Normalize to first of month
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,
})
# ---- Batch upsert ----
if batch_rows:
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()
batch_num = batch_start // _PRICE_BATCH + 1
total_batches = (len(ticker_list) + _PRICE_BATCH - 1) // _PRICE_BATCH
logger.info(
f"Universe: batch {batch_num}/{total_batches}"
f"{len(batch)} tickers, {len(batch_rows)} snapshot rows"
)
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 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=True,
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)
conditions = [UniverseSnapshot.snapshot_date == snapshot_dt]
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