mirror of
https://ghfast.top/https://github.com/aeroxw/easy-tdx.git
synced 2026-09-12 14:34:15 +08:00
feat: add professional financial data support via calc server
Add CALC_HOSTS, financial file list/record parsing (codec/financial.py), new client methods (get_financial_file_list, get_financial_file, get_financial_records) with async counterparts, and example 09 demo. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
GitHub
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
7fd6e610cf
commit
321043f9a5
@@ -8,7 +8,10 @@
|
||||
"Bash(ruff check *)",
|
||||
"Bash(git add *)",
|
||||
"Bash(git commit *)",
|
||||
"Bash(git push *)"
|
||||
"Bash(git push *)",
|
||||
"mcp__zread__get_repo_structure",
|
||||
"Bash(pip install *)",
|
||||
"mcp__web-reader__webReader"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
"""演示：通过 get_report_file 从服务器下载文件。
|
||||
|
||||
行情服务器（KNOWN_HOSTS）当前稳定提供的文件:
|
||||
'tdxhy.cfg' - 行业映射配置（~149KB）
|
||||
'block_zs.dat' - 行业/指数板块（~330KB）
|
||||
'block_gn.dat' - 概念板块（~757KB）
|
||||
'block_fg.dat' - 风格板块（~453KB）
|
||||
|
||||
计算服务器（CALC_HOSTS）提供专业财务数据:
|
||||
'tdxfin/gpcw.txt' - 文件列表
|
||||
'tdxfin/gpcwYYYYMMDD.zip' - 历史财报
|
||||
|
||||
行情服务器已失效（返回空包）:
|
||||
'base_info.zip', 'gpcw.txt'
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from xmtdx import CALC_HOSTS, TdxClient
|
||||
|
||||
OUTPUT_DIR = Path(__file__).parent / "downloads"
|
||||
|
||||
# 行情服务器可用的文件
|
||||
AVAILABLE_FILES = [
|
||||
"tdxhy.cfg",
|
||||
"block_zs.dat",
|
||||
"block_gn.dat",
|
||||
"block_fg.dat",
|
||||
]
|
||||
|
||||
# 行情服务器已失效的文件
|
||||
PROBE_FILES = ["base_info.zip", "gpcw.txt"]
|
||||
|
||||
# ── 1. 行情服务器文件下载 ─────────────────────────────
|
||||
|
||||
with TdxClient.from_best_host() as c:
|
||||
OUTPUT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
print("=" * 50)
|
||||
print("探测已失效文件（预期返回空包）")
|
||||
print("=" * 50)
|
||||
for filename in PROBE_FILES:
|
||||
data = c.get_report_file(filename)
|
||||
status = "空包" if len(data) == 0 else f"{len(data):,} 字节"
|
||||
print(f" {filename}: {status}")
|
||||
|
||||
print()
|
||||
print("=" * 50)
|
||||
print("下载可用文件")
|
||||
print("=" * 50)
|
||||
for filename in AVAILABLE_FILES:
|
||||
data = c.get_report_file(filename)
|
||||
out_path = OUTPUT_DIR / filename
|
||||
out_path.write_bytes(data)
|
||||
print(f" {filename} ({len(data):,} 字节) 已保存")
|
||||
|
||||
print()
|
||||
print("=" * 50)
|
||||
print("行业板块 (block_zs.dat)")
|
||||
print("=" * 50)
|
||||
blocks = c.get_block_info("block_zs.dat")
|
||||
for b in blocks[:5]:
|
||||
print(f" {b.name:<10} 分类={b.category} 成分={b.count}")
|
||||
print(f" ... 共 {len(blocks)} 个")
|
||||
|
||||
# ── 2. 计算服务器：专业财务数据 ────────────────────────
|
||||
|
||||
print()
|
||||
print("=" * 50)
|
||||
print("专业财务数据（计算服务器）")
|
||||
print("=" * 50)
|
||||
|
||||
calc_host = CALC_HOSTS[0]
|
||||
with TdxClient(calc_host) as c:
|
||||
# 获取文件列表
|
||||
file_list = c.get_financial_file_list()
|
||||
for fi in file_list[:5]:
|
||||
print(f" {fi.filename} {fi.filesize:>12,} 字节 hash={fi.hash[:8]}...")
|
||||
print(f" ... 共 {len(file_list)} 个文件")
|
||||
|
||||
# 下载并解析最近一期有实际数据的财报
|
||||
real_files = [f for f in file_list if f.filesize > 10000]
|
||||
if real_files:
|
||||
latest = real_files[0]
|
||||
fname = f"tdxfin/{latest.filename}"
|
||||
print(f"\n下载: {fname} ({latest.filesize:,} 字节)")
|
||||
|
||||
# 保存原始 .zip
|
||||
zip_data = c.get_financial_file(fname)
|
||||
zip_path = OUTPUT_DIR / latest.filename
|
||||
zip_path.write_bytes(zip_data)
|
||||
print(f" .zip 已保存到 {zip_path}")
|
||||
|
||||
# 解析财报记录
|
||||
records = c.get_financial_records(fname)
|
||||
print(f" 解析出 {len(records)} 只股票")
|
||||
if records:
|
||||
for r in records[:5]:
|
||||
print(f" {r.market.name} {r.code} 报告期={r.report_date} 字段数={len(r.fields)}")
|
||||
print(f" ... 共 {len(records)} 只")
|
||||
r = records[0]
|
||||
print(f" 示例: {r.market.name} {r.code}, 报告期={r.report_date}, 字段数={len(r.fields)}")
|
||||
@@ -21,11 +21,15 @@ asyncio 版本::
|
||||
"""
|
||||
|
||||
from .client import AsyncTdxClient, TdxClient
|
||||
from .ex.client import AsyncExTdxClient, ExTdxClient
|
||||
from .ex.models import KNOWN_EX_HOSTS
|
||||
from .exceptions import TdxCommandError, TdxConnectionError, TdxDecodeError, TdxError
|
||||
from .models import (
|
||||
XDXR_CATEGORY_NAMES,
|
||||
CompanyInfoCategory,
|
||||
FinanceInfo,
|
||||
FinancialFileInfo,
|
||||
FinancialRecord,
|
||||
KlineCategory,
|
||||
Market,
|
||||
MinuteBar,
|
||||
@@ -35,9 +39,7 @@ from .models import (
|
||||
TransactionRecord,
|
||||
XdxrRecord,
|
||||
)
|
||||
from .ex.client import AsyncExTdxClient, ExTdxClient
|
||||
from .ex.models import KNOWN_EX_HOSTS
|
||||
from .transport.sync import KNOWN_HOSTS, ping_all
|
||||
from .transport.sync import CALC_HOSTS, KNOWN_HOSTS, ping_all
|
||||
|
||||
__all__ = [
|
||||
# 客户端
|
||||
@@ -56,6 +58,8 @@ __all__ = [
|
||||
"XDXR_CATEGORY_NAMES",
|
||||
"FinanceInfo",
|
||||
"CompanyInfoCategory",
|
||||
"FinancialFileInfo",
|
||||
"FinancialRecord",
|
||||
# 异常
|
||||
"TdxError",
|
||||
"TdxConnectionError",
|
||||
@@ -68,6 +72,7 @@ __all__ = [
|
||||
# 工具
|
||||
"ping_all",
|
||||
"KNOWN_HOSTS",
|
||||
"CALC_HOSTS",
|
||||
]
|
||||
|
||||
__version__ = "0.1.1"
|
||||
|
||||
+159
-4
@@ -1,17 +1,18 @@
|
||||
"""高层行情 API：TdxClient（同步）和 AsyncTdxClient（asyncio）。"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from collections.abc import Awaitable, Callable
|
||||
from types import TracebackType
|
||||
from typing import TypeVar
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from .codec.block import parse_block_dat
|
||||
from .codec.financial import parse_financial_dat, parse_financial_file_list
|
||||
from .codec.industry import parse_tdxhy_cfg
|
||||
from .codec.price_rules import compute_price_limits, get_no_limit_window_days
|
||||
from .commands.base import BaseCommand
|
||||
@@ -30,13 +31,20 @@ from .commands.xdxr_info import GetXdxrInfoCmd
|
||||
from .exceptions import TdxConnectionError
|
||||
from .models.bar import SecurityBar
|
||||
from .models.enums import KlineCategory, Market
|
||||
from .models.finance import CompanyInfoCategory, FinanceInfo, TdxBlock, XdxrRecord
|
||||
from .models.finance import (
|
||||
CompanyInfoCategory,
|
||||
FinanceInfo,
|
||||
FinancialFileInfo,
|
||||
FinancialRecord,
|
||||
TdxBlock,
|
||||
XdxrRecord,
|
||||
)
|
||||
from .models.quote import SecurityQuote
|
||||
from .models.security import SecurityInfo
|
||||
from .models.stats import FundFlow, HistoricalFundFlow, MarketStat
|
||||
from .models.timeseries import MinuteBar, TransactionRecord
|
||||
from .transport.async_ import AsyncTdxConnection
|
||||
from .transport.sync import KNOWN_HOSTS, TdxConnection, ping_all
|
||||
from .transport.sync import CALC_HOSTS, KNOWN_HOSTS, TdxConnection, ping_all
|
||||
|
||||
_DEFAULT_PORT = 7709
|
||||
_T = TypeVar("_T")
|
||||
@@ -496,6 +504,84 @@ class TdxClient:
|
||||
break
|
||||
return bytes(full_data)
|
||||
|
||||
@staticmethod
|
||||
def _download_from_host(
|
||||
host: str, filename: str, port: int = 7709, timeout: float = 15.0
|
||||
) -> bytes:
|
||||
"""从指定服务器创建临时连接并下载文件。"""
|
||||
conn = TdxConnection(host, port, timeout)
|
||||
try:
|
||||
conn.connect()
|
||||
full_data = bytearray()
|
||||
pos = 0
|
||||
chunk_size = 30000
|
||||
while True:
|
||||
chunk = conn.execute(GetReportFileCmd(filename, pos, chunk_size))
|
||||
if not chunk:
|
||||
break
|
||||
full_data.extend(chunk)
|
||||
pos += len(chunk)
|
||||
if len(chunk) < chunk_size:
|
||||
break
|
||||
return bytes(full_data)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_financial_file_list(
|
||||
self, host: str = CALC_HOSTS[0]
|
||||
) -> list[FinancialFileInfo]:
|
||||
"""获取可用的历史专业财报文件列表。
|
||||
|
||||
连接到计算服务器，下载 tdxfin/gpcw.txt 并解析。
|
||||
"""
|
||||
data = self._download_from_host(host, "tdxfin/gpcw.txt")
|
||||
raw_list = parse_financial_file_list(data)
|
||||
return [FinancialFileInfo(filename=f, hash=h, filesize=s) for f, h, s in raw_list]
|
||||
|
||||
def get_financial_file(self, filename: str, host: str = CALC_HOSTS[0]) -> bytes:
|
||||
"""从计算服务器下载财报 zip 文件。
|
||||
|
||||
Args:
|
||||
filename: 如 'tdxfin/gpcw20260331.zip'
|
||||
"""
|
||||
return self._download_from_host(host, filename)
|
||||
|
||||
def get_financial_records(
|
||||
self, filename: str, host: str = CALC_HOSTS[0]
|
||||
) -> list[FinancialRecord]:
|
||||
"""下载财报 zip 并解析为每只股票的记录列表。
|
||||
|
||||
Args:
|
||||
filename: 如 'tdxfin/gpcw20260331.zip'
|
||||
"""
|
||||
import io
|
||||
import re
|
||||
import zipfile
|
||||
|
||||
zip_data = self.get_financial_file(filename, host)
|
||||
if not zip_data:
|
||||
return []
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(zip_data)) as zf:
|
||||
dat_names = [n for n in zf.namelist() if n.endswith(".dat")]
|
||||
if not dat_names:
|
||||
return []
|
||||
dat_data = zf.read(dat_names[0])
|
||||
|
||||
m = re.search(r"(\d{8})", filename)
|
||||
report_date = int(m.group(1)) if m else 0
|
||||
|
||||
raw_records = parse_financial_dat(dat_data, report_date)
|
||||
records: list[FinancialRecord] = []
|
||||
for code, market_byte, rdate, fields in raw_records:
|
||||
market = Market.SH if market_byte == b"\x01" else Market.SZ
|
||||
records.append(
|
||||
FinancialRecord(
|
||||
code=code, market=market, report_date=rdate, fields=fields
|
||||
)
|
||||
)
|
||||
return records
|
||||
|
||||
def get_market_stat(self) -> MarketStat:
|
||||
"""获取 A 股全市场涨跌统计概况（基于 880005 行情统计）。
|
||||
|
||||
@@ -923,6 +1009,75 @@ class AsyncTdxClient:
|
||||
break
|
||||
return bytes(full_data)
|
||||
|
||||
@staticmethod
|
||||
async def _async_download_from_host(
|
||||
host: str, filename: str, port: int = 7709, timeout: float = 15.0
|
||||
) -> bytes:
|
||||
"""从指定服务器创建临时异步连接并下载文件。"""
|
||||
conn = AsyncTdxConnection(host, port, timeout)
|
||||
try:
|
||||
await conn.connect()
|
||||
full_data = bytearray()
|
||||
pos = 0
|
||||
chunk_size = 30000
|
||||
while True:
|
||||
chunk = await conn.execute(GetReportFileCmd(filename, pos, chunk_size))
|
||||
if not chunk:
|
||||
break
|
||||
full_data.extend(chunk)
|
||||
pos += len(chunk)
|
||||
if len(chunk) < chunk_size:
|
||||
break
|
||||
return bytes(full_data)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
async def get_financial_file_list(
|
||||
self, host: str = CALC_HOSTS[0]
|
||||
) -> list[FinancialFileInfo]:
|
||||
"""获取可用的历史专业财报文件列表（异步）。"""
|
||||
data = await self._async_download_from_host(host, "tdxfin/gpcw.txt")
|
||||
raw_list = parse_financial_file_list(data)
|
||||
return [FinancialFileInfo(filename=f, hash=h, filesize=s) for f, h, s in raw_list]
|
||||
|
||||
async def get_financial_file(
|
||||
self, filename: str, host: str = CALC_HOSTS[0]
|
||||
) -> bytes:
|
||||
"""从计算服务器下载财报 zip 文件（异步）。"""
|
||||
return await self._async_download_from_host(host, filename)
|
||||
|
||||
async def get_financial_records(
|
||||
self, filename: str, host: str = CALC_HOSTS[0]
|
||||
) -> list[FinancialRecord]:
|
||||
"""下载财报 zip 并解析为记录列表（异步）。"""
|
||||
import io
|
||||
import re
|
||||
import zipfile
|
||||
|
||||
zip_data = await self.get_financial_file(filename, host)
|
||||
if not zip_data:
|
||||
return []
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(zip_data)) as zf:
|
||||
dat_names = [n for n in zf.namelist() if n.endswith(".dat")]
|
||||
if not dat_names:
|
||||
return []
|
||||
dat_data = zf.read(dat_names[0])
|
||||
|
||||
m = re.search(r"(\d{8})", filename)
|
||||
report_date = int(m.group(1)) if m else 0
|
||||
|
||||
raw_records = parse_financial_dat(dat_data, report_date)
|
||||
records: list[FinancialRecord] = []
|
||||
for code, market_byte, rdate, fields in raw_records:
|
||||
market = Market.SH if market_byte == b"\x01" else Market.SZ
|
||||
records.append(
|
||||
FinancialRecord(
|
||||
code=code, market=market, report_date=rdate, fields=fields
|
||||
)
|
||||
)
|
||||
return records
|
||||
|
||||
async def get_market_stat(self) -> MarketStat:
|
||||
"""获取 A 股全市场涨跌统计概况（基于 880005 行情统计）。
|
||||
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""专业财务数据解析（tdxfin/gpcw.txt 列表 + .dat 二进制记录）。"""
|
||||
|
||||
import struct
|
||||
|
||||
|
||||
def parse_financial_file_list(data: bytes) -> list[tuple[str, str, int]]:
|
||||
"""解析 tdxfin/gpcw.txt 的内容。
|
||||
|
||||
每行格式: filename,md5hash,filesize
|
||||
|
||||
Returns:
|
||||
[(filename, hash, filesize), ...]
|
||||
"""
|
||||
if not data:
|
||||
return []
|
||||
text = data.decode("utf-8", errors="replace").strip()
|
||||
results: list[tuple[str, str, int]] = []
|
||||
for line in text.split("\n"):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split(",")
|
||||
if len(parts) >= 3:
|
||||
results.append((parts[0], parts[1], int(parts[2])))
|
||||
return results
|
||||
|
||||
|
||||
def parse_financial_dat(
|
||||
data: bytes,
|
||||
report_date: int = 0,
|
||||
) -> list[tuple[str, int, int, list[float]]]:
|
||||
"""解析 gpcw*.zip 内的 .dat 二进制文件。
|
||||
|
||||
二进制格式（参考 pytdx.crawler.history_financial_crawler）：
|
||||
Header: 20 bytes <1h I 1H 3L
|
||||
- unknown: int16 (h)
|
||||
- report_date: uint32 (I)
|
||||
- max_count: uint16 (H) -- 股票索引条目数
|
||||
- unknown1: uint32 (L)
|
||||
- report_size: uint32 (L) -- 每条股票数据字节长度
|
||||
- unknown2: uint32 (L)
|
||||
Index: max_count 条，每条 11 bytes <6s 1c 1L
|
||||
- code: 6 bytes -- 股票代码
|
||||
- market: 1 byte -- 市场标识 (0=SZ, 1=SH)
|
||||
- file_offset: uint32 -- 绝对偏移（从文件开头算）
|
||||
Data: 在 file_offset 位置读取 report_size/4 个 float32
|
||||
|
||||
Args:
|
||||
data: .dat 文件的完整字节
|
||||
report_date: 报告期 YYYYMMDD（从文件名提取，0 则用 header 中的值）
|
||||
|
||||
Returns:
|
||||
[(code, market_byte, report_date, [float, ...]), ...]
|
||||
"""
|
||||
header_fmt = "<1hI1H3L"
|
||||
header_size = struct.calcsize(header_fmt)
|
||||
if len(data) < header_size:
|
||||
return []
|
||||
|
||||
header = struct.unpack(header_fmt, data[:header_size])
|
||||
max_count = header[2]
|
||||
dat_report_date = header[1]
|
||||
report_size = header[4]
|
||||
|
||||
if report_date == 0:
|
||||
report_date = dat_report_date
|
||||
|
||||
num_fields = report_size // 4
|
||||
if num_fields <= 0:
|
||||
return []
|
||||
|
||||
index_fmt = "<6s1c1L"
|
||||
index_size = struct.calcsize(index_fmt)
|
||||
index_base = header_size
|
||||
|
||||
results: list[tuple[str, int, int, list[float]]] = []
|
||||
report_fmt = f"<{num_fields}f"
|
||||
report_pack_size = struct.calcsize(report_fmt)
|
||||
|
||||
for i in range(max_count):
|
||||
idx_pos = index_base + i * index_size
|
||||
if idx_pos + index_size > len(data):
|
||||
break
|
||||
|
||||
code_bytes, market_byte, file_offset = struct.unpack(
|
||||
index_fmt, data[idx_pos : idx_pos + index_size]
|
||||
)
|
||||
code = code_bytes.decode("ascii", errors="replace").rstrip("\x00")
|
||||
|
||||
if not code or file_offset == 0:
|
||||
continue
|
||||
|
||||
# file_offset 是绝对偏移（从文件开头算）
|
||||
data_pos = file_offset
|
||||
if data_pos + report_pack_size > len(data):
|
||||
continue
|
||||
|
||||
floats = list(struct.unpack(report_fmt, data[data_pos : data_pos + report_pack_size]))
|
||||
results.append((code, market_byte, report_date, floats))
|
||||
|
||||
return results
|
||||
@@ -4,6 +4,8 @@ from .finance import (
|
||||
XDXR_CATEGORY_NAMES,
|
||||
CompanyInfoCategory,
|
||||
FinanceInfo,
|
||||
FinancialFileInfo,
|
||||
FinancialRecord,
|
||||
XdxrRecord,
|
||||
)
|
||||
from .quote import SecurityQuote
|
||||
@@ -22,4 +24,6 @@ __all__ = [
|
||||
"XDXR_CATEGORY_NAMES",
|
||||
"FinanceInfo",
|
||||
"CompanyInfoCategory",
|
||||
"FinancialFileInfo",
|
||||
"FinancialRecord",
|
||||
]
|
||||
|
||||
@@ -128,6 +128,25 @@ class CompanyInfoCategory:
|
||||
length: int = 0 # 内容长度（字节）
|
||||
|
||||
|
||||
@dataclass
|
||||
class FinancialFileInfo:
|
||||
"""财报 zip 文件索引条目（来自 tdxfin/gpcw.txt）。"""
|
||||
|
||||
filename: str # "gpcw20260331.zip"
|
||||
hash: str # MD5 hex digest
|
||||
filesize: int # 字节
|
||||
|
||||
|
||||
@dataclass
|
||||
class FinancialRecord:
|
||||
"""单只股票的一期历史专业财报记录。"""
|
||||
|
||||
code: str # 6 位股票代码
|
||||
market: Market # 市场
|
||||
report_date: int # 报告期 YYYYMMDD
|
||||
fields: list[float] # N 个浮点字段（N = report_size / 4）
|
||||
|
||||
|
||||
@dataclass
|
||||
class TdxBlock:
|
||||
"""通达信板块信息（行业、概念、风格等）"""
|
||||
|
||||
@@ -76,6 +76,11 @@ KNOWN_HOSTS: list[str] = [
|
||||
"111.231.113.208",
|
||||
]
|
||||
|
||||
# 计算服务器（用于下载 tdxfin/ 财务数据）
|
||||
CALC_HOSTS: list[str] = [
|
||||
"120.76.152.87",
|
||||
]
|
||||
|
||||
|
||||
def ping_host(
|
||||
host: str,
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
"""离线测试：专业财务数据解析。"""
|
||||
|
||||
import struct
|
||||
|
||||
from xmtdx.codec.financial import parse_financial_dat, parse_financial_file_list
|
||||
from xmtdx.models.finance import FinancialFileInfo, FinancialRecord
|
||||
|
||||
|
||||
class TestParseFinancialFileList:
|
||||
def test_basic(self) -> None:
|
||||
data = b"gpcw20260331.zip,abc123,5034901\ngpcw20251231.zip,def456,5737165\n"
|
||||
result = parse_financial_file_list(data)
|
||||
assert len(result) == 2
|
||||
assert result[0] == ("gpcw20260331.zip", "abc123", 5034901)
|
||||
assert result[1] == ("gpcw20251231.zip", "def456", 5737165)
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert parse_financial_file_list(b"") == []
|
||||
|
||||
def test_blank_lines_skipped(self) -> None:
|
||||
data = b"\ngpcw.zip,hash,100\n\n"
|
||||
result = parse_financial_file_list(data)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestParseFinancialDat:
|
||||
def _build_dat(
|
||||
self,
|
||||
report_date: int = 20260331,
|
||||
stocks: list[tuple[str, int, list[float]]] | None = None,
|
||||
) -> bytes:
|
||||
"""构造一个最小的 .dat 二进制文件。"""
|
||||
if stocks is None:
|
||||
stocks = [("600519", 1, [1.0, 2.0, 3.0])]
|
||||
|
||||
num_fields = len(stocks[0][2])
|
||||
report_size = num_fields * 4
|
||||
max_count = len(stocks)
|
||||
|
||||
# Header: <1h I 1H 3L = 20 bytes
|
||||
header = struct.pack("<1hI1H3L", 0, report_date, max_count, 0, report_size, 0)
|
||||
|
||||
index_fmt = "<6s1c1L"
|
||||
index_size = struct.calcsize(index_fmt)
|
||||
header_size = struct.calcsize("<1hI1H3L")
|
||||
data_start = header_size + max_count * index_size
|
||||
|
||||
report_fmt = f"<{num_fields}f"
|
||||
|
||||
# 先收集所有数据块，计算绝对偏移
|
||||
data_chunks: list[bytes] = []
|
||||
offset = data_start # 绝对偏移
|
||||
offsets: list[int] = []
|
||||
for code, market_byte, fields in stocks:
|
||||
offsets.append(offset)
|
||||
chunk = struct.pack(report_fmt, *fields)
|
||||
data_chunks.append(chunk)
|
||||
offset += len(chunk)
|
||||
|
||||
# 组装 index
|
||||
index_entries: list[bytes] = []
|
||||
for i, (code, market_byte, _) in enumerate(stocks):
|
||||
index_entries.append(
|
||||
struct.pack(
|
||||
index_fmt, code.encode("ascii"), bytes([market_byte]), offsets[i]
|
||||
)
|
||||
)
|
||||
|
||||
return header + b"".join(index_entries) + b"".join(data_chunks)
|
||||
|
||||
def test_single_stock(self) -> None:
|
||||
dat = self._build_dat(stocks=[("600519", 1, [1.5, 2.5, 3.5])])
|
||||
result = parse_financial_dat(dat, report_date=20260331)
|
||||
assert len(result) == 1
|
||||
code, market, rdate, fields = result[0]
|
||||
assert code == "600519"
|
||||
assert market == b"\x01" # SH
|
||||
assert rdate == 20260331
|
||||
assert len(fields) == 3
|
||||
assert abs(fields[0] - 1.5) < 1e-6
|
||||
|
||||
def test_multiple_stocks(self) -> None:
|
||||
stocks = [
|
||||
("000001", 0, [10.0, 20.0]),
|
||||
("600036", 1, [30.0, 40.0]),
|
||||
]
|
||||
dat = self._build_dat(stocks=stocks)
|
||||
result = parse_financial_dat(dat, report_date=20260630)
|
||||
assert len(result) == 2
|
||||
assert result[0][0] == "000001"
|
||||
assert result[0][1] == b"\x00" # SZ
|
||||
assert result[1][0] == "600036"
|
||||
assert result[1][1] == b"\x01" # SH
|
||||
|
||||
def test_empty_data(self) -> None:
|
||||
assert parse_financial_dat(b"") == []
|
||||
assert parse_financial_dat(b"\x00" * 10) == []
|
||||
|
||||
def test_report_date_from_header(self) -> None:
|
||||
dat = self._build_dat(report_date=20251231, stocks=[("000001", 0, [1.0])])
|
||||
result = parse_financial_dat(dat) # report_date=0, should use header
|
||||
assert result[0][2] == 20251231
|
||||
|
||||
|
||||
class TestFinancialModels:
|
||||
def test_file_info(self) -> None:
|
||||
fi = FinancialFileInfo(filename="gpcw.zip", hash="abc", filesize=100)
|
||||
assert fi.filename == "gpcw.zip"
|
||||
assert fi.filesize == 100
|
||||
|
||||
def test_record(self) -> None:
|
||||
from xmtdx.models.enums import Market
|
||||
|
||||
r = FinancialRecord(
|
||||
code="600519", market=Market.SH, report_date=20260331, fields=[1.0, 2.0]
|
||||
)
|
||||
assert r.market == Market.SH
|
||||
assert len(r.fields) == 2
|
||||
Reference in New Issue
Block a user