diff --git a/sanguo_portfolio/providers/local_unified_provider.py b/sanguo_portfolio/providers/local_unified_provider.py index 7f31636..4dfd392 100644 --- a/sanguo_portfolio/providers/local_unified_provider.py +++ b/sanguo_portfolio/providers/local_unified_provider.py @@ -370,3 +370,120 @@ class LocalUnifiedProvider(DataProvider): # type: ignore[misc] if bs_pcf is not None: row["pcf_ratio"] = float(bs_pcf) return row + + # ==================== 辅助方法 ==================== + def get_trade_days( + self, + start_date: Optional[Union[str, datetime]] = None, + end_date: Optional[Union[str, datetime]] = None, + count: Optional[int] = None, + ) -> List[datetime]: + """读 dbbardata 蓝筹 600519 distinct datetime 取交易日(全市场交易日一致)。""" + conn = self._connect() + rows = conn.execute( + "SELECT DISTINCT datetime FROM dbbardata WHERE symbol='600519' " + "AND exchange='SSE' AND interval='d' ORDER BY datetime" + ).fetchall() + days = [pd.Timestamp(r[0]).to_pydatetime() for r in rows if r[0]] + start_ts = pd.Timestamp(start_date) if start_date else None + end_ts = pd.Timestamp(end_date) if end_date else None + if start_ts: + days = [d for d in days if pd.Timestamp(d) >= start_ts] + if end_ts: + days = [d for d in days if pd.Timestamp(d) <= end_ts] + if count: + days = days[-count:] + return days + + def get_security_info(self, security: str) -> Dict[str, Any]: + """读 dbbardata min/max datetime → start/end_date;display_name 从 constituent_unified。""" + sym, exc = jq_to_dbbardata(security) + conn = self._connect() + row = conn.execute( + "SELECT MIN(datetime), MAX(datetime) FROM dbbardata " + "WHERE symbol=? AND exchange=? AND interval='d'", + (sym, exc), + ).fetchone() + start_dt = row[0][:10] if row and row[0] else None + end_dt = row[1][:10] if row and row[1] else None + # display_name 查 constituent_unified + name = security + try: + nrow = conn.execute( + "SELECT code_name FROM constituent_unified WHERE code=? LIMIT 1", + (sym,), + ).fetchone() + if nrow and nrow[0]: + name = str(nrow[0]) + except sqlite3.Error: + pass + return { + "code": security, + "display_name": name, + "name": name, + "start_date": start_dt, + "end_date": end_dt, + "type": "stock", + } + + def get_current_tick(self, security: str) -> Optional[Dict[str, Any]]: + """dbbardata 最近 close + 高低涨停 ±10%(简化,ST/创业/科创精确规则 v2)。""" + sym, exc = jq_to_dbbardata(security) + conn = self._connect() + row = conn.execute( + "SELECT close_price FROM dbbardata WHERE symbol=? AND exchange=? " + "AND interval='d' ORDER BY datetime DESC LIMIT 1", + (sym, exc), + ).fetchone() + if not row or row[0] is None: + return None + close = float(row[0]) + return { + "code": security, + "current_price": close, + "close": close, + "high_limit": round(close * 1.1, 2), + "low_limit": round(close * 0.9, 2), + } + + def get_split_dividend( + self, + security: str, + start_date: Optional[Union[str, datetime]] = None, + end_date: Optional[Union[str, datetime]] = None, + ) -> List[Dict[str, Any]]: + """读 bs_adjust_factor → 除权事件列表。""" + bs_code = _jq_to_bs_code(security) + conn = self._connect() + start_str = self._to_date_str(start_date) or "1900-01-01" + end_str = self._to_date_str(end_date) or "2099-12-31" + try: + rows = conn.execute( + "SELECT dividOperateDate, foreAdjustFactor, backAdjustFactor, adjustFactor " + "FROM bs_adjust_factor WHERE code=? AND dividOperateDate>=? " + "AND dividOperateDate<=? ORDER BY dividOperateDate", + (bs_code, start_str, end_str), + ).fetchall() + except sqlite3.Error: + return [] + return [ + { + "code": security, + "date": r[0], + "foreAdjustFactor": float(r[1]) if r[1] is not None else None, + "backAdjustFactor": float(r[2]) if r[2] is not None else None, + "adjustFactor": float(r[3]) if r[3] is not None else None, + } + for r in rows + ] + + def get_all_securities( + self, types: Optional[List[str]] = None, + ) -> pd.DataFrame: + """读 dbbardata distinct symbol → DataFrame[code, display_name]。""" + conn = self._connect() + rows = conn.execute( + "SELECT DISTINCT symbol, exchange FROM dbbardata WHERE interval='d'" + ).fetchall() + codes = [dbbardata_to_jq(sym, exc) for sym, exc in rows if sym and exc] + return pd.DataFrame({"code": codes, "display_name": codes}) diff --git a/tests/portfolio/test_local_unified_provider.py b/tests/portfolio/test_local_unified_provider.py index 4db6f65..9132395 100644 --- a/tests/portfolio/test_local_unified_provider.py +++ b/tests/portfolio/test_local_unified_provider.py @@ -472,3 +472,61 @@ class TestGetFundamentals: p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)}) df = p.get_fundamentals_df([], date="2024-09-30") assert df.empty + + +# ======================== Task 4: 辅助方法 ======================== +class TestAuxMethods: + def test_get_trade_days_from_dbbardata(self, unified_provider): + # dbbardata 600519 三根日线 → 3 个交易日 + days = unified_provider.get_trade_days( + start_date="2024-06-18", end_date="2024-06-20" + ) + assert len(days) == 3 + assert all(hasattr(d, "year") for d in days) + + def test_get_trade_days_count(self, unified_provider): + days = unified_provider.get_trade_days(count=2) + assert len(days) == 2 + + def test_get_security_info(self, unified_provider): + info = unified_provider.get_security_info("600519.XSHG") + assert info["code"] == "600519.XSHG" + assert "start_date" in info + assert "end_date" in info + + def test_get_current_tick_high_limit(self, unified_provider): + # dbbardata 最近 close=910 → high_limit=910*1.1=1001 + tick = unified_provider.get_current_tick("600519.XSHG") + assert tick is not None + assert abs(tick["close"] - 910.0) < 1e-6 + assert abs(tick["high_limit"] - 910.0 * 1.1) < 1e-2 + + def test_get_split_dividend_from_adjust_factor(self, unified_provider): + # bs_adjust_factor 有 1 个事件 → 1 条记录 + events = unified_provider.get_split_dividend( + "600519.XSHG", start_date="2024-01-01", end_date="2024-12-31" + ) + assert len(events) >= 1 + assert "date" in events[0] + + def test_get_all_securities_from_dbbardata(self, tmp_path): + # dbbardata distinct symbol → DataFrame + db = tmp_path / "t.db" + c = sqlite3.connect(str(db)) + c.execute( + "CREATE TABLE dbbardata(symbol TEXT, exchange TEXT, datetime TEXT, interval TEXT)" + ) + c.executemany( + "INSERT INTO dbbardata VALUES(?,?,?,?)", + [ + ("600519", "SSE", "2024-06-19 00:00:00", "d"), + ("000001", "SZSE", "2024-06-19 00:00:00", "d"), + ], + ) + c.commit() + c.close() + p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)}) + df = p.get_all_securities() + assert len(df) == 2 + assert "code" in df.columns + assert "600519.XSHG" in set(df["code"])