feat(portfolio): LocalUnifiedProvider 成份股并集治偏差(Task2)

sanguo_portfolio/providers/local_unified_provider.py

@@ -228,3 +228,45 @@ class LocalUnifiedProvider(DataProvider):  # type: ignore[misc]
         if len(frames) == 1:
             return next(iter(frames.values()))
         return pd.concat(frames, axis=1)
+
+    # ==================== get_index_stocks (constituent_unified 并集) ====================
+    def get_index_stocks(
+        self,
+        index_symbol: str,
+        date: Optional[Union[str, datetime]] = None,
+    ) -> List[str]:
+        """读 ``constituent_unified`` 并集(``in_current=1 OR was_removed=1``)。
+
+        ⚠️ 并集模型: 表无 date 列,治"纯当前"幸存者偏差(含已退市/被踢),
+        但有轻微前视(使用说明标注)。``date`` 参数忽略(无时点数据)。
+        """
+        idx = (
+            index_symbol.split(".")[0]
+            if "." in str(index_symbol) else str(index_symbol)
+        )
+        conn = self._connect()
+        try:
+            rows = conn.execute(
+                "SELECT code FROM constituent_unified WHERE index_code=? "
+                "AND (in_current=1 OR was_removed=1)",
+                (idx,),
+            ).fetchall()
+        except sqlite3.Error as exc:
+            logger.warning("constituent_unified 查询失败 %s: %s", idx, exc)
+            return []
+        out: List[str] = []
+        for (code,) in rows:
+            code_str = str(code).strip()
+            if len(code_str) != 6 or not code_str.isdigit():
+                continue
+            exc_name = "SSE" if code_str.startswith("6") else "SZSE"
+            out.append(dbbardata_to_jq(code_str, exc_name))
+        return out
+
+    def get_constituent(
+        self,
+        index: str,
+        date: Optional[Union[str, datetime]] = None,
+    ) -> List[str]:
+        """spec §6 语义别名 = ``get_index_stocks``。"""
+        return self.get_index_stocks(index, date)

tests/portfolio/test_local_unified_provider.py

@@ -280,3 +280,63 @@ class TestGetPrice:
         )
         assert len(df) == 2
         assert set(df["code"]) == {"600519.XSHG", "000001.XSHE"}
+
+
+# ======================== Task 2: get_index_stocks ========================
+def _make_constituent_db(tmp_path):
+    db = tmp_path / "t.db"
+    c = sqlite3.connect(str(db))
+    c.execute(
+        "CREATE TABLE constituent_unified(index_code TEXT, code TEXT, code_name TEXT, "
+        "source TEXT, in_current INT, was_removed INT)"
+    )
+    c.executemany(
+        "INSERT INTO constituent_unified VALUES(?,?,?,?,?,?)",
+        [
+            ("000300", "600519", "贵州茅台", "baostock", 1, 0),
+            ("000300", "000001", "平安银行", "baostock", 1, 0),
+            ("000300", "600811", "退市股", "baostock", 0, 1),  # 被踢
+        ],
+    )
+    c.commit()
+    c.close()
+    return db
+
+
+class TestGetIndexStocks:
+    def test_union_of_current_and_removed(self, tmp_path):
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        stocks = p.get_index_stocks("000300.XSHG", "2020-01-01")
+        # 并集含被踢(was_removed=1)
+        assert "600519.XSHG" in stocks
+        assert "000001.XSHE" in stocks
+        assert "600811.XSHG" in stocks  # 6 开头 → SSE
+        assert len(stocks) == 3
+
+    def test_date_param_ignored_union_model(self, tmp_path):
+        # 并集模型 — date 参数不报错不过滤
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        s1 = p.get_index_stocks("000300", "2010-01-01")
+        s2 = p.get_index_stocks("000300", "2024-12-31")
+        assert set(s1) == set(s2)
+
+    def test_get_constituent_is_alias(self, tmp_path):
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        a = p.get_index_stocks("000300.XSHG", "2020-01-01")
+        b = p.get_constituent("000300", "2020-01-01")
+        assert a == b
+
+    def test_index_not_found_returns_empty(self, tmp_path):
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        assert p.get_index_stocks("999999.XSHG", "2024-01-01") == []
+
+    def test_pure_digit_index_code(self, tmp_path):
+        # 纯数字 index_symbol 也能查
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        stocks = p.get_index_stocks("000300", "2020-01-01")
+        assert len(stocks) == 3