feat(data): DataReader parquet 读取 → BarData
This commit is contained in:
@@ -0,0 +1,39 @@
|
|||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Add vnpy source to path for imports
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent.parent / "tests" / "data"))
|
||||||
|
|
||||||
|
import pandas as pd
|
||||||
|
from datetime import datetime
|
||||||
|
from vnpy_mock import BarData, Exchange, Interval
|
||||||
|
|
||||||
|
def read_parquet_daily(symbol: str, start: str, end: str, cfg) -> list[BarData]:
|
||||||
|
daily_dir = Path(cfg.data_paths["daily_dir"])
|
||||||
|
start_dt = datetime.strptime(start, "%Y-%m-%d")
|
||||||
|
end_dt = datetime.strptime(end, "%Y-%m-%d")
|
||||||
|
bars: list[BarData] = []
|
||||||
|
for year in range(start_dt.year, end_dt.year + 1):
|
||||||
|
f = daily_dir / str(year) / f"{symbol}.parquet"
|
||||||
|
if not f.exists():
|
||||||
|
continue
|
||||||
|
df = pd.read_parquet(f)
|
||||||
|
for _, row in df.iterrows():
|
||||||
|
d = pd.to_datetime(row["date"])
|
||||||
|
if start_dt <= d <= end_dt:
|
||||||
|
bars.append(_row_to_bar(symbol, row, Interval.DAILY))
|
||||||
|
return bars
|
||||||
|
|
||||||
|
def _row_to_bar(symbol: str, row, interval: Interval) -> BarData:
|
||||||
|
return BarData(
|
||||||
|
symbol=symbol,
|
||||||
|
exchange=Exchange.SSE, # Task 4 改为 guess_exchange
|
||||||
|
datetime=pd.to_datetime(row["date"]).to_pydatetime(),
|
||||||
|
interval=interval,
|
||||||
|
open_price=float(row["open"]),
|
||||||
|
high_price=float(row["high"]),
|
||||||
|
low_price=float(row["low"]),
|
||||||
|
close_price=float(row["close"]),
|
||||||
|
volume=float(row["volume"]),
|
||||||
|
gateway_name="DATA",
|
||||||
|
)
|
||||||
@@ -1,3 +1,9 @@
|
|||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Add vnpy mock to path for imports
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent))
|
||||||
|
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,24 @@
|
|||||||
|
import pandas as pd
|
||||||
|
from vnpy_mock import Exchange, Interval
|
||||||
|
from sanguo_data.config import DataConfig
|
||||||
|
from sanguo_data.datareader import read_parquet_daily
|
||||||
|
|
||||||
|
def test_read_parquet_daily_returns_bardata(tmp_path):
|
||||||
|
year_dir = tmp_path / "2026"
|
||||||
|
year_dir.mkdir()
|
||||||
|
df = pd.DataFrame({
|
||||||
|
"date": ["2026-01-05", "2026-01-06"],
|
||||||
|
"open": [10.0, 11.0], "high": [10.5, 11.5],
|
||||||
|
"low": [9.8, 10.8], "close": [10.2, 11.2],
|
||||||
|
"volume": [10000, 12000],
|
||||||
|
})
|
||||||
|
df.to_parquet(year_dir / "600000.parquet")
|
||||||
|
|
||||||
|
cfg = DataConfig(
|
||||||
|
data_paths={"daily_dir": str(tmp_path)},
|
||||||
|
data_sources={}, validation={}, performance={},
|
||||||
|
)
|
||||||
|
bars = read_parquet_daily("600000", "2026-01-01", "2026-12-31", cfg)
|
||||||
|
assert len(bars) == 2
|
||||||
|
assert bars[0].symbol == "600000"
|
||||||
|
assert bars[0].open_price == 10.0
|
||||||
@@ -0,0 +1,34 @@
|
|||||||
|
"""Minimal vnpy mocks for testing"""
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import datetime
|
||||||
|
from enum import Enum
|
||||||
|
|
||||||
|
class Exchange(Enum):
|
||||||
|
SSE = "SSE"
|
||||||
|
SZSE = "SZSE"
|
||||||
|
SHFE = "SHFE"
|
||||||
|
DCE = "DCE"
|
||||||
|
CZCE = "CZCE"
|
||||||
|
CFFEX = "CFFEX"
|
||||||
|
GFEX = "GFEX"
|
||||||
|
INE = "INE"
|
||||||
|
|
||||||
|
class Interval(Enum):
|
||||||
|
MINUTE = "1m"
|
||||||
|
HOUR = "1h"
|
||||||
|
DAILY = "d"
|
||||||
|
WEEKLY = "w"
|
||||||
|
MONTHLY = "M"
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class BarData:
|
||||||
|
symbol: str
|
||||||
|
exchange: Exchange
|
||||||
|
datetime: datetime
|
||||||
|
interval: Interval
|
||||||
|
open_price: float
|
||||||
|
high_price: float
|
||||||
|
low_price: float
|
||||||
|
close_price: float
|
||||||
|
volume: float
|
||||||
|
gateway_name: str
|
||||||
Reference in New Issue
Block a user