diff --git a/sanguo_portfolio/providers/local_unified_provider.py b/sanguo_portfolio/providers/local_unified_provider.py index 8f0318d..080333a 100644 --- a/sanguo_portfolio/providers/local_unified_provider.py +++ b/sanguo_portfolio/providers/local_unified_provider.py @@ -118,3 +118,113 @@ class LocalUnifiedProvider(DataProvider): # type: ignore[misc] self.db_path: str = cfg.get("db_path", _DEFAULT_DB) self.data_dir: str = cfg.get("data_dir", _DEFAULT_DATA_DIR) self._conn: Optional[sqlite3.Connection] = None + self._val_bs_cache: Dict[int, pd.DataFrame] = {} # year -> valuation_baostock + self._lpp_helper: Any = None + + def _connect(self) -> sqlite3.Connection: + """惰性连接 dbbardata sqlite(单连接复用)。""" + if self._conn is None: + self._conn = sqlite3.connect(self.db_path, timeout=30) + self._conn.execute("PRAGMA busy_timeout = 30000") + return self._conn + + @staticmethod + def _to_date_str(value: Optional[Union[str, datetime]]) -> Optional[str]: + if value is None: + return None + if isinstance(value, str): + return value[:10] + try: + return value.strftime("%Y-%m-%d") + except AttributeError: + return str(value)[:10] + + # ==================== get_price ==================== + def get_price( + self, + security: Union[str, List[str]], + start_date: Optional[Union[str, datetime]] = None, + end_date: Optional[Union[str, datetime]] = None, + frequency: str = "daily", + fields: Optional[List[str]] = None, + skip_paused: bool = False, + fq: str = "raw", + count: Optional[int] = None, + panel: bool = True, + fill_paused: bool = True, + **kwargs: Any, + ) -> pd.DataFrame: + """读 ``dbbardata('d')`` raw 日线,按需 ``bs_adjust_factor`` 算前复权。 + + 策略契约(all_weather 实证): + - ``panel=False`` 返长表含 ``time`` + ``code`` 列(供 pivot) + - ``fields`` 里缺失列(如 ``high_limit``)补 NaN(降级) + - ``frequency`` 非 day/1d/d → 返空 DataFrame(1m 数据层无) + """ + freq = str(frequency or "").lower() + if freq not in ("daily", "day", "1d", "d"): + return pd.DataFrame() + secs: List[str] = [security] if isinstance(security, str) else list(security or []) + if not secs: + return pd.DataFrame() + conn = self._connect() + start_str = self._to_date_str(start_date) or "1990-01-01" + end_str = self._to_date_str(end_date) or datetime.now().strftime("%Y-%m-%d") + + frames: Dict[str, pd.DataFrame] = {} + for jq_code in secs: + sym, exc = jq_to_dbbardata(jq_code) + q = ( + "SELECT datetime, open_price, high_price, low_price, close_price, " + "volume, turnover FROM dbbardata WHERE symbol=? AND exchange=? " + "AND interval='d' AND datetime>=? AND datetime<=? ORDER BY datetime" + ) + df = pd.read_sql( + q, conn, + params=(sym, exc, start_str + " 00:00:00", end_str + " 23:59:59"), + ) + if df.empty: + frames[jq_code] = df + continue + df["datetime"] = pd.to_datetime(df["datetime"]) + df = df.set_index("datetime") + df.index.name = None + if count: + df = df.tail(count) + # 前复权 + if fq in ("qfq", "pre", "前复权") and not df.empty: + factor = _build_qfq_factor(_jq_to_bs_code(jq_code), conn, pd.Series(df.index)) + for col in ("open_price", "high_price", "low_price", "close_price"): + df[col] = df[col].values * factor.values + # jq 风格字段重命名 + df = df.rename(columns={ + "open_price": "open", "high_price": "high", + "low_price": "low", "close_price": "close", + }) + # 缺失字段(如 high_limit)补 NaN + if fields: + for f in fields: + if f not in df.columns: + df[f] = float("nan") + df = df[[f for f in fields if f in df.columns]] + frames[jq_code] = df + + if not frames or all(f.empty for f in frames.values()): + return pd.DataFrame() + if not panel: + parts: List[pd.DataFrame] = [] + for jq_code, df in frames.items(): + if df.empty: + continue + d = df.reset_index() + # index.name=None 时 reset_index 出 'index' 列; 统一改名 'time' + if "index" in d.columns and "time" not in d.columns: + d = d.rename(columns={"index": "time"}) + elif "datetime" in d.columns: + d = d.rename(columns={"datetime": "time"}) + d.insert(0, "code", jq_code) + parts.append(d) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame() + if len(frames) == 1: + return next(iter(frames.values())) + return pd.concat(frames, axis=1) diff --git a/tests/portfolio/test_local_unified_provider.py b/tests/portfolio/test_local_unified_provider.py index 5be1804..3474ad0 100644 --- a/tests/portfolio/test_local_unified_provider.py +++ b/tests/portfolio/test_local_unified_provider.py @@ -143,3 +143,140 @@ class TestBuildQfqFactor: assert len(f) == 2 assert abs(f.iloc[0] - 1.0) < 1e-6 assert abs(f.iloc[1] - 1.0) < 1e-6 + + +# ======================== Task 1: get_price fixture ======================== +@pytest.fixture +def unified_provider(tmp_path): + """造小样本 sqlite fixture: dbbardata 日线 + bs_adjust_factor。""" + db = tmp_path / "quant_trading.db" + c = sqlite3.connect(str(db)) + c.execute( + "CREATE TABLE dbbardata(symbol TEXT, exchange TEXT, datetime TEXT, " + "interval TEXT, volume REAL, turnover REAL, open_interest REAL, " + "open_price REAL, high_price REAL, low_price REAL, close_price REAL)" + ) + # 600519: 除权日 2024-06-19 raw close 1000 → 900 跳水 + rows = [ + ("600519", "SSE", "2024-06-18 00:00:00", "d", 1000, 1e6, 0, + 1000.0, 1010.0, 990.0, 1000.0), + ("600519", "SSE", "2024-06-19 00:00:00", "d", 1000, 1e6, 0, + 900.0, 910.0, 890.0, 900.0), + ("600519", "SSE", "2024-06-20 00:00:00", "d", 1000, 1e6, 0, + 910.0, 920.0, 900.0, 910.0), + ] + c.executemany("INSERT INTO dbbardata VALUES(?,?,?,?,?,?,?,?,?,?,?)", rows) + # 复权因子: 2024-06-19 起除权, factor=0.9 + c.execute( + "CREATE TABLE bs_adjust_factor(code TEXT, dividOperateDate TEXT, " + "foreAdjustFactor REAL, backAdjustFactor REAL, adjustFactor REAL)" + ) + c.execute( + "INSERT INTO bs_adjust_factor VALUES('sh.600519','2024-06-19',0.9,0,0)" + ) + c.commit() + c.close() + return LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)}) + + +# ======================== Task 1: get_price daily ======================== +class TestGetPrice: + def test_raw_keeps_original_prices(self, unified_provider): + # raw: 除权日 close=900 跳水(原值) + df = unified_provider.get_price( + "600519.XSHG", + start_date="2024-06-18", + end_date="2024-06-20", + fq="raw", + ) + assert len(df) == 3 + assert abs(df.loc["2024-06-19", "close"] - 900.0) < 1e-6 + + def test_qfq_earlier_date_uses_earliest_factor(self, unified_provider): + # 2024-06-18 早于除权日 06-19 → factor=0.9 → 1000*0.9=900 + df = unified_provider.get_price( + "600519.XSHG", + start_date="2024-06-18", + end_date="2024-06-20", + fq="qfq", + ) + assert abs(df.loc["2024-06-18", "close"] - 900.0) < 1e-6 + + def test_qfq_after_event_uses_event_factor(self, unified_provider): + # 2024-06-19/20 ≥ 除权日 → factor=0.9 → 900*0.9=810, 910*0.9=819 + df = unified_provider.get_price( + "600519.XSHG", + start_date="2024-06-19", + end_date="2024-06-20", + fq="qfq", + ) + assert abs(df.loc["2024-06-19", "close"] - 810.0) < 1e-6 + assert abs(df.loc["2024-06-20", "close"] - 819.0) < 1e-6 + + def test_panel_false_returns_long_table(self, unified_provider): + # panel=False → 长表含 time + code 列 + df = unified_provider.get_price( + "600519.XSHG", + end_date="2024-06-20", + count=2, + panel=False, + fields=["close"], + ) + assert "code" in df.columns + assert "time" in df.columns + assert len(df) == 2 + assert "600519.XSHG" in set(df["code"]) + + def test_fields_with_missing_column_fills_nan(self, unified_provider): + # high_limit 不在 dbbardata → NaN 降级(策略 prepare_stock_list 涨停识别降级) + df = unified_provider.get_price( + "600519.XSHG", + end_date="2024-06-20", + count=1, + panel=False, + fields=["close", "high_limit"], + ) + assert "high_limit" in df.columns + # high_limit NaN(不崩) + assert pd.isna(df.iloc[0]["high_limit"]) or df.iloc[0]["high_limit"] != df.iloc[0]["high_limit"] + + def test_minute_frequency_returns_empty(self, unified_provider): + # 1m 频率无数据 → 返空 DataFrame + df = unified_provider.get_price( + "600519.XSHG", + end_date="2024-06-20", + frequency="1m", + count=1, + panel=False, + ) + assert isinstance(df, pd.DataFrame) + assert df.empty + + def test_multi_stocks_panel_false(self, tmp_path): + # 多股 panel=False → 长表含 code 列区分 + db = tmp_path / "t.db" + c = sqlite3.connect(str(db)) + c.execute( + "CREATE TABLE dbbardata(symbol TEXT, exchange TEXT, datetime TEXT, " + "interval TEXT, volume REAL, turnover REAL, open_interest REAL, " + "open_price REAL, high_price REAL, low_price REAL, close_price REAL)" + ) + rows = [ + ("600519", "SSE", "2024-06-19 00:00:00", "d", 1000, 1e6, 0, + 900.0, 910.0, 890.0, 900.0), + ("000001", "SZSE", "2024-06-19 00:00:00", "d", 1000, 1e6, 0, + 10.0, 10.5, 9.8, 10.2), + ] + c.executemany("INSERT INTO dbbardata VALUES(?,?,?,?,?,?,?,?,?,?,?)", rows) + c.commit() + c.close() + p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)}) + df = p.get_price( + ["600519.XSHG", "000001.XSHE"], + end_date="2024-06-19", + count=1, + panel=False, + fields=["close"], + ) + assert len(df) == 2 + assert set(df["code"]) == {"600519.XSHG", "000001.XSHE"}