diff --git a/.claude/settings.local.json b/.claude/settings.local.json index bbe7f0b..35e19d1 100644 --- a/.claude/settings.local.json +++ b/.claude/settings.local.json @@ -15,7 +15,9 @@ "Bash(python3 *)", "Bash(where mypy *)", "Bash(pip list *)", - "Bash(uv run *)" + "Bash(uv run *)", + "Bash(git stash *)", + "Bash(dir /s /b src\\\\xmtdx)" ] } } diff --git a/CLAUDE.md b/CLAUDE.md index 2818a58..ab9c729 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,4 +51,5 @@ commands 层不依赖 transport,可独立单测。修改 codec 或 commands - ruff: line-length 100, target py310, rules: E/F/I/UP - mypy strict mode -- 纯标准库,零运行时依赖。新增代码不要引入第三方库。 +- 所有 `get_*` 公开方法返回 `pd.DataFrame`(通过 `_df._to_df()` 转换)。内部方法仍使用 dataclass 列表。 +- 依赖:pandas(>=2.0)、tzdata(>=2024.1)。 diff --git a/examples/01_connection/async_connect.py b/examples/01_connection/async_connect.py index a69f4c8..0f6a0cf 100644 --- a/examples/01_connection/async_connect.py +++ b/examples/01_connection/async_connect.py @@ -1,7 +1,8 @@ """演示:异步客户端连接与基本用法。""" import asyncio -from easy_tdx import AsyncTdxClient, Market, KlineCategory + +from easy_tdx import AsyncTdxClient, KlineCategory, Market async def main(): @@ -12,13 +13,8 @@ async def main(): # 自动优选服务器 async with AsyncTdxClient.from_best_host() as c: - bars = await c.get_security_bars( - Market.SH, "600000", KlineCategory.DAY, 0, 5 - ) - for bar in bars: - print(f"{bar.year}-{bar.month:02d}-{bar.day:02d} " - f"开:{bar.open:.2f} 高:{bar.high:.2f} " - f"低:{bar.low:.2f} 收:{bar.close:.2f}") + df = await c.get_security_bars(Market.SH, "600000", KlineCategory.DAY, 0, 5) + print(df.to_string(index=False)) asyncio.run(main()) diff --git a/examples/02_market_info/market_stat.py b/examples/02_market_info/market_stat.py index 0f805df..22f9110 100644 --- a/examples/02_market_info/market_stat.py +++ b/examples/02_market_info/market_stat.py @@ -1,19 +1,7 @@ """演示:获取全市场涨跌统计概况。""" -import pandas as pd from easy_tdx import TdxClient with TdxClient.from_best_host() as c: stat = c.get_market_stat() - df = pd.DataFrame([{ - "上涨": stat.up_count, - "下跌": stat.down_count, - "平盘": stat.neutral_count, - "停牌(估算)": stat.suspended_count, - "总计": stat.total_count, - "成交额(亿)": round(stat.total_amount / 1e8, 2), - "总市值(万亿)": round(stat.total_market_cap / 1e12, 4), - "涨停": stat.limit_up_count, - "跌停": stat.limit_down_count, - }]) - print(df.T.to_string(header=False)) + print(stat) diff --git a/examples/02_market_info/security_list.py b/examples/02_market_info/security_list.py index 7ea135f..6660110 100644 --- a/examples/02_market_info/security_list.py +++ b/examples/02_market_info/security_list.py @@ -1,40 +1,63 @@ -"""演示:获取市场证券列表(分页)。 - -展示 SecurityInfo 全部字段的中文映射与表结构。 -""" +"""演示:获取市场证券列表(分页)。""" import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: - stocks = c.get_security_list(Market.SH, start=0) + df = c.get_security_list(Market.SH, start=0) # 表结构说明 print("=" * 70) print("SecurityInfo 表结构(字段中英文对照)") print("=" * 70) - schema = pd.DataFrame([ - {"英文字段": "market", "中文含义": "市场", "类型": "Market", "说明": "SZ=深圳 SH=上海 BJ=北京"}, - {"英文字段": "code", "中文含义": "证券代码", "类型": "str", "说明": "6位代码,如 600000"}, - {"英文字段": "name", "中文含义": "证券名称", "类型": "str", "说明": "GBK 解码"}, - {"英文字段": "volunit", "中文含义": "成交量单位", "类型": "int", "说明": "1手 = volunit 股"}, - {"英文字段": "decimal_point", "中文含义": "价格小数位", "类型": "int", "说明": "通常为 2"}, - {"英文字段": "pre_close", "中文含义": "昨收价", "类型": "float", "说明": "通达信自定义浮点"}, - {"英文字段": "industry_tdx", "中文含义": "通达信行业", "类型": "str", "说明": "需 get_security_list_all()"}, - {"英文字段": "industry_sw", "中文含义": "申万行业", "类型": "str", "说明": "需 get_security_list_all()"}, - ]) + schema = pd.DataFrame( + [ + { + "英文字段": "market", + "中文含义": "市场", + "类型": "Market", + "说明": "SZ=深圳 SH=上海 BJ=北京", + }, + { + "英文字段": "code", + "中文含义": "证券代码", + "类型": "str", + "说明": "6位代码,如 600000", + }, + {"英文字段": "name", "中文含义": "证券名称", "类型": "str", "说明": "GBK 解码"}, + { + "英文字段": "volunit", + "中文含义": "成交量单位", + "类型": "int", + "说明": "1手 = volunit 股", + }, + { + "英文字段": "decimal_point", + "中文含义": "价格小数位", + "类型": "int", + "说明": "通常为 2", + }, + { + "英文字段": "pre_close", + "中文含义": "昨收价", + "类型": "float", + "说明": "通达信自定义浮点", + }, + { + "英文字段": "industry_tdx", + "中文含义": "通达信行业", + "类型": "str", + "说明": "需 get_security_list_all()", + }, + { + "英文字段": "industry_sw", + "中文含义": "申万行业", + "类型": "str", + "说明": "需 get_security_list_all()", + }, + ] + ) print(schema.to_string(index=False)) - # 全字段中文 DataFrame - print(f"\n沪市第 1 页,共 {len(stocks)} 只:") - df = pd.DataFrame([{ - "市场": s.market.name, - "代码": s.code, - "名称": s.name, - "成交量单位(股/手)": s.volunit, - "价格小数位": s.decimal_point, - "昨收价": s.pre_close, - "通达信行业": s.industry_tdx or "", - "申万行业": s.industry_sw or "", - } for s in stocks]) + print(f"\n沪市第 1 页,共 {len(df)} 只:") print(df.head(20).to_string(index=False)) diff --git a/examples/02_market_info/security_list_all.py b/examples/02_market_info/security_list_all.py index 5a35120..408179a 100644 --- a/examples/02_market_info/security_list_all.py +++ b/examples/02_market_info/security_list_all.py @@ -1,10 +1,10 @@ """演示:获取沪深 A 股完整列表(含行业映射)。 -展示 SecurityInfo 全部字段(含扩展行业字段)的中文映射。 注意:此方法需要拉取 tdxhy.cfg 并遍历全部证券,耗时较长。 """ import logging + import pandas as pd from easy_tdx import TdxClient @@ -13,34 +13,60 @@ logging.basicConfig(level=logging.INFO, format="%(message)s") # timeout 调大到 30 秒,避免全量拉取时分页请求超时 with TdxClient.from_best_host(timeout=30.0) as c: - all_stocks = c.get_security_list_all() + df = c.get_security_list_all() # 表结构说明 print("=" * 70) print("SecurityInfo 表结构(字段中英文对照)") print("=" * 70) - schema = pd.DataFrame([ - {"英文字段": "market", "中文含义": "市场", "类型": "Market", "说明": "SZ=深圳 SH=上海 BJ=北京"}, - {"英文字段": "code", "中文含义": "证券代码", "类型": "str", "说明": "6位代码,如 600000"}, - {"英文字段": "name", "中文含义": "证券名称", "类型": "str", "说明": "GBK 解码"}, - {"英文字段": "volunit", "中文含义": "成交量单位", "类型": "int", "说明": "1手 = volunit 股"}, - {"英文字段": "decimal_point", "中文含义": "价格小数位", "类型": "int", "说明": "通常为 2"}, - {"英文字段": "pre_close", "中文含义": "昨收价", "类型": "float", "说明": "通达信自定义浮点"}, - {"英文字段": "industry_tdx", "中文含义": "通达信行业", "类型": "str", "说明": "如 T1001,来自 tdxhy.cfg"}, - {"英文字段": "industry_sw", "中文含义": "申万行业", "类型": "str", "说明": "如 X500102,来自 tdxhy.cfg"}, - ]) + schema = pd.DataFrame( + [ + { + "英文字段": "market", + "中文含义": "市场", + "类型": "Market", + "说明": "SZ=深圳 SH=上海 BJ=北京", + }, + { + "英文字段": "code", + "中文含义": "证券代码", + "类型": "str", + "说明": "6位代码,如 600000", + }, + {"英文字段": "name", "中文含义": "证券名称", "类型": "str", "说明": "GBK 解码"}, + { + "英文字段": "volunit", + "中文含义": "成交量单位", + "类型": "int", + "说明": "1手 = volunit 股", + }, + { + "英文字段": "decimal_point", + "中文含义": "价格小数位", + "类型": "int", + "说明": "通常为 2", + }, + { + "英文字段": "pre_close", + "中文含义": "昨收价", + "类型": "float", + "说明": "通达信自定义浮点", + }, + { + "英文字段": "industry_tdx", + "中文含义": "通达信行业", + "类型": "str", + "说明": "如 T1001,来自 tdxhy.cfg", + }, + { + "英文字段": "industry_sw", + "中文含义": "申万行业", + "类型": "str", + "说明": "如 X500102,来自 tdxhy.cfg", + }, + ] + ) print(schema.to_string(index=False)) - # 全字段中文 DataFrame - print(f"\n沪深 A 股总数: {len(all_stocks)}") - df = pd.DataFrame([{ - "市场": s.market.name, - "代码": s.code, - "名称": s.name, - "成交量单位(股/手)": s.volunit, - "价格小数位": s.decimal_point, - "昨收价": s.pre_close, - "通达信行业": s.industry_tdx or "", - "申万行业": s.industry_sw or "", - } for s in all_stocks]) + print(f"\n沪深 A 股总数: {len(df)}") print(df.head(20).to_string(index=False)) diff --git a/examples/02_market_info/security_quotes.py b/examples/02_market_info/security_quotes.py index 763be38..dda7333 100644 --- a/examples/02_market_info/security_quotes.py +++ b/examples/02_market_info/security_quotes.py @@ -1,7 +1,6 @@ """演示:批量获取实时五档行情。最多支持 80 只/次。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: stocks = [ @@ -10,16 +9,10 @@ with TdxClient.from_best_host() as c: (Market.SZ, "000001"), # 平安银行 (Market.SZ, "000858"), # 五粮液 ] - quotes = c.get_security_quotes(stocks) - df = pd.DataFrame([{ - "代码": q.code, - "现价": q.price, - "涨跌幅%": (q.price - q.pre_close) / q.pre_close * 100, - "今开": q.open, - "最高": q.high, - "最低": q.low, - "昨收": q.pre_close, - "成交量(手)": q.vol, - "成交额": q.amount, - } for q in quotes]) - print(df.to_string(index=False)) + df = c.get_security_quotes(stocks) + df["change_pct"] = (df["price"] - df["pre_close"]) / df["pre_close"] * 100 + print( + df[ + ["code", "price", "change_pct", "open", "high", "low", "pre_close", "vol", "amount"] + ].to_string(index=False) + ) diff --git a/examples/03_kline/index_bars.py b/examples/03_kline/index_bars.py index 1ee5b03..57284e7 100644 --- a/examples/03_kline/index_bars.py +++ b/examples/03_kline/index_bars.py @@ -6,20 +6,9 @@ 创业板指: Market.SZ, "399006" """ -import pandas as pd -from easy_tdx import TdxClient, Market, KlineCategory +from easy_tdx import KlineCategory, Market, TdxClient with TdxClient.from_best_host() as c: - bars = c.get_index_bars(Market.SH, "999999", KlineCategory.DAY, 0, 10) - df = pd.DataFrame([{ - "日期": f"{b.year}-{b.month:02d}-{b.day:02d}", - "开盘": b.open, - "最高": b.high, - "最低": b.low, - "收盘": b.close, - "成交量": b.vol, - "成交额": b.amount, - } for b in reversed(bars)]) + df = c.get_index_bars(Market.SH, "999999", KlineCategory.DAY, 0, 10) print("上证指数 日K线:") - fmt = {"成交量": lambda x: f"{x:,.0f}", "成交额": lambda x: f"{x:,.0f}"} - print(df.to_string(index=False, formatters=fmt)) + print(df.to_string(index=False)) diff --git a/examples/03_kline/security_bars.py b/examples/03_kline/security_bars.py index fa4d6e3..d52ca2b 100644 --- a/examples/03_kline/security_bars.py +++ b/examples/03_kline/security_bars.py @@ -5,19 +5,9 @@ K 线类别: KlineCategory.DAY / WEEK / MONTH / YEAR """ -import pandas as pd -from easy_tdx import TdxClient, Market, KlineCategory +from easy_tdx import KlineCategory, Market, TdxClient with TdxClient.from_best_host() as c: - bars = c.get_security_bars(Market.SZ, "002176", KlineCategory.DAY, 0, 100) - df = pd.DataFrame([{ - "日期": f"{b.year}-{b.month:02d}-{b.day:02d}", - "开盘": b.open, - "最高": b.high, - "最低": b.low, - "收盘": b.close, - "成交量": b.vol, - "成交额": b.amount, - } for b in reversed(bars)]) - print("上证指数 日K线:") + df = c.get_security_bars(Market.SZ, "002176", KlineCategory.DAY, 0, 100) + print("江特电机 日K线:") print(df.to_string(index=False)) diff --git a/examples/04_minute/history_minute_data.py b/examples/04_minute/history_minute_data.py index 0612f3f..ea83929 100644 --- a/examples/04_minute/history_minute_data.py +++ b/examples/04_minute/history_minute_data.py @@ -1,15 +1,9 @@ """演示:获取历史某日分时数据。date 参数为 YYYYMMDD 格式的整数。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: date = 20250110 - bars = c.get_history_minute_time_data(Market.SH, "600000", date) - df = pd.DataFrame([{ - "序号": i + 1, - "价格": bar.price, - "成交量": bar.vol, - } for i, bar in enumerate(bars)]) + df = c.get_history_minute_time_data(Market.SH, "600000", date) print(f"浦发银行 {date} 分时数据,共 {len(df)} 条:") - print(df.to_string(index=False)) + print(df.head(20).to_string(index=False)) diff --git a/examples/04_minute/minute_time_data.py b/examples/04_minute/minute_time_data.py index 61498f1..be5c2c6 100644 --- a/examples/04_minute/minute_time_data.py +++ b/examples/04_minute/minute_time_data.py @@ -1,14 +1,8 @@ """演示:获取今日分时数据(240 条)。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: - bars = c.get_minute_time_data(Market.SH, "600000") - df = pd.DataFrame([{ - "序号": i + 1, - "价格": bar.price, - "成交量": bar.vol, - } for i, bar in enumerate(bars)]) + df = c.get_minute_time_data(Market.SH, "600000") print(f"浦发银行今日分时,共 {len(df)} 条:") - print(df.to_string(index=False)) + print(df.head(20).to_string(index=False)) diff --git a/examples/05_transaction/history_transaction.py b/examples/05_transaction/history_transaction.py index e90a10f..8150dbc 100644 --- a/examples/05_transaction/history_transaction.py +++ b/examples/05_transaction/history_transaction.py @@ -1,16 +1,10 @@ """演示:获取历史逐笔成交数据。date 参数为 YYYYMMDD 格式的整数。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: date = 20250110 - records = c.get_history_transaction_data(Market.SH, "600000", date, 0, 20) - df = pd.DataFrame([{ - "时间": f"{r.hour:02d}:{r.minute:02d}", - "成交价": r.price, - "成交量": r.vol, - "方向": "买" if r.buyorsell == 0 else "卖", - } for r in records]) + df = c.get_history_transaction_data(Market.SH, "600000", date, 0, 20) + df["方向"] = df["buyorsell"].map({0: "买", 1: "卖", 2: "中性", 8: "集合竞价"}) print(f"浦发银行 {date} 最近 {len(df)} 笔成交:") - print(df.to_string(index=False)) + print(df[["datetime", "price", "vol", "方向"]].to_string(index=False)) diff --git a/examples/05_transaction/transaction_data.py b/examples/05_transaction/transaction_data.py index 45e3ce6..395f13e 100644 --- a/examples/05_transaction/transaction_data.py +++ b/examples/05_transaction/transaction_data.py @@ -1,15 +1,9 @@ """演示:获取当日逐笔成交数据。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: - records = c.get_transaction_data(Market.SH, "600000", 0, 20) - df = pd.DataFrame([{ - "时间": f"{r.hour:02d}:{r.minute:02d}", - "成交价": r.price, - "成交量": r.vol, - "方向": "买" if r.buyorsell == 0 else "卖", - } for r in records]) + df = c.get_transaction_data(Market.SH, "600000", 0, 20) + df["方向"] = df["buyorsell"].map({0: "买", 1: "卖", 2: "中性", 8: "集合竞价"}) print(f"浦发银行最近 {len(df)} 笔成交:") - print(df.to_string(index=False)) + print(df[["datetime", "price", "vol", "方向"]].to_string(index=False)) diff --git a/examples/06_finance/company_info.py b/examples/06_finance/company_info.py index c72d280..16b3ffe 100644 --- a/examples/06_finance/company_info.py +++ b/examples/06_finance/company_info.py @@ -1,7 +1,6 @@ """演示:获取公司信息目录与各个分类的详细内容。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient CODE = "600519" NAME = "贵州茅台" @@ -28,51 +27,36 @@ SHOW_CATEGORIES = [ ] -def show_categories(categories): - """显示公司信息目录。""" - df = pd.DataFrame([{ - "目录名": cat.name, - "文件名": cat.filename, - "起始偏移": cat.start, - "内容长度": cat.length, - } for cat in categories]) - print(f"{NAME} 公司信息目录:") - print(df.to_string(index=False)) - - def show_category_content(client, categories, category_name, max_chars=500): """获取并展示指定分类的内容。""" - cat = next((c for c in categories if c.name == category_name), None) - if not cat: + row = categories[categories["name"] == category_name] + if row.empty: print(f" 未找到分类: {category_name}") return + r = row.iloc[0] content = client.get_company_info_content( - MARKET, CODE, cat.filename, cat.start, cat.length + MARKET, CODE, r["filename"], int(r["start"]), int(r["length"]) ) text = content.strip() if len(text) > max_chars: text = text[:max_chars] + f"\n... (共 {len(content.strip())} 字,仅显示前 {max_chars} 字)" - print(f"\n{'='*60}") - print(f"【{cat.name}】 (共 {cat.length} 字节)") - print(f"{'='*60}") + print(f"\n{'=' * 60}") + print(f"【{r['name']}】 (共 {r['length']} 字节)") + print(f"{'=' * 60}") print(text) -def show_all_categories(client, categories): - """依次展示所有 SHOW_CATEGORIES 中列出的分类内容。""" - for name in SHOW_CATEGORIES: - show_category_content(client, categories, name) - - with TdxClient.from_best_host() as c: categories = c.get_company_info_category(MARKET, CODE) # 1. 显示目录 - show_categories(categories) + print(f"{NAME} 公司信息目录:") + print(categories.to_string(index=False)) # 2. 显示所有分类内容(每个分类默认只显示前500字) - show_all_categories(c, categories) + for name in SHOW_CATEGORIES: + show_category_content(c, categories, name) # 3. 也可以单独获取某个分类的完整内容,例如: # show_category_content(c, categories, "公司概况", max_chars=99999) diff --git a/examples/06_finance/finance_info.py b/examples/06_finance/finance_info.py index 2ce90cc..02ee3cc 100644 --- a/examples/06_finance/finance_info.py +++ b/examples/06_finance/finance_info.py @@ -1,21 +1,8 @@ """演示:获取最新财务数据。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: info = c.get_finance_info(Market.SH, "600519") - df = pd.DataFrame([ - {"项目": "总股本(万股)", "数值": info.zong_guben}, - {"项目": "流通股本(万股)", "数值": info.liutong_guben}, - {"项目": "每股净资产", "数值": info.meigujing_zichan}, - {"项目": "净利润(元)", "数值": info.jing_lirun}, - {"项目": "主营收入(元)", "数值": info.zhuying_shouru}, - {"项目": "主营利润(元)", "数值": info.zhuying_lirun}, - {"项目": "净资产(元)", "数值": info.jing_zichan}, - {"项目": "总资产(元)", "数值": info.zong_zichan}, - {"项目": "股东人数", "数值": info.gudong_renshu}, - {"项目": "上市日期", "数值": info.ipo_date}, - ]) print("贵州茅台 最新财务数据:") - print(df.to_string(index=False, formatters={"数值": lambda x: f"{x:,.0f}"})) + print(info.T.to_string(header=False)) diff --git a/examples/06_finance/price_limits.py b/examples/06_finance/price_limits.py index b77d75d..d49370d 100644 --- a/examples/06_finance/price_limits.py +++ b/examples/06_finance/price_limits.py @@ -1,23 +1,16 @@ """演示:计算个股涨跌停价格。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient CODE = "600519" NAME = "贵州茅台" with TdxClient.from_best_host() as c: quotes = c.get_security_quotes([(Market.SH, CODE)]) - if quotes: - q = quotes[0] - limit_up, limit_down = c.get_price_limits( - Market.SH, CODE, NAME, q.pre_close - ) - df = pd.DataFrame([{ - "代码": CODE, - "名称": NAME, - "昨收": q.pre_close, - "涨停价": limit_up, - "跌停价": limit_down, - }]) - print(df.to_string(index=False)) + if not quotes.empty: + q = quotes.iloc[0] + limit_up, limit_down = c.get_price_limits(Market.SH, CODE, NAME, q["pre_close"]) + print(f"代码: {CODE} 名称: {NAME}") + print(f"昨收: {q['pre_close']}") + print(f"涨停价: {limit_up}") + print(f"跌停价: {limit_down}") diff --git a/examples/06_finance/xdxr_info.py b/examples/06_finance/xdxr_info.py index 5f0431f..cecea17 100644 --- a/examples/06_finance/xdxr_info.py +++ b/examples/06_finance/xdxr_info.py @@ -1,17 +1,8 @@ """演示:获取除权除息历史记录。""" -import pandas as pd -from easy_tdx import TdxClient, Market, XDXR_CATEGORY_NAMES +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: - records = c.get_xdxr_info(Market.SH, "600519") - df = pd.DataFrame([{ - "日期": f"{r.year}-{r.month:02d}-{r.day:02d}", - "类型": XDXR_CATEGORY_NAMES.get(r.category, f"未知({r.category})"), - "每股分红(元)": r.fenhong, - "送转股比例": r.songzhuangu, - "配股价": r.peigujia, - "配股比例": r.peigu, - } for r in records]) + df = c.get_xdxr_info(Market.SH, "600519") print(f"贵州茅台 除权除息记录,共 {len(df)} 条:") print(df.tail(10).to_string(index=False)) diff --git a/examples/07_block/block_info.py b/examples/07_block/block_info.py index 8cee2ad..2f6af52 100644 --- a/examples/07_block/block_info.py +++ b/examples/07_block/block_info.py @@ -6,16 +6,9 @@ 'block_fg.dat' - 风格板块 """ -import pandas as pd from easy_tdx import TdxClient with TdxClient.from_best_host() as c: - blocks = c.get_block_info("block_gn.dat") - df = pd.DataFrame([{ - "板块名称": b.name, - "分类": b.category, - "成分股数": b.count, - "代码(前5)": ", ".join(b.codes[:5]), - } for b in blocks]) + df = c.get_block_info("block_gn.dat") print(f"概念板块,共 {len(df)} 个:") - print(df.head(20).to_string(index=False)) + print(df[["name", "category", "count"]].head(20).to_string(index=False)) diff --git a/examples/08_fund_flow/fund_flow.py b/examples/08_fund_flow/fund_flow.py index a41b7a8..75da0dc 100644 --- a/examples/08_fund_flow/fund_flow.py +++ b/examples/08_fund_flow/fund_flow.py @@ -4,16 +4,20 @@ """ import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: flow = c.get_fund_flow(Market.SH, "600519") - df = pd.DataFrame([ - {"级别": "超大单", "流入(亿)": flow.super_in / 1e8, "流出(亿)": flow.super_out / 1e8}, - {"级别": "大单", "流入(亿)": flow.large_in / 1e8, "流出(亿)": flow.large_out / 1e8}, - {"级别": "中单", "流入(亿)": flow.medium_in / 1e8, "流出(亿)": flow.medium_out / 1e8}, - {"级别": "小单", "流入(亿)": flow.small_in / 1e8, "流出(亿)": flow.small_out / 1e8}, - ]) + # flow 是单行 DataFrame,转换为万元便于阅读 + in_cols = ["super_in", "large_in", "medium_in", "small_in"] + out_cols = ["super_out", "large_out", "medium_out", "small_out"] + df = pd.DataFrame( + { + "级别": ["超大单", "大单", "中单", "小单"], + "流入(亿)": [flow[c].iloc[0] / 1e8 for c in in_cols], + "流出(亿)": [flow[c].iloc[0] / 1e8 for c in out_cols], + } + ) df["净流入(亿)"] = df["流入(亿)"] - df["流出(亿)"] print("贵州茅台 当日资金流向:") print(df.to_string(index=False)) diff --git a/examples/08_fund_flow/history_fund_flow.py b/examples/08_fund_flow/history_fund_flow.py index 2eae460..f595af9 100644 --- a/examples/08_fund_flow/history_fund_flow.py +++ b/examples/08_fund_flow/history_fund_flow.py @@ -1,15 +1,8 @@ """演示:获取个股历史日线资金流向序列。""" -import pandas as pd -from easy_tdx import TdxClient, Market +from easy_tdx import Market, TdxClient with TdxClient.from_best_host() as c: - flows = c.get_history_fund_flow(Market.SH, "600519", 0, 10) - df = pd.DataFrame([{ - "日期": f"{f.year}-{f.month:02d}-{f.day:02d}", - "超大单净流入(亿)": (f.super_in - f.super_out) / 1e8, - "大单净流入(亿)": (f.large_in - f.large_out) / 1e8, - "主力净流入(亿)": f.main_net_inflow / 1e8, - } for f in flows]) + df = c.get_history_fund_flow(Market.SH, "600519", 0, 10) print(f"贵州茅台 历史资金流向,共 {len(df)} 天:") print(df.to_string(index=False)) diff --git a/examples/09_file_download/report_file.py b/examples/09_file_download/report_file.py index 2d7cbe2..bf3e989 100644 --- a/examples/09_file_download/report_file.py +++ b/examples/09_file_download/report_file.py @@ -59,8 +59,7 @@ with TdxClient.from_best_host() as c: 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(blocks[["name", "category", "count"]].head(5).to_string(index=False)) print(f" ... 共 {len(blocks)} 个") # ── 2. 计算服务器:专业财务数据 ──────────────────────── @@ -74,29 +73,25 @@ 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(file_list.head(5).to_string(index=False)) 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:,} 字节)") + real_files = file_list[file_list["filesize"] > 10000] + if not real_files.empty: + latest = real_files.iloc[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 = 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)}") + if not records.empty: + print(records[["market", "code", "report_date"]].head(5).to_string(index=False)) print(f" ... 共 {len(records)} 只") - r = records[0] - print(f" 示例: {r.market.name} {r.code}, 报告期={r.report_date}, 字段数={len(r.fields)}") diff --git a/examples/10_offline/block_data.py b/examples/10_offline/block_data.py index 2f8f52b..7a7d8ec 100644 --- a/examples/10_offline/block_data.py +++ b/examples/10_offline/block_data.py @@ -54,8 +54,15 @@ if need_fetch: print(f"\n正在连接服务器获取 {len(need_fetch)} 个板块文件...") with TdxClient.from_best_host() as c: for name in need_fetch: - blocks = c.get_block_info(name) - _print_blocks(blocks, f"{block_labels[name]} ({name}, 网络)") + df = c.get_block_info(name) + print(f"\n{block_labels[name]} ({name}, 网络) ({len(df)} 个板块):") + for _, row in df.head(5).iterrows(): + codes = row["codes"] + codes_preview = ", ".join(str(c) for c in codes[:5]) + suffix = "..." if len(codes) > 5 else "" + print(f" {row['name']} ({row['count']}只): {codes_preview}{suffix}") + if len(df) > 5: + print(f" ... 还有 {len(df) - 5} 个板块") # --- 自定义板块 --- print(f"\n{'=' * 60}") diff --git a/pyproject.toml b/pyproject.toml index b5ef404..9650f07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,11 +8,10 @@ version = "1.0.0" description = "通达信 TCP 协议行情数据客户端,支持在线行情与离线本地数据读取" readme = "README.md" requires-python = ">=3.10" -dependencies = ["tzdata>=2024.1"] +dependencies = ["pandas>=2.0", "tzdata>=2024.1"] [project.optional-dependencies] dev = ["pytest>=8.0", "pytest-cov", "mypy>=1.9", "ruff>=0.4"] -pandas = ["pandas>=2.0"] [tool.hatch.build.targets.wheel] packages = ["src/easy_tdx"] diff --git a/src/easy_tdx/_df.py b/src/easy_tdx/_df.py new file mode 100644 index 0000000..9fe80ce --- /dev/null +++ b/src/easy_tdx/_df.py @@ -0,0 +1,111 @@ +"""Dataclass → DataFrame 转换工具。""" + +from __future__ import annotations + +from dataclasses import asdict, is_dataclass +from typing import Any + +import pandas as pd + + +def _to_df(data: Any) -> pd.DataFrame: + """将 list[dataclass] 或单个 dataclass 转为 DataFrame。 + + 自动丢弃以 ``_`` 开头的内部字段(如 ``_raw``)。 + 仅处理 year/month/day(无 hour/minute)→ date 的合并; + SecurityBar 的完整 datetime 合并由调用方按周期决定。 + """ + if isinstance(data, list): + if not data: + return pd.DataFrame() + rows = [] + for item in data: + d = _clean_dict(item) + rows.append(d) + return pd.DataFrame(rows) + if is_dataclass(data) and not isinstance(data, type): + return pd.DataFrame([_clean_dict(data)]) + raise TypeError(f"不支持转换为 DataFrame 的类型: {type(data)}") + + +def _clean_dict(item: Any) -> dict[str, Any]: + d = asdict(item) + d = {k: v for k, v in d.items() if not k.startswith("_")} + return _merge_datetime_fields(d) + + +def _merge_datetime_fields(d: dict[str, Any]) -> dict[str, Any]: + """将仅含 year/month/day(无 hour/minute)的模型合并为 date 列。""" + if all(k in d for k in ("year", "month", "day")) and not all( + k in d for k in ("hour", "minute") + ): + dt = pd.Timestamp(year=d["year"], month=d["month"], day=d["day"]) + result: dict[str, Any] = {"date": dt} + result.update({k: v for k, v in d.items() if k not in {"year", "month", "day"}}) + return result + return d + + +def _merge_bar_datetime(df: pd.DataFrame, daily_plus: bool) -> pd.DataFrame: + """根据 K 线周期将 SecurityBar 的分散字段合并为 date 或 datetime。 + + Args: + daily_plus: True 表示日线及以上周期(DAY/WEEK/MONTH/YEAR),只保留 date; + False 表示分钟线(MIN_1/5/15/30/60),保留完整 datetime。 + """ + if df.empty or "year" not in df.columns: + return df + date_str = ( + df["year"].astype(str) + + "-" + + df["month"].astype(str).str.zfill(2) + + "-" + + df["day"].astype(str).str.zfill(2) + ) + if daily_plus: + df.insert(0, "date", pd.to_datetime(date_str)) + else: + full_str = ( + date_str + + " " + + df["hour"].astype(str).str.zfill(2) + + ":" + + df["minute"].astype(str).str.zfill(2) + ) + df.insert(0, "datetime", pd.to_datetime(full_str)) + df.drop(columns=["year", "month", "day", "hour", "minute"], inplace=True) + return df + + +def _merge_txn_datetime(df: pd.DataFrame, date_int: int) -> pd.DataFrame: + """将逐笔成交的 date + hour:minute 合并为 datetime 列。""" + if df.empty or "hour" not in df.columns: + return df + year = date_int // 10000 + month = (date_int // 100) % 100 + day = date_int % 100 + base = pd.Timestamp(year=year, month=month, day=day) + offsets = pd.to_timedelta(df["hour"] * 3600 + df["minute"] * 60, unit="s") + df.insert(0, "datetime", base + offsets) + df.drop(columns=["hour", "minute"], inplace=True) + return df + + +def _add_minute_datetime(df: pd.DataFrame, date_int: int) -> pd.DataFrame: + """为分时 DataFrame 添加 datetime 列(从 bar 索引计算时间)。 + + A 股分时 240 条:0-119 = 9:30~11:29(上午),120-239 = 13:00~14:59(下午)。 + """ + if df.empty: + return df + year = date_int // 10000 + month = (date_int // 100) % 100 + day = date_int % 100 + base = pd.Timestamp(year=year, month=month, day=day) + n = len(df) + morning = list(range(9 * 60 + 30, 9 * 60 + 30 + 120)) + afternoon = list(range(13 * 60, 13 * 60 + 120)) + all_minutes = (morning + afternoon)[:n] + offsets = pd.to_timedelta(all_minutes, unit="m") + df.insert(0, "datetime", base + offsets) + return df diff --git a/src/easy_tdx/client.py b/src/easy_tdx/client.py index 6f9d1f0..737f9fd 100644 --- a/src/easy_tdx/client.py +++ b/src/easy_tdx/client.py @@ -11,6 +11,9 @@ from types import TracebackType from typing import TypeVar from zoneinfo import ZoneInfo +import pandas as pd + +from ._df import _add_minute_datetime, _merge_bar_datetime, _merge_txn_datetime, _to_df 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 @@ -20,7 +23,7 @@ from .commands.block_info import GetBlockInfoCmd, GetBlockInfoMetaCmd from .commands.company_info import GetCompanyInfoCategoryCmd, GetCompanyInfoContentCmd from .commands.finance_info import GetFinanceInfoCmd from .commands.fund_flow import GetHistoryFundFlowCmd -from .commands.minute_time import GetHistoryMinuteTimeDataCmd, GetMinuteTimeDataCmd +from .commands.minute_time import GetHistoryMinuteTimeDataCmd from .commands.report_file import GetReportFileCmd from .commands.security_bars import GetIndexBarsCmd, GetSecurityBarsCmd from .commands.security_count import GetSecurityCountCmd @@ -32,23 +35,27 @@ from .exceptions import TdxConnectionError from .models.bar import SecurityBar from .models.enums import KlineCategory, Market 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 .models.timeseries import TransactionRecord from .transport.async_ import AsyncTdxConnection from .transport.sync import CALC_HOSTS, KNOWN_HOSTS, TdxConnection, ping_all _DEFAULT_PORT = 7709 _T = TypeVar("_T") _SHANGHAI_TZ = ZoneInfo("Asia/Shanghai") +_DAILY_PLUS = frozenset( + { + KlineCategory.DAY, + KlineCategory.WEEK, + KlineCategory.MONTH, + KlineCategory.YEAR, + KlineCategory.YEAR_ALT, + } +) def _today_in_shanghai() -> int: @@ -88,9 +95,7 @@ def _classify_fund_flow(records: list[TransactionRecord]) -> FundFlow: for record in records: amount = record.price * record.vol * 100.0 - direction = ( - "in" if record.buyorsell == 0 else "out" if record.buyorsell == 1 else None - ) + direction = "in" if record.buyorsell == 0 else "out" if record.buyorsell == 1 else None if not direction: continue @@ -278,11 +283,11 @@ class TdxClient: """获取市场证券总数。""" return self._execute(GetSecurityCountCmd(market)) - def get_security_list(self, market: Market, start: int) -> list[SecurityInfo]: + def get_security_list(self, market: Market, start: int) -> pd.DataFrame: """获取证券列表(每页约1000条,按 start 分页)。""" - return self._execute(GetSecurityListCmd(market, start)) + return _to_df(self._execute(GetSecurityListCmd(market, start))) - def get_security_list_all(self, pages: int | str = "all") -> list[SecurityInfo]: + def get_security_list_all(self, pages: int | str = "all") -> pd.DataFrame: """获取沪深 A 股完整证券列表,并自动挂载行业信息。 Args: @@ -299,7 +304,7 @@ class TdxClient: cached = _load_cache() if cached is not None: log.info("从缓存加载沪深 A 股列表,共 %d 只", len(cached)) - return cached + return _to_df(cached) # 计算每个市场的最大起始偏移 def _max_start(count: int) -> int: @@ -324,15 +329,18 @@ class TdxClient: total_pages = (limit + 999) // 1000 for page_idx, start in enumerate(range(0, limit, 1000)): try: - stocks = self.get_security_list(market, start) + stocks = self._execute(GetSecurityListCmd(market, start)) except Exception: - log.warning("%s 第 %d/%d 页获取失败,跳过", market.name, page_idx + 1, total_pages) + log.warning( + "%s 第 %d/%d 页获取失败,跳过", market.name, page_idx + 1, total_pages + ) continue - log.info("%s 第 %d/%d 页: %d 条", market.name, page_idx + 1, total_pages, len(stocks)) + log.info( + "%s 第 %d/%d 页: %d 条", market.name, page_idx + 1, total_pages, len(stocks) + ) for s in stocks: - is_a_share = ( - (market == Market.SH and s.code.startswith(("60", "68"))) - or (market == Market.SZ and s.code.startswith(("00", "30"))) + is_a_share = (market == Market.SH and s.code.startswith(("60", "68"))) or ( + market == Market.SZ and s.code.startswith(("00", "30")) ) if is_a_share: if s.code in industry_map: @@ -344,13 +352,11 @@ class TdxClient: if pages == "all": _save_cache(all_stocks) - return all_stocks + return _to_df(all_stocks) - def get_security_quotes( - self, stocks: list[tuple[Market, str]] - ) -> list[SecurityQuote]: + def get_security_quotes(self, stocks: list[tuple[Market, str]]) -> pd.DataFrame: """批量获取实时五档行情(最多80只/次)。""" - return self._execute(GetSecurityQuotesCmd(stocks)) + return _to_df(self._execute(GetSecurityQuotesCmd(stocks))) def get_price_limits( self, market: Market, code: str, name: str, pre_close: float @@ -363,8 +369,8 @@ class TdxClient: no_limit_window_days = get_no_limit_window_days(market, code, name) if no_limit_window_days > 0: try: - bars = self.get_security_bars( - market, code, KlineCategory.DAY, 0, no_limit_window_days + 1 + bars = self._execute( + GetSecurityBarsCmd(market, code, KlineCategory.DAY, 0, no_limit_window_days + 1) ) listed_days = len(bars) except Exception: @@ -389,9 +395,10 @@ class TdxClient: category: KlineCategory, start: int, count: int = 800, - ) -> list[SecurityBar]: + ) -> pd.DataFrame: """获取 K 线数据(最多800条/次,按 start 分页)。""" - return self._execute(GetSecurityBarsCmd(market, code, category, start, count)) + df = _to_df(self._execute(GetSecurityBarsCmd(market, code, category, start, count))) + return _merge_bar_datetime(df, category in _DAILY_PLUS) def get_index_bars( self, @@ -400,30 +407,25 @@ class TdxClient: category: KlineCategory, start: int, count: int = 800, - ) -> list[SecurityBar]: + ) -> pd.DataFrame: """获取指数 K 线数据。""" - return self._execute(GetIndexBarsCmd(market, code, category, start, count)) + df = _to_df(self._execute(GetIndexBarsCmd(market, code, category, start, count))) + return _merge_bar_datetime(df, category in _DAILY_PLUS) # ------------------------------------------------------------------ # # 分时 # ------------------------------------------------------------------ # - def get_minute_time_data(self, market: Market, code: str) -> list[MinuteBar]: - """获取今日分时数据(240条)。""" + def get_minute_time_data(self, market: Market, code: str) -> pd.DataFrame: + """获取今日分时数据(240条,走历史分时接口)。""" today = _today_in_shanghai() - try: - bars = self.get_history_minute_time_data(market, code, today) - if bars: - return bars - except Exception: - pass - return self._execute(GetMinuteTimeDataCmd(market, code)) + bars = self._execute(GetHistoryMinuteTimeDataCmd(market, code, today)) + return _add_minute_datetime(_to_df(bars), today) - def get_history_minute_time_data( - self, market: Market, code: str, date: int - ) -> list[MinuteBar]: + def get_history_minute_time_data(self, market: Market, code: str, date: int) -> pd.DataFrame: """获取历史某日分时数据(date: YYYYMMDD)。""" - return self._execute(GetHistoryMinuteTimeDataCmd(market, code, date)) + bars = self._execute(GetHistoryMinuteTimeDataCmd(market, code, date)) + return _add_minute_datetime(_to_df(bars), date) # ------------------------------------------------------------------ # # 逐笔成交 @@ -431,45 +433,41 @@ class TdxClient: def get_transaction_data( self, market: Market, code: str, start: int, count: int = 800 - ) -> list[TransactionRecord]: + ) -> pd.DataFrame: """获取当日逐笔成交(分页)。""" - return self._execute(GetTransactionDataCmd(market, code, start, count)) + df = _to_df(self._execute(GetTransactionDataCmd(market, code, start, count))) + return _merge_txn_datetime(df, _today_in_shanghai()) def get_history_transaction_data( self, market: Market, code: str, date: int, start: int, count: int = 800 - ) -> list[TransactionRecord]: + ) -> pd.DataFrame: """获取历史逐笔成交(date: YYYYMMDD,分页)。""" - return self._execute( - GetHistoryTransactionDataCmd(market, code, date, start, count) - ) + df = _to_df(self._execute(GetHistoryTransactionDataCmd(market, code, date, start, count))) + return _merge_txn_datetime(df, date) # ------------------------------------------------------------------ # # 财务 / 公司 # ------------------------------------------------------------------ # - def get_xdxr_info(self, market: Market, code: str) -> list[XdxrRecord]: + def get_xdxr_info(self, market: Market, code: str) -> pd.DataFrame: """获取除权除息历史记录。""" - return self._execute(GetXdxrInfoCmd(market, code)) + return _to_df(self._execute(GetXdxrInfoCmd(market, code))) - def get_finance_info(self, market: Market, code: str) -> FinanceInfo: + def get_finance_info(self, market: Market, code: str) -> pd.DataFrame: """获取最新财务数据。""" - return self._execute(GetFinanceInfoCmd(market, code)) + return _to_df(self._execute(GetFinanceInfoCmd(market, code))) - def get_company_info_category( - self, market: Market, code: str - ) -> list[CompanyInfoCategory]: + def get_company_info_category(self, market: Market, code: str) -> pd.DataFrame: """获取公司信息文件目录。""" - return self._execute(GetCompanyInfoCategoryCmd(market, code)) + return _to_df(self._execute(GetCompanyInfoCategoryCmd(market, code))) def get_company_info_content( self, market: Market, code: str, filename: str, offset: int, length: int ) -> str: """读取公司信息文本。""" - return self._execute( - GetCompanyInfoContentCmd(market, code, filename, offset, length) - ) + return self._execute(GetCompanyInfoContentCmd(market, code, filename, offset, length)) - def get_block_info(self, filename: str) -> list[TdxBlock]: + def get_block_info(self, filename: str) -> pd.DataFrame: """获取并解析板块文件(行业、概念、风格等)。 常用文件名: @@ -487,7 +485,7 @@ class TdxClient: break full_data.extend(chunk) pos += len(chunk) - return parse_block_dat(bytes(full_data), filename) + return _to_df(parse_block_dat(bytes(full_data), filename)) def get_report_file(self, filename: str) -> bytes: """从服务器拉取大文件(如 'base_info.zip')。""" @@ -527,16 +525,14 @@ class TdxClient: finally: conn.close() - def get_financial_file_list( - self, host: str = CALC_HOSTS[0] - ) -> list[FinancialFileInfo]: + def get_financial_file_list(self, host: str = CALC_HOSTS[0]) -> pd.DataFrame: """获取可用的历史专业财报文件列表。 连接到计算服务器,下载 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] + return _to_df([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 文件。 @@ -546,9 +542,7 @@ class TdxClient: """ return self._download_from_host(host, filename) - def get_financial_records( - self, filename: str, host: str = CALC_HOSTS[0] - ) -> list[FinancialRecord]: + def get_financial_records(self, filename: str, host: str = CALC_HOSTS[0]) -> pd.DataFrame: """下载财报 zip 并解析为每只股票的记录列表。 Args: @@ -560,12 +554,12 @@ class TdxClient: zip_data = self.get_financial_file(filename, host) if not zip_data: - return [] + return pd.DataFrame() 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 [] + return pd.DataFrame() dat_data = zf.read(dat_names[0]) m = re.search(r"(\d{8})", filename) @@ -576,13 +570,11 @@ class TdxClient: 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 - ) + FinancialRecord(code=code, market=market, report_date=rdate, fields=fields) ) - return records + return _to_df(records) - def get_market_stat(self) -> MarketStat: + def get_market_stat(self) -> pd.DataFrame: """获取 A 股全市场涨跌统计概况(基于 880005 行情统计)。 注意: @@ -590,9 +582,11 @@ class TdxClient: 用于保证计数守恒,不应视为协议已明确验证的停牌字段。 """ # 通达信中 880005 是全市场行情统计,880001 是总市值指数,880006 是涨跌停统计 - quotes = self.get_security_quotes([ - (Market.SH, "880005"), (Market.SH, "880001"), (Market.SH, "880006"), - ]) + quotes = self._execute( + GetSecurityQuotesCmd( + [(Market.SH, "880005"), (Market.SH, "880001"), (Market.SH, "880006")] + ) + ) if not quotes: raise RuntimeError("无法获取市场统计数据") q = quotes[0] @@ -603,17 +597,19 @@ class TdxClient: market_cap = quotes[1].price * 1e10 if len(quotes) > 1 else 0.0 limit_down = int(quotes[2].open) if len(quotes) > 2 else 0 limit_up = int(quotes[2].price) if len(quotes) > 2 else 0 - return MarketStat( - up_count=up, - down_count=down, - neutral_count=neutral, - suspended_count=max(0, total - up - down - neutral), - total_count=total, - total_amount=q.amount, - total_volume=q.vol, - total_market_cap=market_cap, - limit_up_count=limit_up, - limit_down_count=limit_down, + return _to_df( + MarketStat( + up_count=up, + down_count=down, + neutral_count=neutral, + suspended_count=max(0, total - up - down - neutral), + total_count=total, + total_amount=q.amount, + total_volume=q.vol, + total_market_cap=market_cap, + limit_up_count=limit_up, + limit_down_count=limit_down, + ) ) def _collect_transaction_records( @@ -659,41 +655,43 @@ class TdxClient: return all_recs - def get_fund_flow(self, market: Market, code: str) -> FundFlow: + def get_fund_flow(self, market: Market, code: str) -> pd.DataFrame: """获取个股当日资金流向分布(基于 L1 逐笔数据统计)。""" records = self._collect_transaction_records( - lambda start, page_size: self.get_transaction_data(market, code, start, page_size), + lambda start, page_size: self._execute( + GetTransactionDataCmd(market, code, start, page_size) + ), 2000, ) - return _classify_fund_flow(records) + return _to_df(_classify_fund_flow(records)) def get_history_fund_flow( self, market: Market, code: str, start: int, count: int - ) -> list[HistoricalFundFlow]: + ) -> pd.DataFrame: """获取个股历史日线资金流向序列。 优先走 Category 22 直连接口;若服务器返回空列表,则自动回退为 - “日 K 线取日期 + 历史逐笔成交重算资金流”的兼容实现。 + "日 K 线取日期 + 历史逐笔成交重算资金流"的兼容实现。 """ try: direct = self._execute(GetHistoryFundFlowCmd(market, code, start, count)) except Exception: direct = [] if direct: - return direct + return _to_df(direct) - bars = self.get_security_bars(market, code, KlineCategory.DAY, start, count) + bars = self._execute(GetSecurityBarsCmd(market, code, KlineCategory.DAY, start, count)) results: list[HistoricalFundFlow] = [] for bar in bars: date = _date_from_bar(bar) records = self._collect_transaction_records( - lambda page_start, page_size: self.get_history_transaction_data( - market, code, date, page_start, page_size + lambda page_start, page_size: self._execute( + GetHistoryTransactionDataCmd(market, code, date, page_start, page_size) ), 800, ) results.append(_historical_fund_flow_from_records(date, records)) - return results + return _to_df(results) # ============================================================ @@ -822,10 +820,10 @@ class AsyncTdxClient: async def get_security_count(self, market: Market) -> int: return await self._execute(GetSecurityCountCmd(market)) - async def get_security_list(self, market: Market, start: int) -> list[SecurityInfo]: - return await self._execute(GetSecurityListCmd(market, start)) + async def get_security_list(self, market: Market, start: int) -> pd.DataFrame: + return _to_df(await self._execute(GetSecurityListCmd(market, start))) - async def get_security_list_all(self, pages: int | str = "all") -> list[SecurityInfo]: + async def get_security_list_all(self, pages: int | str = "all") -> pd.DataFrame: """获取沪深 A 股完整证券列表,并自动挂载行业信息。 Args: @@ -842,7 +840,7 @@ class AsyncTdxClient: cached = _load_cache() if cached is not None: log.info("从缓存加载沪深 A 股列表,共 %d 只", len(cached)) - return cached + return _to_df(cached) def _max_start(count: int) -> int: if pages == "all": @@ -865,15 +863,18 @@ class AsyncTdxClient: total_pages = (limit + 999) // 1000 for page_idx, start in enumerate(range(0, limit, 1000)): try: - stocks = await self.get_security_list(market, start) + stocks = await self._execute(GetSecurityListCmd(market, start)) except Exception: - log.warning("%s 第 %d/%d 页获取失败,跳过", market.name, page_idx + 1, total_pages) + log.warning( + "%s 第 %d/%d 页获取失败,跳过", market.name, page_idx + 1, total_pages + ) continue - log.info("%s 第 %d/%d 页: %d 条", market.name, page_idx + 1, total_pages, len(stocks)) + log.info( + "%s 第 %d/%d 页: %d 条", market.name, page_idx + 1, total_pages, len(stocks) + ) for s in stocks: - is_a_share = ( - (market == Market.SH and s.code.startswith(("60", "68"))) - or (market == Market.SZ and s.code.startswith(("00", "30"))) + is_a_share = (market == Market.SH and s.code.startswith(("60", "68"))) or ( + market == Market.SZ and s.code.startswith(("00", "30")) ) if is_a_share: if s.code in industry_map: @@ -883,12 +884,10 @@ class AsyncTdxClient: log.info("沪深 A 股总数: %d", len(all_stocks)) if pages == "all": _save_cache(all_stocks) - return all_stocks + return _to_df(all_stocks) - async def get_security_quotes( - self, stocks: list[tuple[Market, str]] - ) -> list[SecurityQuote]: - return await self._execute(GetSecurityQuotesCmd(stocks)) + async def get_security_quotes(self, stocks: list[tuple[Market, str]]) -> pd.DataFrame: + return _to_df(await self._execute(GetSecurityQuotesCmd(stocks))) async def get_price_limits( self, market: Market, code: str, name: str, pre_close: float @@ -898,8 +897,8 @@ class AsyncTdxClient: no_limit_window_days = get_no_limit_window_days(market, code, name) if no_limit_window_days > 0: try: - bars = await self.get_security_bars( - market, code, KlineCategory.DAY, 0, no_limit_window_days + 1 + bars = await self._execute( + GetSecurityBarsCmd(market, code, KlineCategory.DAY, 0, no_limit_window_days + 1) ) listed_days = len(bars) except Exception: @@ -920,10 +919,9 @@ class AsyncTdxClient: category: KlineCategory, start: int, count: int = 800, - ) -> list[SecurityBar]: - return await self._execute( - GetSecurityBarsCmd(market, code, category, start, count) - ) + ) -> pd.DataFrame: + df = _to_df(await self._execute(GetSecurityBarsCmd(market, code, category, start, count))) + return _merge_bar_datetime(df, category in _DAILY_PLUS) async def get_index_bars( self, @@ -932,55 +930,50 @@ class AsyncTdxClient: category: KlineCategory, start: int, count: int = 800, - ) -> list[SecurityBar]: - return await self._execute(GetIndexBarsCmd(market, code, category, start, count)) + ) -> pd.DataFrame: + df = _to_df(await self._execute(GetIndexBarsCmd(market, code, category, start, count))) + return _merge_bar_datetime(df, category in _DAILY_PLUS) - async def get_minute_time_data(self, market: Market, code: str) -> list[MinuteBar]: + async def get_minute_time_data(self, market: Market, code: str) -> pd.DataFrame: today = _today_in_shanghai() - try: - bars = await self.get_history_minute_time_data(market, code, today) - if bars: - return bars - except Exception: - pass - return await self._execute(GetMinuteTimeDataCmd(market, code)) + bars = await self._execute(GetHistoryMinuteTimeDataCmd(market, code, today)) + return _add_minute_datetime(_to_df(bars), today) async def get_history_minute_time_data( self, market: Market, code: str, date: int - ) -> list[MinuteBar]: - return await self._execute(GetHistoryMinuteTimeDataCmd(market, code, date)) + ) -> pd.DataFrame: + bars = await self._execute(GetHistoryMinuteTimeDataCmd(market, code, date)) + return _add_minute_datetime(_to_df(bars), date) async def get_transaction_data( self, market: Market, code: str, start: int, count: int = 800 - ) -> list[TransactionRecord]: - return await self._execute(GetTransactionDataCmd(market, code, start, count)) + ) -> pd.DataFrame: + df = _to_df(await self._execute(GetTransactionDataCmd(market, code, start, count))) + return _merge_txn_datetime(df, _today_in_shanghai()) async def get_history_transaction_data( self, market: Market, code: str, date: int, start: int, count: int = 800 - ) -> list[TransactionRecord]: - return await self._execute( - GetHistoryTransactionDataCmd(market, code, date, start, count) + ) -> pd.DataFrame: + df = _to_df( + await self._execute(GetHistoryTransactionDataCmd(market, code, date, start, count)) ) + return _merge_txn_datetime(df, date) - async def get_xdxr_info(self, market: Market, code: str) -> list[XdxrRecord]: - return await self._execute(GetXdxrInfoCmd(market, code)) + async def get_xdxr_info(self, market: Market, code: str) -> pd.DataFrame: + return _to_df(await self._execute(GetXdxrInfoCmd(market, code))) - async def get_finance_info(self, market: Market, code: str) -> FinanceInfo: - return await self._execute(GetFinanceInfoCmd(market, code)) + async def get_finance_info(self, market: Market, code: str) -> pd.DataFrame: + return _to_df(await self._execute(GetFinanceInfoCmd(market, code))) - async def get_company_info_category( - self, market: Market, code: str - ) -> list[CompanyInfoCategory]: - return await self._execute(GetCompanyInfoCategoryCmd(market, code)) + async def get_company_info_category(self, market: Market, code: str) -> pd.DataFrame: + return _to_df(await self._execute(GetCompanyInfoCategoryCmd(market, code))) async def get_company_info_content( self, market: Market, code: str, filename: str, offset: int, length: int ) -> str: - return await self._execute( - GetCompanyInfoContentCmd(market, code, filename, offset, length) - ) + return await self._execute(GetCompanyInfoContentCmd(market, code, filename, offset, length)) - async def get_block_info(self, filename: str) -> list[TdxBlock]: + async def get_block_info(self, filename: str) -> pd.DataFrame: """获取并解析板块文件(行业、概念、风格等)。""" size, _hash = await self._execute(GetBlockInfoMetaCmd(filename)) full_data = bytearray() @@ -992,7 +985,7 @@ class AsyncTdxClient: break full_data.extend(chunk) pos += len(chunk) - return parse_block_dat(bytes(full_data), filename) + return _to_df(parse_block_dat(bytes(full_data), filename)) async def get_report_file(self, filename: str) -> bytes: """从服务器拉取大文件。""" @@ -1032,23 +1025,17 @@ class AsyncTdxClient: finally: await conn.close() - async def get_financial_file_list( - self, host: str = CALC_HOSTS[0] - ) -> list[FinancialFileInfo]: + async def get_financial_file_list(self, host: str = CALC_HOSTS[0]) -> pd.DataFrame: """获取可用的历史专业财报文件列表(异步)。""" 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] + return _to_df([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: + 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]: + async def get_financial_records(self, filename: str, host: str = CALC_HOSTS[0]) -> pd.DataFrame: """下载财报 zip 并解析为记录列表(异步)。""" import io import re @@ -1056,12 +1043,12 @@ class AsyncTdxClient: zip_data = await self.get_financial_file(filename, host) if not zip_data: - return [] + return pd.DataFrame() 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 [] + return pd.DataFrame() dat_data = zf.read(dat_names[0]) m = re.search(r"(\d{8})", filename) @@ -1072,13 +1059,11 @@ class AsyncTdxClient: 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 - ) + FinancialRecord(code=code, market=market, report_date=rdate, fields=fields) ) - return records + return _to_df(records) - async def get_market_stat(self) -> MarketStat: + async def get_market_stat(self) -> pd.DataFrame: """获取 A 股全市场涨跌统计概况(基于 880005 行情统计)。 注意: @@ -1086,9 +1071,11 @@ class AsyncTdxClient: 用于保证计数守恒,不应视为协议已明确验证的停牌字段。 """ # 通达信中 880005 是全市场行情统计,880001 是总市值指数,880006 是涨跌停统计 - quotes = await self.get_security_quotes([ - (Market.SH, "880005"), (Market.SH, "880001"), (Market.SH, "880006"), - ]) + quotes = await self._execute( + GetSecurityQuotesCmd( + [(Market.SH, "880005"), (Market.SH, "880001"), (Market.SH, "880006")] + ) + ) if not quotes: raise RuntimeError("无法获取市场统计数据") q = quotes[0] @@ -1099,17 +1086,19 @@ class AsyncTdxClient: market_cap = quotes[1].price * 1e10 if len(quotes) > 1 else 0.0 limit_down = int(quotes[2].open) if len(quotes) > 2 else 0 limit_up = int(quotes[2].price) if len(quotes) > 2 else 0 - return MarketStat( - up_count=up, - down_count=down, - neutral_count=neutral, - suspended_count=max(0, total - up - down - neutral), - total_count=total, - total_amount=q.amount, - total_volume=q.vol, - total_market_cap=market_cap, - limit_up_count=limit_up, - limit_down_count=limit_down, + return _to_df( + MarketStat( + up_count=up, + down_count=down, + neutral_count=neutral, + suspended_count=max(0, total - up - down - neutral), + total_count=total, + total_amount=q.amount, + total_volume=q.vol, + total_market_cap=market_cap, + limit_up_count=limit_up, + limit_down_count=limit_down, + ) ) async def _collect_transaction_records( @@ -1155,40 +1144,42 @@ class AsyncTdxClient: return all_recs - async def get_fund_flow(self, market: Market, code: str) -> FundFlow: + async def get_fund_flow(self, market: Market, code: str) -> pd.DataFrame: """获取个股当日资金流向分布(基于 L1 逐笔数据统计)。""" records = await self._collect_transaction_records( - lambda start, page_size: self.get_transaction_data( - market, code, start, page_size + lambda start, page_size: self._execute( + GetTransactionDataCmd(market, code, start, page_size) ), 2000, ) - return _classify_fund_flow(records) + return _to_df(_classify_fund_flow(records)) async def get_history_fund_flow( self, market: Market, code: str, start: int, count: int - ) -> list[HistoricalFundFlow]: + ) -> pd.DataFrame: """获取个股历史日线资金流向序列。 优先走 Category 22 直连接口;若服务器返回空列表,则自动回退为 - “日 K 线取日期 + 历史逐笔成交重算资金流”的兼容实现。 + "日 K 线取日期 + 历史逐笔成交重算资金流"的兼容实现。 """ try: direct = await self._execute(GetHistoryFundFlowCmd(market, code, start, count)) except Exception: direct = [] if direct: - return direct + return _to_df(direct) - bars = await self.get_security_bars(market, code, KlineCategory.DAY, start, count) + bars = await self._execute( + GetSecurityBarsCmd(market, code, KlineCategory.DAY, start, count) + ) results: list[HistoricalFundFlow] = [] for bar in bars: date = _date_from_bar(bar) records = await self._collect_transaction_records( - lambda page_start, page_size: self.get_history_transaction_data( - market, code, date, page_start, page_size + lambda page_start, page_size: self._execute( + GetHistoryTransactionDataCmd(market, code, date, page_start, page_size) ), 800, ) results.append(_historical_fund_flow_from_records(date, records)) - return results + return _to_df(results) diff --git a/src/easy_tdx/commands/minute_time.py b/src/easy_tdx/commands/minute_time.py index 4265d47..ba7d711 100644 --- a/src/easy_tdx/commands/minute_time.py +++ b/src/easy_tdx/commands/minute_time.py @@ -62,7 +62,7 @@ def _parse_minute_body(body: bytes, skip: int = 4) -> list[MinuteBar]: MinuteBar( price=last_price / 100.0, vol=vol, - unknown_1=unknown_1, + _unknown_1=unknown_1, _raw=body[record_start:pos], ) ) diff --git a/src/easy_tdx/models/timeseries.py b/src/easy_tdx/models/timeseries.py index 4b1db4e..e9c96d8 100644 --- a/src/easy_tdx/models/timeseries.py +++ b/src/easy_tdx/models/timeseries.py @@ -10,11 +10,11 @@ class MinuteBar: unknown_1: 协议中第二个变长整数,含义未明(疑似均价的编码形式)。 """ - price: float # 价格 - vol: int # 成交量 + price: float # 价格 + vol: int # 成交量 # pytdx 中被完全丢弃的字段,保留以供分析 - unknown_1: int = field(default=0, repr=False) # 原 reversed1 + _unknown_1: int = field(default=0, repr=False) # 原 reversed1 _raw: bytes = field(default=b"", repr=False, compare=False) diff --git a/tests/fixtures/minute_time.json b/tests/fixtures/minute_time.json index 837d7d6..40963ca 100644 --- a/tests/fixtures/minute_time.json +++ b/tests/fixtures/minute_time.json @@ -3,6 +3,6 @@ "first": { "price": 0.01, "vol": 48, - "unknown_1": 54 + "_unknown_1": 54 } } \ No newline at end of file diff --git a/tests/unit/test_a_share_extensions.py b/tests/unit/test_a_share_extensions.py index d4540ce..b42ae03 100644 --- a/tests/unit/test_a_share_extensions.py +++ b/tests/unit/test_a_share_extensions.py @@ -4,132 +4,175 @@ import asyncio import struct from unittest.mock import patch +import pandas as pd + from easy_tdx import AsyncTdxClient, Market, TdxClient from easy_tdx.client import _classify_fund_flow +from easy_tdx.commands.minute_time import ( + GetHistoryMinuteTimeDataCmd, +) +from easy_tdx.commands.security_bars import GetSecurityBarsCmd +from easy_tdx.commands.security_list import GetSecurityListCmd +from easy_tdx.commands.security_quotes import GetSecurityQuotesCmd +from easy_tdx.commands.transaction import ( + GetHistoryTransactionDataCmd, + GetTransactionDataCmd, +) from easy_tdx.models.bar import SecurityBar from easy_tdx.models.quote import SecurityQuote from easy_tdx.models.security import SecurityInfo -from easy_tdx.models.stats import HistoricalFundFlow -from easy_tdx.models.timeseries import MinuteBar -from easy_tdx.models.timeseries import TransactionRecord +from easy_tdx.models.timeseries import MinuteBar, TransactionRecord @patch("easy_tdx.client.TdxConnection") def test_get_fund_flow_logic(_mock_conn_cls): """测试资金流分类计算逻辑。""" client = TdxClient("127.0.0.1") - - # 构造模拟 Tick 数据 + mock_recs = [ - TransactionRecord(10, 0, 100.0, 101, 0, 0), # super_in (100*101*100 = 101w) - TransactionRecord(10, 1, 10.0, 250, 1, 0), # large_out (10*250*100 = 25w) - TransactionRecord(10, 2, 10.0, 10, 0, 0), # small_in (10*10*100 = 1w) + TransactionRecord(10, 0, 100.0, 101, 0, 0), # super_in + TransactionRecord(10, 1, 10.0, 250, 1, 0), # large_out + TransactionRecord(10, 2, 10.0, 10, 0, 0), # small_in ] - - with patch.object(TdxClient, "get_transaction_data", return_value=mock_recs): + + def mock_execute(cmd): + if isinstance(cmd, GetTransactionDataCmd): + return mock_recs + return [] + + with patch.object(TdxClient, "_execute", side_effect=mock_execute): flow = client.get_fund_flow(Market.SH, "600000") - - assert flow.super_in == 1010000.0 - assert flow.large_out == 250000.0 - assert flow.small_in == 10000.0 - assert flow.main_net_inflow == 1010000.0 - 250000.0 + assert isinstance(flow, pd.DataFrame) + assert flow["super_in"].iloc[0] == 1010000.0 + assert flow["large_out"].iloc[0] == 250000.0 + assert flow["small_in"].iloc[0] == 10000.0 def test_classify_fund_flow_exact_thresholds_use_lower_bucket(): """恰好命中阈值时,应落入较低一档。""" - flow = _classify_fund_flow([ - TransactionRecord(10, 0, 100.0, 100, 0, 0), # 100w -> large - TransactionRecord(10, 1, 100.0, 20, 0, 0), # 20w -> medium - TransactionRecord(10, 2, 100.0, 4, 0, 0), # 4w -> small - ]) + flow = _classify_fund_flow( + [ + TransactionRecord(10, 0, 100.0, 100, 0, 0), # 100w -> large + TransactionRecord(10, 1, 100.0, 20, 0, 0), # 20w -> medium + TransactionRecord(10, 2, 100.0, 4, 0, 0), # 4w -> small + ] + ) assert flow.super_in == 0.0 assert flow.large_in == 1000000.0 assert flow.medium_in == 200000.0 assert flow.small_in == 40000.0 + @patch("easy_tdx.client.TdxConnection") def test_get_security_list_all_filtering(_mock_conn_cls): """测试三市 A 股过滤与行业挂载逻辑。""" client = TdxClient("127.0.0.1") - - # 模拟行业配置 tdxhy.cfg + industry_cfg = b"1|600000|T01|||X01\n0|000001|T02|||X02\n2|830000|T03|||X03" - - # 模拟各市场返回 - def mock_get_list(market, start): - if market == Market.SH: - return [ - SecurityInfo(Market.SH, "600000", "SH_A", 100, 2, 10.0), - SecurityInfo(Market.SH, "999999", "INDEX", 100, 2, 3000.0), # 应被过滤 - ] - if market == Market.SZ: - return [SecurityInfo(Market.SZ, "000001", "SZ_A", 100, 2, 10.0)] - if market == Market.BJ: - return [SecurityInfo(Market.BJ, "830000", "BJ_A", 100, 2, 10.0)] + + def mock_execute(cmd): + if isinstance(cmd, GetSecurityListCmd): + if cmd.market == Market.SH: + return [ + SecurityInfo(Market.SH, "600000", "SH_A", 100, 2, 10.0), + SecurityInfo(Market.SH, "999999", "INDEX", 100, 2, 3000.0), + ] + if cmd.market == Market.SZ: + return [SecurityInfo(Market.SZ, "000001", "SZ_A", 100, 2, 10.0)] + return [] return [] - with patch.object(TdxClient, "get_report_file", return_value=industry_cfg), \ - patch.object(TdxClient, "get_security_count", return_value=1), \ - patch.object(TdxClient, "get_security_list", side_effect=mock_get_list): - - all_stocks = client.get_security_list_all() + with ( + patch.object(TdxClient, "_execute", side_effect=mock_execute), + patch.object(TdxClient, "get_report_file", return_value=industry_cfg), + patch.object(TdxClient, "get_security_count", return_value=1), + ): + all_stocks = client.get_security_list_all(pages=1) - # 预期只有 SH 和 SZ,BJ 已在扫描中降级移除 + assert isinstance(all_stocks, pd.DataFrame) assert len(all_stocks) == 2 - codes = [s.code for s in all_stocks] + codes = all_stocks["code"].tolist() assert "600000" in codes assert "000001" in codes - assert "830000" not in codes - s0 = next(s for s in all_stocks if s.code == "600000") - assert s0.industry_tdx == "T01" + assert "830000" not in codes + row = all_stocks[all_stocks["code"] == "600000"].iloc[0] + assert row["industry_tdx"] == "T01" + @patch("easy_tdx.client.TdxConnection") def test_get_market_stat_mapping(_mock_conn_cls): """测试市场统计字段映射。""" client = TdxClient("127.0.0.1") - + mock_quote = SecurityQuote( - Market.SH, "880005", - price=3000.0, # up - pre_close=2000.0, # down - open=0, - high=5500.0, # total - low=500.0, # neutral (low=500 -> neutral_count=500) - vol=1000000.0, cur_vol=0, amount=50000000.0, - s_vol=0, b_vol=0, active1=0, active2=0, - bid1=0, bid_vol1=0, bid2=0, bid_vol2=0, bid3=0, bid_vol3=0, - bid4=0, bid_vol4=0, bid5=0, bid_vol5=0, - ask1=0, ask_vol1=0, ask2=0, ask_vol2=0, ask3=0, ask_vol3=0, - ask4=0, ask_vol4=0, ask5=0, ask_vol5=0, - rise_speed=0, limit_up=0, limit_down=0 + Market.SH, + "880005", + price=3000.0, # up = int(price) + pre_close=0, + open=2000.0, # down = int(open) + high=5500.0, # total = int(high) + low=500.0, # neutral = int(low) + vol=1000000.0, + cur_vol=0, + amount=50000000.0, + s_vol=0, + b_vol=0, + active1=0, + active2=0, + bid1=0, + bid_vol1=0, + bid2=0, + bid_vol2=0, + bid3=0, + bid_vol3=0, + bid4=0, + bid_vol4=0, + bid5=0, + bid_vol5=0, + ask1=0, + ask_vol1=0, + ask2=0, + ask_vol2=0, + ask3=0, + ask_vol3=0, + ask4=0, + ask_vol4=0, + ask5=0, + ask_vol5=0, + rise_speed=0, + limit_up=0, + limit_down=0, ) - - with patch.object(TdxClient, "get_security_quotes", return_value=[mock_quote]): + + def mock_execute(cmd): + if isinstance(cmd, GetSecurityQuotesCmd): + return [mock_quote] + return [] + + with patch.object(TdxClient, "_execute", side_effect=mock_execute): stat = client.get_market_stat() - assert stat.up_count == 3000 - assert stat.down_count == 2000 - assert stat.neutral_count == 500 - assert stat.total_count == 5500 + assert isinstance(stat, pd.DataFrame) + assert stat["up_count"].iloc[0] == 3000 + assert stat["down_count"].iloc[0] == 2000 + assert stat["neutral_count"].iloc[0] == 500 + assert stat["total_count"].iloc[0] == 5500 + def test_get_history_fund_flow_parsing(): """测试历史资金流序列解析逻辑。""" from easy_tdx.commands.fund_flow import GetHistoryFundFlowCmd - - # 模拟 Category 22 响应 (Header 9 + Count 2 + Body 36) + body = bytearray(9) - body.extend(struct.pack(" 0: + def mock_execute(cmd): + if isinstance(cmd, GetHistoryFundFlowCmd): return [] - return txn_map[date] + if isinstance(cmd, GetSecurityBarsCmd): + return bars + if isinstance(cmd, GetHistoryTransactionDataCmd): + if cmd.start > 0: + return [] + return txn_map.get(cmd.date, []) + return [] - with patch.object(TdxClient, "_execute", return_value=[]), patch.object( - TdxClient, "get_security_bars", return_value=bars - ), patch.object( - TdxClient, "get_history_transaction_data", side_effect=mock_history_txn - ): + with patch.object(TdxClient, "_execute", side_effect=mock_execute): flows = client.get_history_fund_flow(Market.SH, "600000", 0, 2) - assert flows == [ - HistoricalFundFlow( - year=2025, - month=1, - day=8, - super_in=1010000.0, - super_out=0.0, - large_in=0.0, - large_out=250000.0, - medium_in=0.0, - medium_out=0.0, - small_in=0.0, - small_out=0.0, - ), - HistoricalFundFlow( - year=2025, - month=1, - day=9, - super_in=0.0, - super_out=0.0, - large_in=0.0, - large_out=0.0, - medium_in=0.0, - medium_out=0.0, - small_in=10000.0, - small_out=0.0, - ), - ] + assert isinstance(flows, pd.DataFrame) + assert len(flows) == 2 + row0 = flows.iloc[0] + assert row0["super_in"] == 1010000.0 + assert row0["large_out"] == 250000.0 + row1 = flows.iloc[1] + assert row1["small_in"] == 10000.0 @patch("easy_tdx.client.TdxConnection") @@ -202,21 +228,23 @@ def test_get_price_limits_uses_listing_window(_mock_conn_cls): """client.get_price_limits 应结合日 K 条数判断上市初期限价窗口。""" client = TdxClient("127.0.0.1") - with patch.object( - TdxClient, - "get_security_bars", - return_value=[SecurityBar(0, 0, 0, 0, 0, 0, 2025, 1, 1, 15, 0)] * 5, - ): + def mock_execute_5(cmd): + if isinstance(cmd, GetSecurityBarsCmd): + return [SecurityBar(0, 0, 0, 0, 0, 0, 2025, 1, 1, 15, 0)] * 5 + return [] + + with patch.object(TdxClient, "_execute", side_effect=mock_execute_5): assert client.get_price_limits(Market.SH, "600001", "主板新股", 10.0) == ( None, None, ) - with patch.object( - TdxClient, - "get_security_bars", - return_value=[SecurityBar(0, 0, 0, 0, 0, 0, 2025, 1, 1, 15, 0)] * 6, - ): + def mock_execute_6(cmd): + if isinstance(cmd, GetSecurityBarsCmd): + return [SecurityBar(0, 0, 0, 0, 0, 0, 2025, 1, 1, 15, 0)] * 6 + return [] + + with patch.object(TdxClient, "_execute", side_effect=mock_execute_6): assert client.get_price_limits(Market.SH, "600001", "主板老股", 10.0) == ( 11.0, 9.0, @@ -224,67 +252,50 @@ def test_get_price_limits_uses_listing_window(_mock_conn_cls): @patch("easy_tdx.client.TdxConnection") -def test_get_minute_time_data_prefers_history_endpoint(_mock_conn_cls): - """今日分时优先走历史分时接口,规避当前分时协议歧义。""" +def test_get_minute_time_data_uses_history_endpoint(_mock_conn_cls): + """今日分时走历史分时接口。""" client = TdxClient("127.0.0.1") expected = [MinuteBar(price=9.7, vol=13694)] - with patch("easy_tdx.client._today_in_shanghai", return_value=20260422), patch.object( - TdxClient, - "get_history_minute_time_data", - return_value=expected, - ) as mock_history, patch.object( - TdxClient, - "_execute", - side_effect=AssertionError("should not hit current-minute command"), + def mock_execute(cmd): + if isinstance(cmd, GetHistoryMinuteTimeDataCmd): + return expected + return [] + + with ( + patch("easy_tdx.client._today_in_shanghai", return_value=20260422), + patch.object(TdxClient, "_execute", side_effect=mock_execute) as mock_exec, ): result = client.get_minute_time_data(Market.SH, "600000") - mock_history.assert_called_once_with(Market.SH, "600000", 20260422) - assert result == expected + assert isinstance(result, pd.DataFrame) + assert result["price"].iloc[0] == 9.7 + history_calls = [ + c for c in mock_exec.call_args_list if isinstance(c[0][0], GetHistoryMinuteTimeDataCmd) + ] + assert len(history_calls) == 1 -@patch("easy_tdx.client.TdxConnection") -def test_get_minute_time_data_falls_back_to_current_endpoint(_mock_conn_cls): - """历史分时失败时,仍回退到原今日分时命令。""" - client = TdxClient("127.0.0.1") - fallback = [MinuteBar(price=9.61, vol=10698)] - - with patch("easy_tdx.client._today_in_shanghai", return_value=20260422), patch.object( - TdxClient, - "get_history_minute_time_data", - side_effect=RuntimeError("history unavailable"), - ) as mock_history, patch.object( - TdxClient, - "_execute", - return_value=fallback, - ) as mock_execute: - result = client.get_minute_time_data(Market.SH, "600000") - - mock_history.assert_called_once_with(Market.SH, "600000", 20260422) - mock_execute.assert_called_once() - assert result == fallback - - -def test_async_get_minute_time_data_prefers_history_endpoint(): - """异步客户端应与同步客户端保持同一回退策略。""" +def test_async_get_minute_time_data_uses_history_endpoint(): + """异步客户端走历史分时接口。""" expected = [MinuteBar(price=9.7, vol=13694)] async def run_test() -> None: with patch("easy_tdx.client.AsyncTdxConnection"): client = AsyncTdxClient("127.0.0.1") - with patch("easy_tdx.client._today_in_shanghai", return_value=20260422), patch.object( - AsyncTdxClient, - "get_history_minute_time_data", - return_value=expected, - ) as mock_history, patch.object( - AsyncTdxClient, - "_execute", - side_effect=AssertionError("should not hit current-minute command"), + + async def mock_execute(cmd): + if isinstance(cmd, GetHistoryMinuteTimeDataCmd): + return expected + return [] + + with ( + patch("easy_tdx.client._today_in_shanghai", return_value=20260422), + patch.object(AsyncTdxClient, "_execute", side_effect=mock_execute), ): result = await client.get_minute_time_data(Market.SH, "600000") - mock_history.assert_called_once_with(Market.SH, "600000", 20260422) - assert result == expected + assert isinstance(result, pd.DataFrame) + assert result["price"].iloc[0] == 9.7 asyncio.run(run_test()) diff --git a/tests/unit/test_block_info.py b/tests/unit/test_block_info.py index 23d7c26..f4786cd 100644 --- a/tests/unit/test_block_info.py +++ b/tests/unit/test_block_info.py @@ -4,6 +4,8 @@ import asyncio import struct from unittest.mock import patch +import pandas as pd + from easy_tdx.client import AsyncTdxClient, TdxClient from easy_tdx.codec.block import parse_block_dat from easy_tdx.models.finance import TdxBlock @@ -13,10 +15,11 @@ from easy_tdx.models.finance import TdxBlock def test_async_get_block_info_logic(mock_conn_cls): """测试 AsyncTdxClient.get_block_info 的异步拉取逻辑。""" mock_conn = mock_conn_cls.return_value - + # 模拟异步 execute async def mock_execute(cmd): from easy_tdx.commands.block_info import GetBlockInfoCmd, GetBlockInfoMetaCmd + if isinstance(cmd, GetBlockInfoMetaCmd): return 100, "hash" if isinstance(cmd, GetBlockInfoCmd): @@ -32,9 +35,9 @@ def test_async_get_block_info_logic(mock_conn_cls): with patch("easy_tdx.client.parse_block_dat") as mock_parse: mock_parse.return_value = [] res = await client.get_block_info("test.dat") - - assert isinstance(res, list) - assert mock_conn.execute.call_count == 2 # 1 meta + 1 data + + assert isinstance(res, pd.DataFrame) + assert mock_conn.execute.call_count == 2 # 1 meta + 1 data asyncio.run(main()) @@ -49,20 +52,20 @@ def test_parse_block_dat_basic(): # Header 384 + Count 2 + Record 2813 data = bytearray(384) data.extend(struct.pack(" bytes: # security_count # --------------------------------------------------------------------------- + def test_security_count_parse(): from easy_tdx.commands.security_count import GetSecurityCountCmd from easy_tdx.models.enums import Market @@ -38,6 +40,7 @@ def test_security_count_parse(): # security_list # --------------------------------------------------------------------------- + def test_security_list_parse(): from easy_tdx.commands.security_list import GetSecurityListCmd from easy_tdx.models.enums import Market @@ -61,18 +64,15 @@ def test_security_list_pre_close_uses_tdx_float_for_a_share(): from easy_tdx.commands.security_list import GetSecurityListCmd from easy_tdx.models.enums import Market - body = ( - struct.pack(" 0 # fixed values assert abs(b0.price - 0.01) < 0.001 assert b0.vol == 48 - assert b0.unknown_1 == 54 + assert b0._unknown_1 == 54 # --------------------------------------------------------------------------- # history_minute_time # --------------------------------------------------------------------------- + def test_history_minute_time_parse(): from easy_tdx.commands.minute_time import GetHistoryMinuteTimeDataCmd from easy_tdx.models.enums import Market @@ -202,7 +206,7 @@ def test_history_minute_time_parse(): b0 = bars[0] assert abs(b0.price - 10.29) < 0.01 assert b0.vol == 10044 - assert hasattr(b0, "unknown_1") + assert hasattr(b0, "_unknown_1") assert len(b0._raw) > 0 @@ -210,6 +214,7 @@ def test_history_minute_time_parse(): # transaction (current day) # --------------------------------------------------------------------------- + def test_transaction_parse(): from easy_tdx.commands.transaction import GetTransactionDataCmd from easy_tdx.models.enums import Market @@ -239,6 +244,7 @@ def test_transaction_parse(): # history_transaction # --------------------------------------------------------------------------- + def test_history_transaction_parse(): from easy_tdx.commands.transaction import GetHistoryTransactionDataCmd from easy_tdx.models.enums import Market @@ -266,6 +272,7 @@ def test_history_transaction_parse(): # xdxr_info # --------------------------------------------------------------------------- + def test_xdxr_info_parse(): from easy_tdx.commands.xdxr_info import GetXdxrInfoCmd from easy_tdx.models.enums import Market @@ -329,6 +336,7 @@ def test_xdxr_info_category_1_normalizes_per_10_share_fields(): # finance_info # --------------------------------------------------------------------------- + def test_finance_info_parse(): from easy_tdx.commands.finance_info import GetFinanceInfoCmd from easy_tdx.models.enums import Market @@ -354,6 +362,7 @@ def test_finance_info_parse(): # company_info_category # --------------------------------------------------------------------------- + def test_company_info_category_parse(): from easy_tdx.commands.company_info import GetCompanyInfoCategoryCmd from easy_tdx.models.enums import Market @@ -375,6 +384,7 @@ def test_company_info_category_parse(): # company_info_content # --------------------------------------------------------------------------- + def test_company_info_content_parse(): from easy_tdx.commands.company_info import GetCompanyInfoContentCmd from easy_tdx.models.enums import Market