quant-trading/universe.py

"""
Stock universe providers — fetch index constituents for backtesting.

Each function returns a dict with 'tickers' (list[str]) and 'benchmark' (str)
compatible with Yahoo Finance ticker format.  Results are cached locally so
Wikipedia / web scrapes only happen once per day.
"""

import io
import json
import os
import urllib.request
from datetime import date

import pandas as pd

_HEADERS = {"User-Agent": "Mozilla/5.0 (quant-backtest)"}


def _fetch_html_tables(url: str) -> list[pd.DataFrame]:
    """Fetch HTML tables from a URL with a proper User-Agent."""
    req = urllib.request.Request(url, headers=_HEADERS)
    with urllib.request.urlopen(req) as resp:
        html = resp.read().decode("utf-8")
    return pd.read_html(io.StringIO(html))

CACHE_DIR = "data"


def _read_cache(name: str) -> list[str] | None:
    path = os.path.join(CACHE_DIR, f"universe_{name}.json")
    if not os.path.exists(path):
        return None
    try:
        with open(path) as f:
            data = json.load(f)
        if data.get("date") == str(date.today()):
            return data["tickers"]
    except Exception:
        pass
    return None


def _write_cache(name: str, tickers: list[str]) -> None:
    os.makedirs(CACHE_DIR, exist_ok=True)
    path = os.path.join(CACHE_DIR, f"universe_{name}.json")
    with open(path, "w") as f:
        json.dump({"date": str(date.today()), "tickers": tickers}, f)


# ---------------------------------------------------------------------------
# S&P 500
# ---------------------------------------------------------------------------

def get_sp500() -> list[str]:
    """Fetch current S&P 500 constituents from Wikipedia."""
    cached = _read_cache("sp500")
    if cached:
        print(f"--- Loaded S&P 500 universe from cache ({len(cached)} tickers) ---")
        return cached

    print("--- Fetching S&P 500 constituents from Wikipedia ---")
    url = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"
    tables = _fetch_html_tables(url)
    df = tables[0]
    tickers = sorted(df["Symbol"].str.replace(".", "-", regex=False).tolist())
    _write_cache("sp500", tickers)
    print(f"--- Got {len(tickers)} S&P 500 tickers ---")
    return tickers


# ---------------------------------------------------------------------------
# CSI 300
# ---------------------------------------------------------------------------

# Yahoo Finance format: Shanghai → XXXXXX.SS, Shenzhen → XXXXXX.SZ
# Wikipedia CSI 300 page has a table with stock codes and exchanges.

CSI300_CONS_URL = ("https://oss-ch.csindex.com.cn/static/html/csindex/public/uploads/"
                   "file/autofile/cons/000300cons.xls")


def _csi300_official() -> list[str]:
    """Current constituents from China Securities Index's published file."""
    req = urllib.request.Request(CSI300_CONS_URL, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        payload = resp.read()
    df = pd.read_excel(io.BytesIO(payload), dtype=str)
    code_col = next(c for c in df.columns if "Constituent Code" in c)
    exch_col = next(c for c in df.columns if c.startswith("交易所Exchange"))
    tickers = [f"{code}.{'SS' if '上海' in exch else 'SZ'}"
               for code, exch in zip(df[code_col], df[exch_col])]
    if len(tickers) < 250:
        raise RuntimeError(f"CSI 300 file has only {len(tickers)} rows")
    return sorted(set(tickers))


def _csi300_latest_snapshot() -> list[str]:
    with open(os.path.join(CACHE_DIR, "csi300_history.json")) as f:
        snaps = json.load(f)["snapshots"]
    return snaps[max(snaps)]


def _csi300_wikipedia() -> list[str]:
    tables = _fetch_html_tables("https://en.wikipedia.org/wiki/CSI_300_Index")
    # The constituent table has ~300 rows and a "Ticker" column like "SSE: 600519"
    df = next((t for t in tables if len(t) >= 200 and "Ticker" in t.columns), None)
    if df is None:
        raise RuntimeError("Could not find CSI 300 constituent table on Wikipedia")
    tickers = []
    for raw in df["Ticker"].astype(str):
        parts = raw.split(":")
        if len(parts) == 2 and parts[1].strip().isdigit():
            suffix = ".SS" if parts[0].strip().upper() == "SSE" else ".SZ"
            tickers.append(parts[1].strip() + suffix)
    return sorted(set(tickers))


def get_csi300() -> list[str]:
    """Current CSI 300 constituents.

    Source order: China Securities Index's published constituent file, then
    the latest snapshot in data/csi300_history.json, then Wikipedia (whose
    table lagged the real index by ~60 names in 2026).
    """
    cached = _read_cache("csi300")
    if cached:
        print(f"--- Loaded CSI 300 universe from cache ({len(cached)} tickers) ---")
        return cached

    for label, source in (("csindex.com.cn", _csi300_official),
                          ("stored CSI 300 history", _csi300_latest_snapshot),
                          ("Wikipedia", _csi300_wikipedia)):
        try:
            tickers = source()
        except Exception as exc:  # noqa: BLE001 — fall through to the next source
            print(f"--- CSI 300 from {label} failed: {exc} ---")
            continue
        _write_cache("csi300", tickers)
        print(f"--- Got {len(tickers)} CSI 300 tickers from {label} ---")
        return tickers
    raise RuntimeError("No CSI 300 constituent source available")


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

UNIVERSES = {
    "us": {
        "fetch": get_sp500,
        "benchmark": "SPY",
        "benchmark_label": "SPY (Benchmark)",
    },
    "cn": {
        "fetch": get_csi300,
        # Yahoo's 000300.SS index series is truncated (and at times a single
        # row); the CSI 300 ETF has a continuous history back to 2012.
        "benchmark": "510300.SS",
        "benchmark_label": "CSI 300 ETF 510300 (Benchmark)",
    },
}