feat: CSI500 universe + reversal_vol signal + argparse CLI

data/universe.py

@@ -1,4 +1,4 @@
-"""CSI 300 (HS300) universe helpers."""
+"""CSI 300 (HS300) and CSI 500 (ZZ500) universe helpers."""
 import logging
 
 import baostock as bs
@@ -19,6 +19,20 @@ SYMBOLS = [
 ]
 
 
+# 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.
 
@@ -36,3 +50,22 @@ def get_hs300_stocks() -> pd.DataFrame:
     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