mirror of
https://ghfast.top/https://github.com/aeroxw/easy-tdx.git
synced 2026-09-12 16:54:17 +08:00
feat: concurrent scanning with ProcessPoolExecutor
- Add workers param to SignalScanner.scan() (default=0 for serial) - workers=2+ uses ProcessPoolExecutor for parallel .day file processing - Extract _scan_one_file as top-level function for pickle compatibility - Results identical between serial and parallel modes - Add 4 tests with synthetic .day file fixtures
This commit is contained in:
@@ -87,6 +87,7 @@ class SignalScanner:
|
||||
self,
|
||||
universe: str = "all",
|
||||
progress_callback: Any = None,
|
||||
workers: int = 0,
|
||||
) -> list[ScanResult]:
|
||||
"""扫描全市场,返回触发买入信号的股票列表。
|
||||
|
||||
@@ -97,6 +98,10 @@ class SignalScanner:
|
||||
- "sz": 仅深圳
|
||||
- 文件路径: 每行一个 "市场 代码"(如 "SZ 000001")
|
||||
progress_callback: 进度回调函数(current, total, filename)
|
||||
workers: 并发工作进程数
|
||||
- 0: 串行模式(默认,向后兼容)
|
||||
- 1: 串行但使用 executor 基础设施
|
||||
- 2+: ProcessPoolExecutor 并发执行
|
||||
|
||||
Returns:
|
||||
触发买入信号的 ScanResult 列表
|
||||
@@ -107,9 +112,24 @@ class SignalScanner:
|
||||
if not files:
|
||||
return []
|
||||
|
||||
results: list[ScanResult] = []
|
||||
total = len(files)
|
||||
|
||||
# 串行模式(workers=0,向后兼容)
|
||||
if workers <= 0:
|
||||
return self._scan_serial(files, total, progress_callback)
|
||||
|
||||
# 并发模式
|
||||
return self._scan_parallel(files, total, workers, progress_callback)
|
||||
|
||||
def _scan_serial(
|
||||
self,
|
||||
files: list[tuple[Path, str, str]],
|
||||
total: int,
|
||||
progress_callback: Any,
|
||||
) -> list[ScanResult]:
|
||||
"""串行扫描(原有逻辑)。"""
|
||||
results: list[ScanResult] = []
|
||||
|
||||
for idx, (filepath, market, code) in enumerate(files):
|
||||
if progress_callback:
|
||||
progress_callback(idx, total, filepath.name)
|
||||
@@ -119,7 +139,6 @@ class SignalScanner:
|
||||
if result is not None:
|
||||
results.append(result)
|
||||
except Exception:
|
||||
# 单个文件出错不中断整体扫描
|
||||
continue
|
||||
|
||||
if progress_callback:
|
||||
@@ -127,6 +146,49 @@ class SignalScanner:
|
||||
|
||||
return results
|
||||
|
||||
def _scan_parallel(
|
||||
self,
|
||||
files: list[tuple[Path, str, str]],
|
||||
total: int,
|
||||
workers: int,
|
||||
progress_callback: Any,
|
||||
) -> list[ScanResult]:
|
||||
"""并发扫描(ProcessPoolExecutor)。"""
|
||||
import concurrent.futures
|
||||
|
||||
# 构建参数:每个任务需要的独立数据
|
||||
tasks = [
|
||||
(str(filepath), market, code, self._strategy_cls, self._cash, self._commission)
|
||||
for filepath, market, code in files
|
||||
]
|
||||
|
||||
results: list[ScanResult] = []
|
||||
completed = 0
|
||||
|
||||
with concurrent.futures.ProcessPoolExecutor(max_workers=workers) as executor:
|
||||
future_to_idx = {
|
||||
executor.submit(_scan_one_file, *task): i for i, task in enumerate(tasks)
|
||||
}
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_idx):
|
||||
completed += 1
|
||||
idx = future_to_idx[future]
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(completed, total, files[idx][0].name)
|
||||
|
||||
try:
|
||||
result = future.result()
|
||||
if result is not None:
|
||||
results.append(result)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(total, total, "done")
|
||||
|
||||
return results
|
||||
|
||||
def _collect_files(self, universe: str) -> list[tuple[Path, str, str]]:
|
||||
"""收集需要扫描的 .day 文件列表。
|
||||
|
||||
@@ -318,3 +380,60 @@ def _bars_to_df(bars: list[Any]) -> pd.DataFrame:
|
||||
)
|
||||
|
||||
return pd.DataFrame(rows)
|
||||
|
||||
|
||||
def _scan_one_file(
|
||||
filepath: str,
|
||||
market: str,
|
||||
code: str,
|
||||
strategy_cls: type[Strategy],
|
||||
cash: float,
|
||||
commission: float,
|
||||
) -> ScanResult | None:
|
||||
"""顶层扫描函数(供 ProcessPoolExecutor 调用)。
|
||||
|
||||
将实例方法逻辑提取为独立函数,避免 pickle 实例方法的问题。
|
||||
逻辑与 SignalScanner._scan_one 完全一致。
|
||||
|
||||
Args:
|
||||
filepath: .day 文件路径字符串
|
||||
market: 市场代码(SZ/SH)
|
||||
code: 6 位股票代码
|
||||
strategy_cls: Strategy 子类
|
||||
cash: 初始资金
|
||||
commission: 佣金率
|
||||
|
||||
Returns:
|
||||
ScanResult 如果触发信号,否则 None
|
||||
"""
|
||||
bars = read_daily_bars(filepath)
|
||||
if len(bars) < 30:
|
||||
return None
|
||||
|
||||
df = _bars_to_df(bars)
|
||||
if df.empty:
|
||||
return None
|
||||
|
||||
try:
|
||||
factor_signals = extract_factor_signals(
|
||||
strategy_cls,
|
||||
df,
|
||||
cash=cash,
|
||||
commission=commission,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
if not factor_signals.buy_mask[-1]:
|
||||
return None
|
||||
|
||||
last_bar = bars[-1]
|
||||
signal_date = last_bar.year * 10000 + last_bar.month * 100 + last_bar.day
|
||||
last_close = last_bar.close
|
||||
|
||||
return ScanResult(
|
||||
code=code,
|
||||
market=market,
|
||||
signal_date=signal_date,
|
||||
last_close=last_close,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""单元测试:信号扫描引擎。
|
||||
|
||||
测试 SignalScanner 的并发扫描和增量扫描功能。
|
||||
使用临时目录构造 .day 文件 fixture,无需真实数据。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from easy_tdx.backtest.strategy import Strategy
|
||||
from easy_tdx.screen.scanner import SignalScanner
|
||||
|
||||
|
||||
class AlwaysBuyStrategy(Strategy):
|
||||
"""策略:每个 bar 都产生买入信号(用于扫描测试)。"""
|
||||
|
||||
def init(self) -> None:
|
||||
pass
|
||||
|
||||
def next(self) -> None:
|
||||
self.buy(size=0)
|
||||
|
||||
|
||||
def _write_day_file(
|
||||
path: Path,
|
||||
n_bars: int = 50,
|
||||
base_price: float = 10.0,
|
||||
) -> None:
|
||||
"""写一个最小的 .day 文件(通达信日线格式)。
|
||||
|
||||
格式: date(I) open(I) high(I) low(I) close(I) amount(f) vol(I) reserved(I)
|
||||
每条 32 字节, 小端序. 价格以 0.01 为系数存储.
|
||||
"""
|
||||
dates = pd.date_range("2024-01-01", periods=n_bars, freq="D")
|
||||
|
||||
with open(path, "wb") as f:
|
||||
for i in range(n_bars):
|
||||
dt = dates[i]
|
||||
day = dt.year * 10000 + dt.month * 100 + dt.day
|
||||
price = base_price + i * 0.01
|
||||
f.write(
|
||||
struct.pack(
|
||||
"<IIIIIfII",
|
||||
day,
|
||||
int(price * 100),
|
||||
int((price + 0.5) * 100),
|
||||
int((price - 0.5) * 100),
|
||||
int(price * 100),
|
||||
float(1000000 + i * 100),
|
||||
10000 + i * 10,
|
||||
0,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vipdoc(tmp_path: Path) -> Path:
|
||||
"""创建包含 .day 文件的临时 vipdoc 目录."""
|
||||
sz_lday = tmp_path / "sz" / "lday"
|
||||
sz_lday.mkdir(parents=True)
|
||||
|
||||
for code in ("000001", "000002", "000003"):
|
||||
_write_day_file(sz_lday / f"sz{code}.day", n_bars=50)
|
||||
|
||||
# 指数文件 (应被过滤)
|
||||
_write_day_file(sz_lday / "sz399001.day", n_bars=50)
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestConcurrentScan:
|
||||
"""测试并发扫描."""
|
||||
|
||||
def test_scan_produces_results(self, vipdoc: Path) -> None:
|
||||
"""基本扫描应返回触发信号的股票."""
|
||||
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
|
||||
results = scanner.scan(universe="all")
|
||||
|
||||
assert len(results) >= 1, f"Expected >= 1 result, got {len(results)}"
|
||||
|
||||
def test_concurrent_same_as_serial(self, vipdoc: Path) -> None:
|
||||
"""并发扫描结果应与串行扫描一致."""
|
||||
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
|
||||
|
||||
serial = scanner.scan(universe="all", workers=1)
|
||||
parallel = scanner.scan(universe="all", workers=2)
|
||||
|
||||
serial_codes = sorted(r.code for r in serial)
|
||||
parallel_codes = sorted(r.code for r in parallel)
|
||||
assert serial_codes == parallel_codes
|
||||
|
||||
def test_scan_with_zero_workers_uses_serial(self, vipdoc: Path) -> None:
|
||||
"""workers=0 应退回串行模式."""
|
||||
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
|
||||
results = scanner.scan(universe="all", workers=0)
|
||||
|
||||
assert len(results) >= 1
|
||||
|
||||
def test_progress_callback(self, vipdoc: Path) -> None:
|
||||
"""进度回调应被正确调用."""
|
||||
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
|
||||
progress: list[tuple[int, int, str]] = []
|
||||
|
||||
def on_progress(current: int, total: int, name: str) -> None:
|
||||
progress.append((current, total, name))
|
||||
|
||||
scanner.scan(universe="all", progress_callback=on_progress)
|
||||
|
||||
assert len(progress) >= 2
|
||||
assert progress[-1][2] == "done"
|
||||
Reference in New Issue
Block a user