From b46a0c26c8f9938ba7ae2f6e382d2ae6427b762d Mon Sep 17 00:00:00 2001 From: claude_dev Date: Thu, 23 Jul 2026 08:07:37 +0800 Subject: [PATCH] =?UTF-8?q?feat(portfolio):=20LocalUnifiedProvider=20?= =?UTF-8?q?=E6=88=90=E4=BB=BD=E8=82=A1=E5=B9=B6=E9=9B=86=E6=B2=BB=E5=81=8F?= =?UTF-8?q?=E5=B7=AE(Task2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../providers/local_unified_provider.py | 42 +++++++++++++ .../portfolio/test_local_unified_provider.py | 60 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/sanguo_portfolio/providers/local_unified_provider.py b/sanguo_portfolio/providers/local_unified_provider.py index 080333a..f42fdd4 100644 --- a/sanguo_portfolio/providers/local_unified_provider.py +++ b/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) diff --git a/tests/portfolio/test_local_unified_provider.py b/tests/portfolio/test_local_unified_provider.py index 3474ad0..5c3e3a3 100644 --- a/tests/portfolio/test_local_unified_provider.py +++ b/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