feat(portfolio): LocalUnifiedProvider get_price+前复权(Task1)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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"}
|
||||
|
||||
Reference in New Issue
Block a user