diff --git a/sanguo_data/datareader.py b/sanguo_data/datareader.py new file mode 100644 index 0000000..0fef4d9 --- /dev/null +++ b/sanguo_data/datareader.py @@ -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", + ) diff --git a/tests/data/conftest.py b/tests/data/conftest.py index c387191..0363402 100644 --- a/tests/data/conftest.py +++ b/tests/data/conftest.py @@ -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 pytest diff --git a/tests/data/test_datareader.py b/tests/data/test_datareader.py new file mode 100644 index 0000000..384869f --- /dev/null +++ b/tests/data/test_datareader.py @@ -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 diff --git a/tests/data/vnpy_mock.py b/tests/data/vnpy_mock.py new file mode 100644 index 0000000..7d267c5 --- /dev/null +++ b/tests/data/vnpy_mock.py @@ -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