From d01b11fa74cafffe96d83ca2b2bc2c0ee1872f4e Mon Sep 17 00:00:00 2001 From: Justin Gu <97915@qq.com> Date: Sun, 7 Jun 2026 21:13:49 +0800 Subject: [PATCH] feat: add offline data write-back and sync commands, bump to v1.6.0 - Add write_daily.py: encode/append daily bars to .day files - Add write_ex_daily.py: encode/append extended market daily bars - Add write_min_bar.py: encode/append minute bars (.5/.lc1/.lc5) - Add sync-daily CLI: sync single stock with pagination support - Add sync-all CLI: one-command sync for all SH/SZ .day files - Update README with sync commands and Python write API docs - 50 new unit tests covering encode round-trip, append dedup, edge cases - Bump version 1.5.0 -> 1.6.0 --- README.md | 56 +++- pyproject.toml | 4 +- src/easy_tdx/cli/cmd_offline.py | 255 +++++++++++++++++- src/easy_tdx/offline/__init__.py | 45 +++- src/easy_tdx/offline/write_daily.py | 145 +++++++++++ src/easy_tdx/offline/write_ex_daily.py | 93 +++++++ src/easy_tdx/offline/write_min_bar.py | 186 +++++++++++++ tests/unit/test_write_daily.py | 345 +++++++++++++++++++++++++ tests/unit/test_write_ex_daily.py | 190 ++++++++++++++ tests/unit/test_write_min_bar.py | 276 ++++++++++++++++++++ 10 files changed, 1583 insertions(+), 12 deletions(-) create mode 100644 src/easy_tdx/offline/write_daily.py create mode 100644 src/easy_tdx/offline/write_ex_daily.py create mode 100644 src/easy_tdx/offline/write_min_bar.py create mode 100644 tests/unit/test_write_daily.py create mode 100644 tests/unit/test_write_ex_daily.py create mode 100644 tests/unit/test_write_min_bar.py diff --git a/README.md b/README.md index 600109a..17769fe 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) [![PyPI](https://img.shields.io/pypi/v/easy-tdx.svg)](https://pypi.org/project/easy-tdx/) -通达信 TCP 行情协议客户端。支持 A 股、港股、美股、期货全市场；内置 `easy-tdx` CLI 工具，默认 JSON 输出，天然适配 Claude Code、OpenClaw、Hermes 等 AI Agent 工具链。提供同步 + asyncio 双接口；strict mypy 通过；每一层编解码都有离线 fixture 测试覆盖。 +通达信 TCP 行情协议客户端。支持 A 股、港股、美股、期货全市场；内置 `easy-tdx` CLI 工具，默认 JSON 输出，天然适配 Claude Code、OpenClaw、Hermes 等 AI Agent 工具链。提供同步 + asyncio 双接口；支持离线数据读取与写入同步；strict mypy 通过；每一层编解码都有离线 fixture 测试覆盖。 ## 安装 @@ -213,7 +213,7 @@ easy-tdx ex quote-list HK_MAIN_BOARD --table # 港股商品列表 easy-tdx ex tick HK_MAIN_BOARD 00700 --table # 港股分时 ``` -### 离线数据（无需网络） +### 离线数据（读取 + 写入同步） 从本地通达信安装目录直接读取数据文件，无需网络连接： @@ -228,6 +228,19 @@ easy-tdx offline financial C:\new_jyplug\vipdoc\fin\gpcw20260331.dat # 历史 easy-tdx offline blocks C:\new_jyplug\T0002\blocknew --table # 自定义板块 ``` +从服务端获取最新日线并写入本地 .day 文件，替代通达信内置下载功能： + +```bash +# 同步单只股票日线（自动增量/全量） +easy-tdx offline sync-daily SZ 000001 +easy-tdx offline sync-daily SH 600519 --vipdoc C:\new_jyplug\vipdoc + +# 一键同步沪深全市场（每天一条命令） +easy-tdx offline sync-all +``` + +> 建议在通达信关闭时执行 sync 命令，避免文件被锁定。空文件自动全量下载，已有数据只做增量追加。 + ## CLI 命令汇总 | 命令 | 说明 | @@ -261,6 +274,8 @@ easy-tdx offline blocks C:\new_jyplug\T0002\blocknew --table # 自定 | `ex markets` | 列出可用扩展市场 | | `offline home` | 检测通达信安装目录 | | `offline daily` | A 股日线（本地 .day 文件） | +| `offline sync-daily` | 从服务端同步单只股票日线到本地 .day 文件 | +| `offline sync-all` | 一键同步沪深全市场日线（扫描本地 .day 文件） | | `offline min` | 分钟线（本地 .5/.lc1/.lc5 文件） | | `offline ex-files` | 列出扩展市场可用文件 | | `offline ex-daily` | 扩展市场日线（期货/港股/外盘） | @@ -543,6 +558,29 @@ bars = read_daily_bars(filepath) 支持：日线、分钟线、扩展市场日线、板块、股本变迁、历史财务数据。 +### 离线数据写入同步 + +从服务端获取最新数据并追加写入本地通达信数据文件： + +```python +from easy_tdx.offline import ( + encode_daily_bar, append_daily_bars, get_last_bar_date, + encode_5min_bar, append_5min_bars, + encode_lc_min_bar, append_lc_min_bars, +) +from easy_tdx import Market +from easy_tdx.client import TdxClient + +# 追加日线到 .day 文件（自动跳过重复日期） +from easy_tdx.offline import sync_daily_bars_from_security_bars + +# 手动编码单条记录 +bar_bytes = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + +# 获取文件末尾日期 +last_date = get_last_bar_date("C:/new_jyplug/vipdoc/sh/lday/sh600000.day") +``` + v1.5.0 起可通过 CLI 直接使用： ```bash @@ -720,7 +758,7 @@ src/easy_tdx/ ├── commands/ # 标准协议命令（无 IO） ├── codec/ # price / volume / datetime / frame / bitmap 编解码 ├── models/ # 纯 dataclass，无业务逻辑 -├── offline/ # 离线数据读取模块 +├── offline/ # 离线数据读写模块（读取 + 写入同步） └── cli/ # easy-tdx CLI（click） ``` @@ -747,6 +785,18 @@ ruff format --check src/ tests/ # format check ## Changelog +### 1.6.0 (2026-06-07) + +**离线数据写入同步** — 从服务端获取最新日线数据并写入本地通达信 .day 文件，替代通达信内置下载功能。 + +- 新增 `offline sync-daily` CLI 命令：同步单只股票日线，自动增量/全量判断，支持分页获取完整历史 +- 新增 `offline sync-all` CLI 命令：一键扫描沪深全市场 .day 文件并同步 +- 新增 `write_daily.py` 模块：日线编解码（`encode_daily_bar`）、追加写入（`append_daily_bars`）、末尾日期检测 +- 新增 `write_ex_daily.py` 模块：扩展市场日线写入（期货/港股，价格 float32） +- 新增 `write_min_bar.py` 模块：分钟线写入（.5/.lc1/.lc5 格式） +- 写入自动跳过重复日期，空文件自动全量下载，已有数据只做增量追加 +- 50 个新增单元测试覆盖编解码 round-trip、追加去重、边界条件 + ### 1.5.0 (2026-06-02) **离线数据 CLI 命令** — 新增 `offline` 命令组，无需网络即可通过 CLI 读取本地通达信数据文件。 diff --git a/pyproject.toml b/pyproject.toml index a2e8186..4e3805d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "hatchling.build" [project] name = "easy-tdx" -version = "1.5.0" -description = "通达信 TCP 协议行情数据客户端，支持在线行情与离线本地数据读取" +version = "1.6.0" +description = "通达信 TCP 协议行情数据客户端，支持在线行情、离线数据读取与写入同步" readme = "README.md" requires-python = ">=3.10" dependencies = ["pandas>=2.0", "tzdata>=2024.1", "click>=8.0"] diff --git a/src/easy_tdx/cli/cmd_offline.py b/src/easy_tdx/cli/cmd_offline.py index 59d7bb9..eddf56b 100644 --- a/src/easy_tdx/cli/cmd_offline.py +++ b/src/easy_tdx/cli/cmd_offline.py @@ -1,15 +1,25 @@ -"""离线本地数据读取命令（无需网络，读取本地通达信数据文件）。""" +"""离线本地数据读写命令 —— 读取本地通达信数据文件 & 从服务端同步写入。""" from __future__ import annotations +from pathlib import Path +from typing import TYPE_CHECKING + import click +if TYPE_CHECKING: + import pandas as pd + + from ..client import TdxClient + from ..models.bar import SecurityBar + @click.group() def offline() -> None: - """离线本地数据读取（无需网络，读取本地通达信数据文件）。 + """离线本地数据读写（读取本地通达信数据文件 & 从服务端同步写入）。 - 需要本地已安装通达信并下载过对应数据。 + 读取需要本地已安装通达信并下载过对应数据。 + sync-daily 可从服务端获取最新日线并追加到本地 .day 文件。 示例： @@ -22,6 +32,10 @@ def offline() -> None: easy-tdx offline ex-files --table easy-tdx offline ex-daily 29#A1801 --table + + easy-tdx offline sync-daily SZ 000001 + + easy-tdx offline sync-daily SH 600519 --vipdoc C:\\new_jyplug\\vipdoc """ pass @@ -412,3 +426,238 @@ def blocks( for b in result ] print_output(pd.DataFrame(rows), fmt) + + +# --------------------------------------------------------------------------- +# sync-daily：从服务端同步日线到本地 .day 文件 +# --------------------------------------------------------------------------- + + +def _df_to_bars(df: pd.DataFrame) -> list[SecurityBar]: + """将日线 DataFrame 转换为 SecurityBar 列表（按日期升序）。""" + from ..models.bar import SecurityBar + + bars: list[SecurityBar] = [] + for _, row in df.iterrows(): + dt = row["date"] + bars.append( + SecurityBar( + open=row["open"], + close=row["close"], + high=row["high"], + low=row["low"], + vol=row["vol"], + amount=row["amount"], + year=dt.year, + month=dt.month, + day=dt.day, + hour=0, + minute=0, + ) + ) + bars.sort(key=lambda b: b.year * 10000 + b.month * 100 + b.day) + return bars + + +def _fetch_all_daily_bars( + client: TdxClient, market: int, code: str, need_full: bool = False +) -> list[SecurityBar]: + """从服务端分页获取全部日线数据。 + + Args: + client: 已连接的 TdxClient。 + market: 市场代码（0=SZ, 1=SH）。 + code: 6 位股票代码。 + need_full: True 表示拉取全量历史（空文件场景）， + False 表示只拉最近一页（增量更新）。 + + Returns: + SecurityBar 列表（按日期升序）。 + """ + from ..models.enums import KlineCategory + + all_bars: list[SecurityBar] = [] + start = 0 + page_size = 800 + max_pages = 50 if need_full else 1 # 50 页 = 40000 条，足够覆盖 A 股全部历史 + + for _ in range(max_pages): + df = client.get_security_bars(market, code, KlineCategory.DAY, start, page_size) + if df.empty: + break + all_bars.extend(_df_to_bars(df)) + if len(df) < page_size: + break # 最后一页不满，已无更多数据 + start += page_size + + # 去重并按日期升序排列（跨页可能有重叠） + seen: set[tuple[int, int, int]] = set() + unique: list[SecurityBar] = [] + for b in all_bars: + key = (b.year, b.month, b.day) + if key not in seen: + seen.add(key) + unique.append(b) + unique.sort(key=lambda b: b.year * 10000 + b.month * 100 + b.day) + return unique + + +def _sync_one_daily(client: TdxClient, filepath: Path) -> tuple[int, str]: + """同步单只股票日线，返回 (写入条数, 状态消息)。""" + from ..offline import append_daily_bars, get_last_bar_date + from ..offline.daily_bar import _SECURITY_COEFFICIENTS, _detect_security_type + + # 文件名 → 市场 + 代码 + name = filepath.name.lower() # e.g. sh600000.day + exchange = name[:2] # "sh" or "sz" + code = name[2:8] + market = 1 if exchange == "sh" else 0 # Market.SH=1, Market.SZ=0 + + # 判断是否需要全量拉取（空文件 → 全量，有数据 → 增量） + last_date = get_last_bar_date(filepath) + need_full = last_date is None + + # 从服务端分页获取日线 + bars = _fetch_all_daily_bars(client, market, code, need_full=need_full) + if not bars: + return 0, "服务端无数据" + + # 检测证券类型获取系数 + sec_type = _detect_security_type(filepath.name) + price_coeff, vol_coeff = _SECURITY_COEFFICIENTS.get(sec_type, (0.01, 0.01)) + + # 追加写入 + written = append_daily_bars(filepath, bars, price_coeff, vol_coeff) + if written > 0: + return written, f"+{written}" + return 0, "已是最新" + + +@offline.command("sync-daily") +@click.argument("market") +@click.argument("code") +@click.option("--vipdoc", default=None, help="vipdoc 目录路径（默认自动检测）") +def sync_daily(market: str, code: str, vipdoc: str | None) -> None: + """从服务端同步日线数据到本地 .day 文件。 + + 自动检测本地文件末尾日期，从服务端分页获取缺失的数据并追加写入。 + 空文件自动全量下载，已有数据只做增量更新。 + 建议在通达信关闭时执行，避免文件被锁定。 + + MARKET: 市场代码（SZ/SH） + CODE: 6 位股票代码 + + 示例： + + easy-tdx offline sync-daily SZ 000001 + + easy-tdx offline sync-daily SH 000001 --vipdoc C:\\new_jyplug\\vipdoc + """ + from ..client import TdxClient + from ..offline import find_daily_bar_file + from .output import print_error + from .parsers import parse_market + + mkt = parse_market(market) + + try: + filepath = find_daily_bar_file(mkt, code, vipdoc) + click.echo(f"目标文件: {filepath}") + + click.echo("正在连接服务端获取日线数据...") + with TdxClient.from_best_host() as client: + written, msg = _sync_one_daily(client, filepath) + + if written > 0: + click.echo(f"✓ 成功写入 {written} 条新记录") + else: + click.echo(f"本地已是最新，无需写入 ({msg})") + + except PermissionError: + click.echo(f"✗ 文件被锁定，请关闭通达信后重试: {filepath}", err=True) + raise SystemExit(1) + except Exception as e: + print_error(str(e)) + raise SystemExit(1) + + +# --------------------------------------------------------------------------- +# sync-all：一键同步全部日线 +# --------------------------------------------------------------------------- + + +@offline.command("sync-all") +@click.option("--vipdoc", default=None, help="vipdoc 目录路径（默认自动检测）") +def sync_all(vipdoc: str | None) -> None: + """一键同步全部本地日线数据（沪深全市场）。 + + 扫描 vipdoc 下所有 .day 文件，自动连接服务端获取最新数据并追加写入。 + 建议在通达信关闭时执行，避免文件被锁定。 + + 示例： + + easy-tdx offline sync-all + + easy-tdx offline sync-all --vipdoc C:\\new_jyplug\\vipdoc + """ + import time + + from ..client import TdxClient + from ..offline.paths import resolve_vipdoc + + try: + vipdoc_path = resolve_vipdoc(vipdoc) + except Exception as e: + click.echo(f"✗ {e}", err=True) + raise SystemExit(1) + + # 1. 扫描所有 .day 文件 + all_files: list[Path] = [] + for exchange in ("sh", "sz"): + lday_dir = vipdoc_path / exchange / "lday" + if lday_dir.is_dir(): + all_files.extend(sorted(lday_dir.glob("*.day"))) + + if not all_files: + click.echo("未找到任何 .day 文件，请确认 vipdoc 路径正确") + raise SystemExit(0) + + total = len(all_files) + click.echo(f"发现 {total} 个 .day 文件，开始同步...") + + # 2. 连接服务端，逐个同步 + success = 0 + skipped = 0 + failed = 0 + total_written = 0 + + with TdxClient.from_best_host() as client: + for idx, filepath in enumerate(all_files, 1): + name = filepath.name + try: + written, msg = _sync_one_daily(client, filepath) + total_written += written + if written > 0: + success += 1 + else: + skipped += 1 + click.echo(f" [{idx}/{total}] {name}: {msg}") + except PermissionError: + failed += 1 + click.echo(f" [{idx}/{total}] {name}: ✗ 文件被锁定", err=True) + except Exception as e: + failed += 1 + click.echo(f" [{idx}/{total}] {name}: ✗ {e}", err=True) + + # 每 100 只暂停一小段，避免请求过快被服务器断开 + if idx % 100 == 0: + time.sleep(0.2) + + # 3. 汇总 + click.echo("") + summary = ( + f"同步完成: {total} 只 | " + f"更新 {success} | 已是最新 {skipped} | " + f"失败 {failed} | 共写入 {total_written} 条" + ) + click.echo(summary) diff --git a/src/easy_tdx/offline/__init__.py b/src/easy_tdx/offline/__init__.py index 6600ef0..fa5bf7e 100644 --- a/src/easy_tdx/offline/__init__.py +++ b/src/easy_tdx/offline/__init__.py @@ -1,4 +1,4 @@ -"""离线数据读取模块 —— 从本地通达信安装目录读取数据文件。""" +"""离线数据读写模块 —— 从本地通达信安装目录读取/写入数据文件。""" from .block import CustomerBlock, read_block_dat, read_customer_blocks from .daily_bar import find_daily_bar_file, read_daily_bars @@ -8,21 +8,58 @@ from .gbbq import GbbqRecord, read_gbbq from .history_financial import read_history_financial from .min_bar import read_5min_bars, read_lc_min_bars from .paths import detect_tdx_home, resolve_vipdoc +from .write_daily import ( + append_daily_bars, + encode_daily_bar, + get_last_bar_date, + sync_daily_bars_from_security_bars, +) +from .write_ex_daily import ( + append_ex_daily_bars, + encode_ex_daily_bar, + get_last_ex_bar_date, + sync_ex_daily_bars, +) +from .write_min_bar import ( + append_5min_bars, + append_lc_min_bars, + encode_5min_bar, + encode_lc_min_bar, + get_last_5min_bar_datetime, + get_last_lc_min_bar_datetime, +) __all__ = [ # 路径 "detect_tdx_home", "resolve_vipdoc", - # 日线 + # 日线读取 "read_daily_bars", "find_daily_bar_file", - # 分钟线 + # 日线写入 + "encode_daily_bar", + "append_daily_bars", + "get_last_bar_date", + "sync_daily_bars_from_security_bars", + # 扩展市场日线写入 + "encode_ex_daily_bar", + "append_ex_daily_bars", + "get_last_ex_bar_date", + "sync_ex_daily_bars", + # 分钟线写入 + "encode_5min_bar", + "encode_lc_min_bar", + "append_5min_bars", + "append_lc_min_bars", + "get_last_5min_bar_datetime", + "get_last_lc_min_bar_datetime", + # 分钟线读取 "read_5min_bars", "read_lc_min_bars", "find_5min_bar_file", "find_lc1_bar_file", "find_lc5_bar_file", - # 扩展市场 + # 扩展市场读取 "ExDailyBar", "read_ex_daily_bars", # 板块 diff --git a/src/easy_tdx/offline/write_daily.py b/src/easy_tdx/offline/write_daily.py new file mode 100644 index 0000000..1756957 --- /dev/null +++ b/src/easy_tdx/offline/write_daily.py @@ -0,0 +1,145 @@ +"""离线日线数据写入 —— 将 SecurityBar 编码并追加到 .day 文件。""" + +from __future__ import annotations + +from pathlib import Path + +from ..models.bar import SecurityBar +from .daily_bar import _DAILY_FMT + +__all__ = [ + "encode_daily_bar", + "append_daily_bars", + "get_last_bar_date", + "sync_daily_bars_from_security_bars", +] + + +# --------------------------------------------------------------------------- +# encode +# --------------------------------------------------------------------------- + + +def encode_daily_bar( + bar: SecurityBar, + price_coeff: float, + vol_coeff: float, +) -> bytes: + """将 SecurityBar 编码为 32 字节 .day 记录。 + + Args: + bar: K 线数据（open/close/high/low 为实际价格，非整数）。 + price_coeff: 价格系数（A 股 0.01，基金 0.001 等）。 + vol_coeff: 成交量系数（A 股 0.01，指数 1.0 等）。 + + Returns: + 32 字节的二进制记录。 + """ + date_int = bar.year * 10000 + bar.month * 100 + bar.day + return _DAILY_FMT.pack( + date_int, + int(round(bar.open / price_coeff)), + int(round(bar.high / price_coeff)), + int(round(bar.low / price_coeff)), + int(round(bar.close / price_coeff)), + bar.amount, # float32, 由 struct 自动截断 + int(round(bar.vol / vol_coeff)), + 0, # reserved + ) + + +# --------------------------------------------------------------------------- +# query +# --------------------------------------------------------------------------- + + +def get_last_bar_date(filepath: str | Path) -> int | None: + """读取 .day 文件最后一条记录的日期。 + + Returns: + YYYYMMDD 整数，文件为空或太短时返回 None。 + """ + filepath = Path(filepath) + if not filepath.is_file(): + return None + size = filepath.stat().st_size + if size < _DAILY_FMT.size: + return None + with filepath.open("rb") as f: + f.seek(size - _DAILY_FMT.size) + last_record = f.read(_DAILY_FMT.size) + (date_int, *_) = _DAILY_FMT.unpack(last_record) + return date_int + + +def _bar_date_int(bar: SecurityBar) -> int: + return bar.year * 10000 + bar.month * 100 + bar.day + + +# --------------------------------------------------------------------------- +# append +# --------------------------------------------------------------------------- + + +def append_daily_bars( + filepath: str | Path, + bars: list[SecurityBar], + price_coeff: float, + vol_coeff: float, +) -> int: + """将 bars 追加写入 .day 文件，自动跳过重复日期。 + + Args: + filepath: .day 文件路径。 + bars: 待写入的 K 线列表（按时间升序）。 + price_coeff: 价格系数。 + vol_coeff: 成交量系数。 + + Returns: + 实际写入的记录数。 + """ + filepath = Path(filepath) + + # 获取文件末尾日期，用于去重 + last_date = get_last_bar_date(filepath) + + # 过滤出日期严格大于末尾的新记录 + new_bars = ( + [b for b in bars if _bar_date_int(b) > last_date] if last_date is not None else list(bars) + ) + + if not new_bars: + return 0 + + encoded = b"".join(encode_daily_bar(b, price_coeff, vol_coeff) for b in new_bars) + with filepath.open("ab") as f: + f.write(encoded) + + return len(new_bars) + + +# --------------------------------------------------------------------------- +# sync +# --------------------------------------------------------------------------- + + +def sync_daily_bars_from_security_bars( + filepath: str | Path, + server_bars: list[SecurityBar], + price_coeff: float, + vol_coeff: float, +) -> int: + """将服务端获取的日线数据同步写入本地 .day 文件。 + + 完整流程：读取文件末尾日期 → 过滤新数据 → 追加写入。 + + Args: + filepath: .day 文件路径。 + server_bars: 服务端返回的日线数据（按时间升序）。 + price_coeff: 价格系数。 + vol_coeff: 成交量系数。 + + Returns: + 实际写入的记录数。 + """ + return append_daily_bars(filepath, server_bars, price_coeff, vol_coeff) diff --git a/src/easy_tdx/offline/write_ex_daily.py b/src/easy_tdx/offline/write_ex_daily.py new file mode 100644 index 0000000..6d067d9 --- /dev/null +++ b/src/easy_tdx/offline/write_ex_daily.py @@ -0,0 +1,93 @@ +"""离线扩展市场日线写入 —— 将 ExDailyBar 编码并追加到 .day 文件。""" + +from __future__ import annotations + +from pathlib import Path + +from .ex_daily_bar import _EX_DAILY_FMT, ExDailyBar + +__all__ = [ + "encode_ex_daily_bar", + "append_ex_daily_bars", + "get_last_ex_bar_date", + "sync_ex_daily_bars", +] + + +def encode_ex_daily_bar(bar: ExDailyBar) -> bytes: + """将 ExDailyBar 编码为 32 字节扩展市场 .day 记录。 + + 扩展市场价格直接为 float32，无需系数转换。 + """ + date_int = bar.year * 10000 + bar.month * 100 + bar.day + return _EX_DAILY_FMT.pack( + date_int, + bar.open, + bar.high, + bar.low, + bar.close, + bar.amount, + bar.vol, + bar.settlement, + ) + + +def get_last_ex_bar_date(filepath: str | Path) -> int | None: + """读取扩展市场 .day 文件最后一条记录的日期。 + + Returns: + YYYYMMDD 整数，文件为空或太短时返回 None。 + """ + filepath = Path(filepath) + if not filepath.is_file(): + return None + size = filepath.stat().st_size + if size < _EX_DAILY_FMT.size: + return None + with filepath.open("rb") as f: + f.seek(size - _EX_DAILY_FMT.size) + last_record = f.read(_EX_DAILY_FMT.size) + (date_int, *_) = _EX_DAILY_FMT.unpack(last_record) + return date_int + + +def _bar_date_int(bar: ExDailyBar) -> int: + return bar.year * 10000 + bar.month * 100 + bar.day + + +def append_ex_daily_bars( + filepath: str | Path, + bars: list[ExDailyBar], +) -> int: + """将扩展市场 bars 追加写入 .day 文件，自动跳过重复日期。 + + Returns: + 实际写入的记录数。 + """ + filepath = Path(filepath) + last_date = get_last_ex_bar_date(filepath) + + new_bars = ( + [b for b in bars if _bar_date_int(b) > last_date] if last_date is not None else list(bars) + ) + + if not new_bars: + return 0 + + encoded = b"".join(encode_ex_daily_bar(b) for b in new_bars) + with filepath.open("ab") as f: + f.write(encoded) + + return len(new_bars) + + +def sync_ex_daily_bars( + filepath: str | Path, + server_bars: list[ExDailyBar], +) -> int: + """将服务端获取的扩展市场日线同步写入本地 .day 文件。 + + Returns: + 实际写入的记录数。 + """ + return append_ex_daily_bars(filepath, server_bars) diff --git a/src/easy_tdx/offline/write_min_bar.py b/src/easy_tdx/offline/write_min_bar.py new file mode 100644 index 0000000..df82ca2 --- /dev/null +++ b/src/easy_tdx/offline/write_min_bar.py @@ -0,0 +1,186 @@ +"""离线分钟线写入 —— 将 SecurityBar 编码并追加到 .5 / .lc1 / .lc5 文件。""" + +from __future__ import annotations + +from pathlib import Path + +from ..models.bar import SecurityBar +from .min_bar import _LC_MIN_FMT, _MIN_FMT + +__all__ = [ + "encode_5min_bar", + "encode_lc_min_bar", + "append_5min_bars", + "append_lc_min_bars", + "get_last_5min_bar_datetime", + "get_last_lc_min_bar_datetime", +] + + +# --------------------------------------------------------------------------- +# date/time 编码（与 min_bar._decode_tdx_date/time 互逆） +# --------------------------------------------------------------------------- + + +def _encode_tdx_date(year: int, month: int, day: int) -> int: + """将 (year, month, day) 编码为 2 字节压缩日期。""" + return (year - 2004) * 2048 + month * 100 + day + + +def _encode_tdx_time(hour: int, minute: int) -> int: + """将 (hour, minute) 编码为从 0:00 起的分钟数。""" + return hour * 60 + minute + + +def _bar_datetime_key(bar: SecurityBar) -> tuple[int, int, int, int, int]: + return (bar.year, bar.month, bar.day, bar.hour, bar.minute) + + +# --------------------------------------------------------------------------- +# .5 文件 (OHLC 为整数 / 100) +# --------------------------------------------------------------------------- + + +def encode_5min_bar(bar: SecurityBar) -> bytes: + """将 SecurityBar 编码为 32 字节 .5 记录。 + + OHLC 为整数（实际价格 × 100），amount 为 float32。 + """ + return _MIN_FMT.pack( + _encode_tdx_date(bar.year, bar.month, bar.day), + _encode_tdx_time(bar.hour, bar.minute), + int(round(bar.open * 100)), + int(round(bar.high * 100)), + int(round(bar.low * 100)), + int(round(bar.close * 100)), + bar.amount, # float32 + int(round(bar.vol)), + 0, # reserved + ) + + +def get_last_5min_bar_datetime( + filepath: str | Path, +) -> tuple[int, int, int, int, int] | None: + """读取 .5 文件最后一条记录的日期时间。 + + Returns: + (year, month, day, hour, minute) 元组，文件为空时返回 None。 + """ + filepath = Path(filepath) + if not filepath.is_file(): + return None + size = filepath.stat().st_size + if size < _MIN_FMT.size: + return None + with filepath.open("rb") as f: + f.seek(size - _MIN_FMT.size) + last_record = f.read(_MIN_FMT.size) + from .min_bar import _decode_tdx_date, _decode_tdx_time + + date_num, time_num, *_ = _MIN_FMT.unpack(last_record) + year, month, day = _decode_tdx_date(date_num) + hour, minute = _decode_tdx_time(time_num) + return (year, month, day, hour, minute) + + +def append_5min_bars( + filepath: str | Path, + bars: list[SecurityBar], +) -> int: + """将 5 分钟线 bars 追加写入 .5 文件，自动跳过重复时间点。 + + Returns: + 实际写入的记录数。 + """ + filepath = Path(filepath) + last_dt = get_last_5min_bar_datetime(filepath) + + if last_dt is not None: + new_bars = [b for b in bars if _bar_datetime_key(b) > last_dt] + else: + new_bars = list(bars) + + if not new_bars: + return 0 + + encoded = b"".join(encode_5min_bar(b) for b in new_bars) + with filepath.open("ab") as f: + f.write(encoded) + + return len(new_bars) + + +# --------------------------------------------------------------------------- +# .lc1 / .lc5 文件 (OHLC 为 float32) +# --------------------------------------------------------------------------- + + +def encode_lc_min_bar(bar: SecurityBar) -> bytes: + """将 SecurityBar 编码为 32 字节 .lc1/.lc5 记录。 + + OHLC 为 float32，无需系数转换。 + """ + return _LC_MIN_FMT.pack( + _encode_tdx_date(bar.year, bar.month, bar.day), + _encode_tdx_time(bar.hour, bar.minute), + bar.open, + bar.high, + bar.low, + bar.close, + bar.amount, # float32 + int(round(bar.vol)), + 0, # reserved + ) + + +def get_last_lc_min_bar_datetime( + filepath: str | Path, +) -> tuple[int, int, int, int, int] | None: + """读取 .lc1/.lc5 文件最后一条记录的日期时间。 + + Returns: + (year, month, day, hour, minute) 元组，文件为空时返回 None。 + """ + filepath = Path(filepath) + if not filepath.is_file(): + return None + size = filepath.stat().st_size + if size < _LC_MIN_FMT.size: + return None + with filepath.open("rb") as f: + f.seek(size - _LC_MIN_FMT.size) + last_record = f.read(_LC_MIN_FMT.size) + from .min_bar import _decode_tdx_date, _decode_tdx_time + + date_num, time_num, *_ = _LC_MIN_FMT.unpack(last_record) + year, month, day = _decode_tdx_date(date_num) + hour, minute = _decode_tdx_time(time_num) + return (year, month, day, hour, minute) + + +def append_lc_min_bars( + filepath: str | Path, + bars: list[SecurityBar], +) -> int: + """将分钟线 bars 追加写入 .lc1/.lc5 文件，自动跳过重复时间点。 + + Returns: + 实际写入的记录数。 + """ + filepath = Path(filepath) + last_dt = get_last_lc_min_bar_datetime(filepath) + + if last_dt is not None: + new_bars = [b for b in bars if _bar_datetime_key(b) > last_dt] + else: + new_bars = list(bars) + + if not new_bars: + return 0 + + encoded = b"".join(encode_lc_min_bar(b) for b in new_bars) + with filepath.open("ab") as f: + f.write(encoded) + + return len(new_bars) diff --git a/tests/unit/test_write_daily.py b/tests/unit/test_write_daily.py new file mode 100644 index 0000000..5980585 --- /dev/null +++ b/tests/unit/test_write_daily.py @@ -0,0 +1,345 @@ +"""离线日线写入测试（纯离线，无网络）。""" + +from __future__ import annotations + +import struct +from pathlib import Path + +from easy_tdx.models.bar import SecurityBar +from easy_tdx.offline.daily_bar import ( + _DAILY_FMT, + read_daily_bars, +) +from easy_tdx.offline.write_daily import ( + append_daily_bars, + encode_daily_bar, + get_last_bar_date, + sync_daily_bars_from_security_bars, +) + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _make_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + open_: float = 10.25, + high: float = 10.50, + low: float = 10.10, + close: float = 10.30, + vol: float = 100000.0, + amount: float = 1025000.0, +) -> SecurityBar: + return SecurityBar( + open=open_, + close=close, + high=high, + low=low, + vol=vol, + amount=amount, + year=year, + month=month, + day=day, + hour=0, + minute=0, + ) + + +def _make_raw_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + open_int: int = 1025, + high_int: int = 1050, + low_int: int = 1010, + close_int: int = 1030, + amount: float = 1025000.0, + vol_int: int = 100000, + reserved: int = 0, +) -> bytes: + """直接用原始整数构造一条 32 字节 .day 记录。""" + date_int = year * 10000 + month * 100 + day + return _DAILY_FMT.pack( + date_int, open_int, high_int, low_int, close_int, amount, vol_int, reserved + ) + + +# --------------------------------------------------------------------------- +# encode_daily_bar +# --------------------------------------------------------------------------- + + +class TestEncodeDailyBar: + """编码测试：SecurityBar → 32 字节。""" + + def test_output_length(self) -> None: + bar = _make_bar() + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + assert len(result) == 32 + + def test_date_encoding(self) -> None: + bar = _make_bar(year=2025, month=3, day=15) + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + date_int = struct.unpack_from(" None: + """A股：系数 0.01，即 float × 100 → 整数。""" + bar = _make_bar(open_=10.25) + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + open_int = struct.unpack_from(" None: + """指数：系数 0.01。""" + bar = _make_bar(open_=3250.18) + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=1.0) + open_int = struct.unpack_from(" None: + """基金：系数 0.001，即 float × 1000 → 整数。""" + bar = _make_bar(open_=1.523) + result = encode_daily_bar(bar, price_coeff=0.001, vol_coeff=1.0) + open_int = struct.unpack_from(" None: + """A股量系数 0.01。""" + bar = _make_bar(vol=12345.67) + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + vol_int = struct.unpack_from(" None: + """指数量系数 1.0。""" + bar = _make_bar(vol=123456789.0) + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=1.0) + vol_int = struct.unpack_from(" None: + bar = _make_bar(amount=1_025_000.0) + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + (amt_out,) = struct.unpack_from(" None: + bar = _make_bar() + result = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + reserved = struct.unpack_from(" None: + bar = _make_bar( + open_=10.25, high=10.50, low=10.10, close=10.30, vol=100000.0, amount=1025000.0 + ) + encoded = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=0.01) + + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(encoded) + + bars = read_daily_bars(filepath) + assert len(bars) == 1 + + b = bars[0] + assert b.year == 2026 and b.month == 6 and b.day == 6 + assert abs(b.open - 10.25) < 0.01 + assert abs(b.high - 10.50) < 0.01 + assert abs(b.low - 10.10) < 0.01 + assert abs(b.close - 10.30) < 0.01 + assert abs(b.vol - 100000.0) < 1.0 + assert abs(b.amount - 1025000.0) < 1.0 + + def test_multiple_bars_round_trip(self, tmp_path: Path) -> None: + bars_in = [ + _make_bar(year=2026, month=6, day=4, open_=10.0, close=10.1), + _make_bar(year=2026, month=6, day=5, open_=10.1, close=10.2), + _make_bar(year=2026, month=6, day=6, open_=10.2, close=10.3), + ] + encoded = b"".join(encode_daily_bar(b, price_coeff=0.01, vol_coeff=0.01) for b in bars_in) + + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(encoded) + + bars_out = read_daily_bars(filepath) + assert len(bars_out) == 3 + assert bars_out[0].day == 4 + assert bars_out[2].day == 6 + + def test_index_round_trip(self, tmp_path: Path) -> None: + """指数：价格系数 0.01，量系数 1.0。""" + bar = _make_bar(open_=3250.18, close=3260.5, vol=123456789.0) + encoded = encode_daily_bar(bar, price_coeff=0.01, vol_coeff=1.0) + + filepath = tmp_path / "sh000001.day" + filepath.write_bytes(encoded) + + bars = read_daily_bars(filepath) + assert len(bars) == 1 + assert abs(bars[0].open - 3250.18) < 0.01 + assert abs(bars[0].vol - 123456789.0) < 1.0 + + +# --------------------------------------------------------------------------- +# get_last_bar_date +# --------------------------------------------------------------------------- + + +class TestGetLastBarDate: + def test_returns_last_date(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(_make_raw_bar(year=2026, month=6, day=5)) + assert get_last_bar_date(filepath) == 20260605 + + def test_returns_none_for_empty_file(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(b"") + assert get_last_bar_date(filepath) is None + + def test_returns_none_for_short_file(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(b"\x00" * 16) # < 32 bytes + assert get_last_bar_date(filepath) is None + + def test_returns_last_of_multiple(self, tmp_path: Path) -> None: + data = ( + _make_raw_bar(year=2026, month=6, day=3) + + _make_raw_bar(year=2026, month=6, day=4) + + _make_raw_bar(year=2026, month=6, day=5) + ) + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(data) + assert get_last_bar_date(filepath) == 20260605 + + +# --------------------------------------------------------------------------- +# append_daily_bars +# --------------------------------------------------------------------------- + + +class TestAppendDailyBars: + def test_append_to_existing(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(_make_raw_bar(year=2026, month=6, day=5)) + + new_bar = _make_bar(year=2026, month=6, day=6, open_=10.3, close=10.4) + append_daily_bars(filepath, [new_bar], price_coeff=0.01, vol_coeff=0.01) + + bars = read_daily_bars(filepath) + assert len(bars) == 2 + assert bars[0].day == 5 + assert bars[1].day == 6 + + def test_append_to_empty_file(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(b"") + + new_bar = _make_bar(year=2026, month=6, day=6) + append_daily_bars(filepath, [new_bar], price_coeff=0.01, vol_coeff=0.01) + + bars = read_daily_bars(filepath) + assert len(bars) == 1 + assert bars[0].day == 6 + + def test_append_skips_duplicate_date(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(_make_raw_bar(year=2026, month=6, day=6)) + + new_bar = _make_bar(year=2026, month=6, day=6) # same date + written = append_daily_bars(filepath, [new_bar], price_coeff=0.01, vol_coeff=0.01) + assert written == 0 # skipped + + bars = read_daily_bars(filepath) + assert len(bars) == 1 # no duplicate + + def test_append_multiple_filters_duplicates(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(_make_raw_bar(year=2026, month=6, day=5)) + + bars_to_append = [ + _make_bar(year=2026, month=6, day=5), # dup + _make_bar(year=2026, month=6, day=6), # new + _make_bar(year=2026, month=6, day=7), # new + ] + written = append_daily_bars(filepath, bars_to_append, price_coeff=0.01, vol_coeff=0.01) + assert written == 2 + + bars = read_daily_bars(filepath) + assert len(bars) == 3 + + +# --------------------------------------------------------------------------- +# sync_daily_bars_from_security_bars +# --------------------------------------------------------------------------- + + +class TestSyncDailyBars: + """模拟完整同步流程（不需要真实服务端，手动构造 SecurityBar 列表）。""" + + def test_sync_appends_new_only(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + # 文件已有 6月4日和6月5日 + filepath.write_bytes( + _make_raw_bar(year=2026, month=6, day=4) + _make_raw_bar(year=2026, month=6, day=5) + ) + + # 模拟服务端返回的数据：6月3日~6月7日 + server_bars = [ + _make_bar(year=2026, month=6, day=3), + _make_bar(year=2026, month=6, day=4), + _make_bar(year=2026, month=6, day=5), + _make_bar(year=2026, month=6, day=6), + _make_bar(year=2026, month=6, day=7), + ] + + written = sync_daily_bars_from_security_bars( + filepath, server_bars, price_coeff=0.01, vol_coeff=0.01 + ) + assert written == 2 # only 6/6 and 6/7 + + bars = read_daily_bars(filepath) + assert len(bars) == 4 # 6/4, 6/5 (original) + 6/6, 6/7 (new) + + def test_sync_empty_file(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(b"") + + server_bars = [ + _make_bar(year=2026, month=6, day=4), + _make_bar(year=2026, month=6, day=5), + ] + + written = sync_daily_bars_from_security_bars( + filepath, server_bars, price_coeff=0.01, vol_coeff=0.01 + ) + assert written == 2 + + bars = read_daily_bars(filepath) + assert len(bars) == 2 + + def test_sync_nothing_new(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.day" + filepath.write_bytes(_make_raw_bar(year=2026, month=6, day=7)) + + server_bars = [ + _make_bar(year=2026, month=6, day=6), + _make_bar(year=2026, month=6, day=7), + ] + + written = sync_daily_bars_from_security_bars( + filepath, server_bars, price_coeff=0.01, vol_coeff=0.01 + ) + assert written == 0 diff --git a/tests/unit/test_write_ex_daily.py b/tests/unit/test_write_ex_daily.py new file mode 100644 index 0000000..dcdfce5 --- /dev/null +++ b/tests/unit/test_write_ex_daily.py @@ -0,0 +1,190 @@ +"""离线扩展市场日线写入测试（纯离线，无网络）。""" + +from __future__ import annotations + +import struct +from pathlib import Path + +from easy_tdx.offline.ex_daily_bar import _EX_DAILY_FMT, ExDailyBar, read_ex_daily_bars +from easy_tdx.offline.write_ex_daily import ( + append_ex_daily_bars, + encode_ex_daily_bar, + get_last_ex_bar_date, + sync_ex_daily_bars, +) + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _make_ex_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + open_: float = 3500.0, + high: float = 3550.0, + low: float = 3480.0, + close: float = 3520.0, + amount: int = 123456, + vol: int = 98765, + settlement: float = 3510.0, +) -> ExDailyBar: + return ExDailyBar( + open=open_, + close=close, + high=high, + low=low, + amount=amount, + vol=vol, + settlement=settlement, + hk_stock_amount=0.0, + year=year, + month=month, + day=day, + ) + + +def _make_raw_ex_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + open_: float = 3500.0, + high: float = 3550.0, + low: float = 3480.0, + close: float = 3520.0, + amount: int = 123456, + vol: int = 98765, + settlement: float = 3510.0, +) -> bytes: + date_int = year * 10000 + month * 100 + day + return _EX_DAILY_FMT.pack(date_int, open_, high, low, close, amount, vol, settlement) + + +# --------------------------------------------------------------------------- +# encode_ex_daily_bar +# --------------------------------------------------------------------------- + + +class TestEncodeExDailyBar: + def test_output_length(self) -> None: + bar = _make_ex_bar() + result = encode_ex_daily_bar(bar) + assert len(result) == 32 + + def test_date_encoding(self) -> None: + bar = _make_ex_bar(year=2025, month=12, day=25) + result = encode_ex_daily_bar(bar) + date_int = struct.unpack_from(" None: + bar = _make_ex_bar(open_=3500.5, high=3560.25, low=3479.75, close=3520.0) + result = encode_ex_daily_bar(bar) + op, hi, lo, cl = struct.unpack_from(" None: + bar = _make_ex_bar(settlement=3510.5) + result = encode_ex_daily_bar(bar) + (sett,) = struct.unpack_from(" None: + bar = _make_ex_bar(open_=3500.0, close=3520.0, vol=98765) + encoded = encode_ex_daily_bar(bar) + + filepath = tmp_path / "29#A1801.day" + filepath.write_bytes(encoded) + + bars = read_ex_daily_bars(filepath) + assert len(bars) == 1 + b = bars[0] + assert b.year == 2026 and b.month == 6 and b.day == 6 + assert abs(b.open - 3500.0) < 0.01 + assert abs(b.close - 3520.0) < 0.01 + assert b.vol == 98765 + + def test_multiple_bars_round_trip(self, tmp_path: Path) -> None: + bars_in = [ + _make_ex_bar(year=2026, month=6, day=4, open_=3400.0), + _make_ex_bar(year=2026, month=6, day=5, open_=3450.0), + _make_ex_bar(year=2026, month=6, day=6, open_=3500.0), + ] + encoded = b"".join(encode_ex_daily_bar(b) for b in bars_in) + + filepath = tmp_path / "29#A1801.day" + filepath.write_bytes(encoded) + + bars_out = read_ex_daily_bars(filepath) + assert len(bars_out) == 3 + assert bars_out[0].day == 4 + assert bars_out[2].day == 6 + + +# --------------------------------------------------------------------------- +# get_last_ex_bar_date +# --------------------------------------------------------------------------- + + +class TestGetLastExBarDate: + def test_returns_last_date(self, tmp_path: Path) -> None: + filepath = tmp_path / "test.day" + filepath.write_bytes(_make_raw_ex_bar(year=2026, month=6, day=5)) + assert get_last_ex_bar_date(filepath) == 20260605 + + def test_returns_none_for_empty(self, tmp_path: Path) -> None: + filepath = tmp_path / "test.day" + filepath.write_bytes(b"") + assert get_last_ex_bar_date(filepath) is None + + +# --------------------------------------------------------------------------- +# append & sync +# --------------------------------------------------------------------------- + + +class TestAppendExDailyBars: + def test_append_to_existing(self, tmp_path: Path) -> None: + filepath = tmp_path / "test.day" + filepath.write_bytes(_make_raw_ex_bar(year=2026, month=6, day=5)) + + new_bar = _make_ex_bar(year=2026, month=6, day=6) + written = append_ex_daily_bars(filepath, [new_bar]) + assert written == 1 + + bars = read_ex_daily_bars(filepath) + assert len(bars) == 2 + + def test_skips_duplicate(self, tmp_path: Path) -> None: + filepath = tmp_path / "test.day" + filepath.write_bytes(_make_raw_ex_bar(year=2026, month=6, day=6)) + + new_bar = _make_ex_bar(year=2026, month=6, day=6) + written = append_ex_daily_bars(filepath, [new_bar]) + assert written == 0 + + def test_sync_filters_correctly(self, tmp_path: Path) -> None: + filepath = tmp_path / "test.day" + filepath.write_bytes(_make_raw_ex_bar(year=2026, month=6, day=5)) + + server_bars = [ + _make_ex_bar(year=2026, month=6, day=4), + _make_ex_bar(year=2026, month=6, day=5), + _make_ex_bar(year=2026, month=6, day=6), + ] + written = sync_ex_daily_bars(filepath, server_bars) + assert written == 1 + + bars = read_ex_daily_bars(filepath) + assert len(bars) == 2 diff --git a/tests/unit/test_write_min_bar.py b/tests/unit/test_write_min_bar.py new file mode 100644 index 0000000..ca43b3c --- /dev/null +++ b/tests/unit/test_write_min_bar.py @@ -0,0 +1,276 @@ +"""离线分钟线写入测试（纯离线，无网络）。""" + +from __future__ import annotations + +import struct +from pathlib import Path + +from easy_tdx.models.bar import SecurityBar +from easy_tdx.offline.min_bar import ( + _LC_MIN_FMT, + _MIN_FMT, + _decode_tdx_date, + _decode_tdx_time, + read_5min_bars, + read_lc_min_bars, +) +from easy_tdx.offline.write_min_bar import ( + append_5min_bars, + append_lc_min_bars, + encode_5min_bar, + encode_lc_min_bar, + get_last_5min_bar_datetime, + get_last_lc_min_bar_datetime, +) + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _make_min_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + hour: int = 9, + minute: int = 30, + open_: float = 10.25, + high: float = 10.50, + low: float = 10.10, + close: float = 10.30, + vol: float = 5000.0, + amount: float = 51250.0, +) -> SecurityBar: + return SecurityBar( + open=open_, + close=close, + high=high, + low=low, + vol=vol, + amount=amount, + year=year, + month=month, + day=day, + hour=hour, + minute=minute, + ) + + +def _encode_tdx_date(year: int, month: int, day: int) -> int: + return (year - 2004) * 2048 + month * 100 + day + + +def _encode_tdx_time(hour: int, minute: int) -> int: + return hour * 60 + minute + + +def _make_raw_5min_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + hour: int = 9, + minute: int = 30, + open_int: int = 1025, + high_int: int = 1050, + low_int: int = 1010, + close_int: int = 1030, + amount: float = 51250.0, + vol_int: int = 5000, +) -> bytes: + return _MIN_FMT.pack( + _encode_tdx_date(year, month, day), + _encode_tdx_time(hour, minute), + open_int, + high_int, + low_int, + close_int, + amount, + vol_int, + 0, + ) + + +def _make_raw_lc_min_bar( + year: int = 2026, + month: int = 6, + day: int = 6, + hour: int = 9, + minute: int = 30, + open_: float = 10.25, + high: float = 10.50, + low: float = 10.10, + close: float = 10.30, + amount: float = 51250.0, + vol: int = 5000, +) -> bytes: + return _LC_MIN_FMT.pack( + _encode_tdx_date(year, month, day), + _encode_tdx_time(hour, minute), + open_, + high, + low, + close, + amount, + vol, + 0, + ) + + +# =========================================================================== +# .5 文件 (5 分钟线) +# =========================================================================== + + +class TestEncode5minBar: + def test_output_length(self) -> None: + bar = _make_min_bar() + result = encode_5min_bar(bar) + assert len(result) == 32 + + def test_date_encoding(self) -> None: + bar = _make_min_bar(year=2026, month=6, day=5) + result = encode_5min_bar(bar) + date_num = struct.unpack_from(" None: + bar = _make_min_bar(hour=14, minute=30) + result = encode_5min_bar(bar) + time_num = struct.unpack_from(" None: + """5 分钟线 OHLC 整数 × 100。""" + bar = _make_min_bar(open_=10.25) + result = encode_5min_bar(bar) + open_int = struct.unpack_from(" None: + bar = _make_min_bar(open_=10.25, close=10.30, vol=5000.0, amount=51250.0) + encoded = encode_5min_bar(bar) + + filepath = tmp_path / "sh600000.5" + filepath.write_bytes(encoded) + + bars = read_5min_bars(filepath) + assert len(bars) == 1 + b = bars[0] + assert abs(b.open - 10.25) < 0.01 + assert abs(b.close - 10.30) < 0.01 + assert abs(b.vol - 5000.0) < 1.0 + assert b.hour == 9 and b.minute == 30 + + def test_multiple_bars(self, tmp_path: Path) -> None: + bars_in = [ + _make_min_bar(hour=9, minute=30), + _make_min_bar(hour=9, minute=35), + _make_min_bar(hour=9, minute=40), + ] + encoded = b"".join(encode_5min_bar(b) for b in bars_in) + + filepath = tmp_path / "sh600000.5" + filepath.write_bytes(encoded) + + bars_out = read_5min_bars(filepath) + assert len(bars_out) == 3 + assert bars_out[0].minute == 30 + assert bars_out[2].minute == 40 + + +class TestGetLast5minBarDatetime: + def test_returns_last_datetime(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.5" + filepath.write_bytes(_make_raw_5min_bar(year=2026, month=6, day=5, hour=15, minute=0)) + result = get_last_5min_bar_datetime(filepath) + assert result == (2026, 6, 5, 15, 0) + + def test_returns_none_for_empty(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.5" + filepath.write_bytes(b"") + assert get_last_5min_bar_datetime(filepath) is None + + +class TestAppend5minBars: + def test_append_to_existing(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.5" + filepath.write_bytes(_make_raw_5min_bar(hour=9, minute=30)) + + new_bar = _make_min_bar(hour=9, minute=35) + written = append_5min_bars(filepath, [new_bar]) + assert written == 1 + + bars = read_5min_bars(filepath) + assert len(bars) == 2 + + def test_skips_duplicate(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.5" + filepath.write_bytes(_make_raw_5min_bar(hour=9, minute=30)) + + new_bar = _make_min_bar(hour=9, minute=30) # same datetime + written = append_5min_bars(filepath, [new_bar]) + assert written == 0 + + +# =========================================================================== +# .lc1 / .lc5 文件 +# =========================================================================== + + +class TestEncodeLcMinBar: + def test_output_length(self) -> None: + bar = _make_min_bar() + result = encode_lc_min_bar(bar) + assert len(result) == 32 + + def test_float_prices_preserved(self) -> None: + bar = _make_min_bar(open_=10.255) + result = encode_lc_min_bar(bar) + (op,) = struct.unpack_from(" None: + bar = _make_min_bar(open_=10.255, close=10.31, vol=5000.0, amount=51250.0) + encoded = encode_lc_min_bar(bar) + + filepath = tmp_path / "sh600000.lc1" + filepath.write_bytes(encoded) + + bars = read_lc_min_bars(filepath) + assert len(bars) == 1 + b = bars[0] + assert abs(b.open - 10.255) < 0.001 + assert abs(b.close - 10.31) < 0.001 + assert b.hour == 9 and b.minute == 30 + + +class TestGetLastLcMinBarDatetime: + def test_returns_last_datetime(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.lc1" + filepath.write_bytes(_make_raw_lc_min_bar(year=2026, month=6, day=5, hour=14, minute=55)) + result = get_last_lc_min_bar_datetime(filepath) + assert result == (2026, 6, 5, 14, 55) + + def test_returns_none_for_empty(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.lc1" + filepath.write_bytes(b"") + assert get_last_lc_min_bar_datetime(filepath) is None + + +class TestAppendLcMinBars: + def test_append_to_existing(self, tmp_path: Path) -> None: + filepath = tmp_path / "sh600000.lc1" + filepath.write_bytes(_make_raw_lc_min_bar(hour=9, minute=30)) + + new_bar = _make_min_bar(hour=9, minute=31) + written = append_lc_min_bars(filepath, [new_bar]) + assert written == 1 + + bars = read_lc_min_bars(filepath) + assert len(bars) == 2