diff --git a/src/easy_tdx/screen/scanner.py b/src/easy_tdx/screen/scanner.py index 6e9fb2d..162a1ec 100644 --- a/src/easy_tdx/screen/scanner.py +++ b/src/easy_tdx/screen/scanner.py @@ -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, + ) diff --git a/tests/unit/test_scanner.py b/tests/unit/test_scanner.py new file mode 100644 index 0000000..6b3c5bb --- /dev/null +++ b/tests/unit/test_scanner.py @@ -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( + " 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"