Yuxuan Yan
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""CSI 300 (HS300) and CSI 500 (ZZ500) universe helpers."""
|
|
import logging
|
|
|
|
import baostock as bs
|
|
import pandas as pd
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# First 30 HS300 constituents (large caps) in 'shXXXXXX' / 'szXXXXXX' format.
|
|
# Hardcoded for fast, deterministic smoke tests. Use get_hs300_stocks() for the
|
|
# live, full list — downloading daily bars for all ~300 takes roughly 10 minutes.
|
|
SYMBOLS = [
|
|
"sh600000", "sh600009", "sh600010", "sh600028", "sh600030",
|
|
"sh600036", "sh600048", "sh600050", "sh600104", "sh600276",
|
|
"sh600309", "sh600519", "sh600585", "sh600887", "sh600900",
|
|
"sh601012", "sh601166", "sh601288", "sh601318", "sh601398",
|
|
"sh601628", "sh601668", "sh601857", "sh601888", "sh601988",
|
|
"sz000001", "sz000002", "sz000333", "sz000651", "sz000858",
|
|
]
|
|
|
|
|
|
# First 30 CSI 500 (ZZ500) constituents (mid/small caps) in 'shXXXXXX' /
|
|
# 'szXXXXXX' format. Hardcoded for fast, deterministic smoke tests. Use
|
|
# get_zz500_stocks() for the live, full list. Mean reversion tends to be
|
|
# stronger in these smaller caps than in the HS300 large caps.
|
|
CSI500_SYMBOLS = [
|
|
"sh600006", "sh600008", "sh600017", "sh600020", "sh600021",
|
|
"sh600026", "sh600037", "sh600039", "sh600053", "sh600056",
|
|
"sh600060", "sh600061", "sh600062", "sh600073", "sh600089",
|
|
"sh600095", "sh600118", "sh600125", "sh600126", "sh600143",
|
|
"sh600153", "sh600160", "sh600169", "sh600176", "sh600183",
|
|
"sz000009", "sz000012", "sz000021", "sz000025", "sz000027",
|
|
]
|
|
|
|
|
|
def get_hs300_stocks() -> pd.DataFrame:
|
|
"""Fetch the current CSI 300 constituents from baostock.
|
|
|
|
Returns:
|
|
DataFrame with columns ``code`` (e.g. ``sh600000``), ``name``, ``date``.
|
|
"""
|
|
bs.login()
|
|
try:
|
|
rs = bs.query_hs300_stocks()
|
|
stocks = []
|
|
while rs.next():
|
|
stocks.append(rs.get_row_data())
|
|
finally:
|
|
bs.logout()
|
|
df = pd.DataFrame(stocks, columns=["code", "name", "date"])
|
|
df["code"] = df["code"].str.replace(".", "", regex=False)
|
|
return df
|
|
|
|
|
|
def get_zz500_stocks() -> pd.DataFrame:
|
|
"""Fetch the current CSI 500 (ZZ500) constituents from baostock.
|
|
|
|
Returns:
|
|
DataFrame with columns ``code`` (e.g. ``sh600006``), ``name``, ``date``.
|
|
"""
|
|
bs.login()
|
|
try:
|
|
rs = bs.query_zz500_stocks()
|
|
stocks = []
|
|
while rs.next():
|
|
stocks.append(rs.get_row_data())
|
|
finally:
|
|
bs.logout()
|
|
df = pd.DataFrame(stocks, columns=["code", "name", "date"])
|
|
df["code"] = df["code"].str.replace(".", "", regex=False)
|
|
return df
|