|
|
"""
|
|
|
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 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 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
|