feat: Enhance A-share stock name correction and index generation (#1462)

- Improved `fetch_tushare_stock_list.py` to support A-share name corrections for stocks with prefixes XD/XR/DR/N/C using the `--a-rk` option.
- Updated documentation to reflect the new functionality and usage of the correction script.
- Added a new script `refresh_stock_index.py` to automate the fetching of stock lists and index generation.
- Implemented tests for the new features, ensuring the correct functionality of stock name corrections and index generation.
- Enhanced error handling for missing dependencies and improved user feedback during execution.
This commit is contained in:
LouisHong
2026-05-25 20:26:24 +08:00
committed by GitHub
parent 639796f982
commit 43d05b4179
9 changed files with 859 additions and 266 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
<!-- 新条目格式:- [类型] 描述(类型取值:新功能/改进/修复/文档/测试/chore-->
<!-- 每条独立一行追加到本段末尾,无需分类标题,合并时冲突最小 -->
- [改进] `scripts/fetch_tushare_stock_list.py` 可对 A 股中带 `XD`/`XR`/`DR`/`N`/`C` 前缀的名称进行回填修正,供自动补全刷新流程默认使用。
- [修复] 股票自动补全索引生成缺少 `pypinyin` 时改为直接失败,避免写出缺失拼音字段的降级索引。
- [修复] 归一腾讯实时行情成交量为股口径,避免量能变化倍数被放大并误导分析报告。
- [改进] Web 路由页面改为按需加载,降低首包体积并增加路由加载失败恢复提示。
- [修复] Docker 默认部署移除 `.env` 单文件挂载,避免 WebUI 保存配置时因 `os.replace` 更新挂载点触发 `Device or resource busy`

View File

@@ -22,13 +22,19 @@ TUSHARE_TOKEN=你的tushare_token
python3 scripts/fetch_tushare_stock_list.py
```
如需针对 A 股名称状态做修正,可以加上 `--a-rk`,脚本会保持 `stock_basic` 作为基础来源,再用 `rt_k` 对带 `XD``XR``DR``N``C` 前缀的名称进行回填,并覆盖输出到 `data/stock_list_a.csv`
```bash
python3 scripts/fetch_tushare_stock_list.py --a-rk
```
### 3. 查看输出
数据将保存到 `data/` 目录:
```
data/
├── stock_list_a.csv # A股列表
├── stock_list_a.csv # A股列表--a-rk 时为修正后名称)
├── stock_list_hk.csv # 港股列表
├── stock_list_us.csv # 美股列表
└── README_stock_list.md # 数据说明文档
@@ -54,6 +60,8 @@ data/
### A股stock_list_a.csv
执行 `--a-rk` 时,这个文件会写入修正后的 A 股名称。
```csv
ts_code,symbol,name,area,industry,market,exchange,list_date,...
000001.SZ,000001,平安银行,深圳,银行,主板,SZSE,19910403,...
@@ -97,12 +105,26 @@ stock = a_stocks[a_stocks['ts_code'] == '600519.SH']
print(stock[['name', 'industry', 'list_date']])
```
### 更新自动补全索引
### 刷新股票自动补全索引
获取数据后,可以更新自动补全索引:
推荐直接使用一键刷新脚本,它会默认在抓取 A 股时使用 `--a-rk`,然后生成并同步自动补全索引:
```bash
pip install -r requirements.txt
python3 scripts/refresh_stock_index.py
```
生成自动补全索引依赖 `pypinyin` 写入中文股票的完整拼音和拼音首字母字段;缺少该依赖时脚本会直接失败,避免生成无法支持拼音搜索的降级索引。
如果你只想单独更新 CSV可以先抓取数据
```bash
python3 scripts/fetch_tushare_stock_list.py --a-rk
```
如果已经有新的 CSV只想重新生成索引
```bash
# 将 Tushare CSV 数据生成为前端自动补全索引
python3 scripts/generate_index_from_csv.py --test # 先测试
python3 scripts/generate_index_from_csv.py # 确认后生成
```

View File

@@ -7,6 +7,7 @@ Tushare 股票列表获取脚本
使用方法:
python3 scripts/fetch_tushare_stock_list.py
python3 scripts/fetch_tushare_stock_list.py --a-rk
环境要求:
- 需要在 .env 中配置 TUSHARE_TOKEN
@@ -16,19 +17,21 @@ Tushare 股票列表获取脚本
* 美股120积分试用5000积分正式权限
输出文件:
- data/stock_list_a.csv A股列表
- data/stock_list_a.csv A股列表--a-rk 时会覆盖为修正后名称)
- data/stock_list_hk.csv 港股列表
- data/stock_list_us.csv 美股列表
- data/README_stock_list.md 数据说明文档
"""
import argparse
import os
import sys
import time
import random
import re
from pathlib import Path
from datetime import datetime
from typing import Optional
from typing import Optional, List, Dict
import pandas as pd
from dotenv import load_dotenv
@@ -52,6 +55,9 @@ OUTPUT_DIR = Path(__file__).parent.parent / "data"
PAGE_SIZE = 5000 # 美股每页读取数量API 最大6000设置5000留余量
SLEEP_MIN = 5 # 最小睡眠时间(秒)
SLEEP_MAX = 10 # 最大睡眠时间(秒)
A_RK_BATCH_SIZE = 200
A_RK_FIELDS = "ts_code,name,close,pre_close,trade_time"
A_RK_NAME_PREFIX_RE = re.compile(r"^(XD|XR|DR|N|C)")
def get_tushare_api() -> Optional[ts.pro_api]:
@@ -131,6 +137,128 @@ def fetch_a_stock_list(api: ts.pro_api) -> Optional[pd.DataFrame]:
return None
def should_fix_a_stock_name(name: str) -> bool:
"""
判断 A 股名称是否属于需要修正的状态名。
主要覆盖新股、除权除息等前缀:
XD / XR / DR / N / C
"""
if name is None:
return False
text = str(name).strip()
if not text or text.lower() in {"nan", "none"}:
return False
return bool(A_RK_NAME_PREFIX_RE.match(text))
def chunk_list(items: List[str], chunk_size: int) -> List[List[str]]:
"""将列表按固定大小切片。"""
return [items[i:i + chunk_size] for i in range(0, len(items), chunk_size)]
def fetch_rt_k_names(api: ts.pro_api, ts_codes: List[str]) -> Dict[str, str]:
"""
批量获取 rt_k 返回的股票名称。
参考官方文档:
https://tushare.pro/wctapi/documents/372.md
rt_k 是 A 股实时日线接口,支持按股票代码和股票代码通配符提取
实时日 K 线行情。本脚本只把它用作名称回填的辅助来源,修正
stock_basic 中返回的短期交易状态前缀名称。
"""
if not ts_codes:
return {}
name_map: Dict[str, str] = {}
batches = chunk_list(ts_codes, A_RK_BATCH_SIZE)
print(f"\n[rt_k] 待修正股票数:{len(ts_codes)},分 {len(batches)} 批获取...")
for index, batch in enumerate(batches, start=1):
ts_code_param = ",".join(batch)
print(f" [rt_k] 第 {index}/{len(batches)} 批:{len(batch)} 只股票")
try:
df = api.rt_k(ts_code=ts_code_param, fields=A_RK_FIELDS)
except Exception as e:
print(f" [警告] rt_k 批次 {index} 获取失败: {e}")
continue
if df is None or len(df) == 0:
print(f" [警告] rt_k 批次 {index} 无返回数据")
continue
for _, row in df.iterrows():
code_value = row.get("ts_code", "")
name_value = row.get("name", "")
if pd.isna(code_value) or pd.isna(name_value):
continue
code = str(code_value).strip()
name = str(name_value).strip()
if code and name and code.lower() not in {"nan", "none"} and name.lower() not in {"nan", "none"}:
name_map[code] = name
if index < len(batches):
random_sleep(1, 2)
print(f"[rt_k] 成功获取 {len(name_map)} 条名称映射")
return name_map
def fix_a_stock_names_with_rt_k(api: ts.pro_api, df: pd.DataFrame) -> pd.DataFrame:
"""
使用 rt_k 修正 A 股名称。
仅对名称带有 XD / XR / DR / N / C 前缀的股票进行校正。
"""
if df is None or len(df) == 0:
return df
if "name" not in df.columns or "ts_code" not in df.columns:
print("[警告] A股数据缺少 ts_code/name 列,跳过 rt_k 名称修正")
return df
fix_mask = df["name"].astype(str).map(should_fix_a_stock_name)
fix_df = df.loc[fix_mask, ["ts_code", "name"]].copy()
if fix_df.empty:
print("[rt_k] 未发现需要修正的 A 股名称")
return df
ts_codes = fix_df["ts_code"].astype(str).tolist()
print(f"[rt_k] 发现 {len(ts_codes)} 只待修正 A 股:")
print(" " + ", ".join(ts_codes[:20]) + (" ..." if len(ts_codes) > 20 else ""))
name_map = fetch_rt_k_names(api, ts_codes)
if not name_map:
print("[警告] rt_k 未返回可用名称,保留原始 A 股名称")
return df
fixed_df = df.copy()
fixed_count = 0
for code, new_name in name_map.items():
if not new_name:
continue
match_index = fixed_df.index[fixed_df["ts_code"].astype(str) == code]
if len(match_index) == 0:
continue
old_name = str(fixed_df.loc[match_index[0], "name"])
if old_name != new_name:
fixed_df.loc[match_index[0], "name"] = new_name
fixed_count += 1
print(f"{code}: {old_name} -> {new_name}")
print(f"[rt_k] A 股名称修正完成,共修正 {fixed_count} 只股票")
return fixed_df
def fetch_hk_stock_list(api: ts.pro_api) -> Optional[pd.DataFrame]:
"""
获取港股列表
@@ -263,7 +391,9 @@ def save_to_csv(df: pd.DataFrame, filename: str, market_name: str) -> bool:
def generate_data_documentation(
a_df: Optional[pd.DataFrame],
hk_df: Optional[pd.DataFrame],
us_df: Optional[pd.DataFrame]
us_df: Optional[pd.DataFrame],
a_filename: str = "stock_list_a.csv",
a_title: str = "A股列表"
):
"""
生成数据说明文档
@@ -284,13 +414,13 @@ def generate_data_documentation(
| 文件 | 说明 | 记录数 |
|------|------|--------|
| `stock_list_a.csv` | A股列表 | {len(a_df) if a_df is not None else 0} |
| `{a_filename}` | {a_title} | {len(a_df) if a_df is not None else 0} |
| `stock_list_hk.csv` | 港股列表 | {len(hk_df) if hk_df is not None else 0} |
| `stock_list_us.csv` | 美股列表 | {len(us_df) if us_df is not None else 0} |
---
## A股数据stock_list_a.csv
## A股数据{a_filename}
### 数据接口
- **接口名称**`stock_basic`
@@ -402,7 +532,7 @@ BABA,阿里巴巴,Alibaba Group Holding Ltd.,ADR,20140919,
import pandas as pd
# 读取 A股数据
a_stocks = pd.read_csv('data/stock_list_a.csv')
a_stocks = pd.read_csv('data/{a_filename}')
# 读取港股数据
hk_stocks = pd.read_csv('data/stock_list_hk.csv')
@@ -455,11 +585,26 @@ us_stocks = pd.read_csv('data/stock_list_us.csv')
print(f"[错误] 生成说明文档失败: {e}")
def main():
def build_arg_parser() -> argparse.ArgumentParser:
"""构建命令行参数。"""
parser = argparse.ArgumentParser(description="Tushare 股票列表获取工具")
parser.add_argument(
"--a-rk",
action="store_true",
help="使用 rt_k 修正 A 股中带 XD/XR/DR/N/C 前缀的名称,并覆盖输出到 stock_list_a.csv",
)
return parser
def main(argv: Optional[List[str]] = None):
"""主函数"""
parser = build_arg_parser()
args = parser.parse_args(argv)
print("=" * 60)
print("Tushare 股票列表获取工具")
print("=" * 60)
print(f"[信息] A股名称修正模式{'开启' if args.a_rk else '关闭'}")
# 1. 获取 API 实例
api = get_tushare_api()
@@ -469,7 +614,15 @@ def main():
# 2. 获取 A股数据
a_df = fetch_a_stock_list(api)
if a_df is not None:
save_to_csv(a_df, 'stock_list_a.csv', 'A股')
a_filename = 'stock_list_a.csv'
a_title = 'A股列表'
a_market_name = 'A股'
if args.a_rk:
a_df = fix_a_stock_names_with_rt_k(api, a_df)
a_title = 'A股列表修正后'
save_to_csv(a_df, a_filename, a_market_name)
# 3. 获取港股数据
random_sleep() # 休息后再获取港股
@@ -485,7 +638,9 @@ def main():
# 5. 生成数据说明文档
print("\n正在生成数据说明文档...")
generate_data_documentation(a_df, hk_df, us_df)
a_filename = 'stock_list_a.csv'
a_title = 'A股列表修正后' if args.a_rk else 'A股列表'
generate_data_documentation(a_df, hk_df, us_df, a_filename=a_filename, a_title=a_title)
# 6. 总结
print("\n" + "=" * 60)

View File

@@ -31,9 +31,19 @@ try:
from pypinyin import lazy_pinyin, Style
PYPINYIN_AVAILABLE = True
except ImportError:
lazy_pinyin = None
Style = None
PYPINYIN_AVAILABLE = False
print("[Warning] pypinyin not available, pinyin fields will be empty")
print("[Info] Install with: pip install pypinyin")
def require_pypinyin() -> bool:
"""Ensure pypinyin is available before generating autocomplete assets."""
if PYPINYIN_AVAILABLE:
return True
print("[Error] pypinyin not available; cannot generate stock autocomplete index.")
print("[Info] Install dependencies with: pip install -r requirements.txt")
return False
def load_csv_data(csv_path: Path) -> List[Dict[str, Any]]:
@@ -166,6 +176,11 @@ def load_akshare_data(logs_dir: Path) -> List[Dict[str, Any]]:
Returns:
股票列表
说明:
AkShare 这条输入路径保留其原始 name 字段,不额外套用
Tushare A 股那套 XD / XR / DR 状态前缀修正逻辑。这里的目标是
复用 AkShare 已输出的展示名,而不是对其做二次归一化。
"""
csv_files = list(logs_dir.glob("stock_basic_*.csv"))
@@ -214,7 +229,7 @@ def generate_pinyin(name: str) -> tuple:
Tuple of (pinyin_full, pinyin_abbr)
"""
if not PYPINYIN_AVAILABLE:
return (None, None)
raise RuntimeError("pypinyin is required to generate stock autocomplete index")
try:
normalized_name = normalize_name_for_pinyin(name)
@@ -251,6 +266,23 @@ def normalize_name_for_pinyin(name: str) -> str:
return normalized.strip() or unicodedata.normalize('NFKC', name).strip()
def normalize_stock_name_for_index(name: str, market: str) -> str:
"""
Normalize stock names before writing the long-lived autocomplete index.
For A-shares (including BSE), ``XD``/``XR``/``DR`` are
ex-dividend/ex-rights trading-day prefixes. They should not be stored in
the official static index because they can become stale almost immediately.
New-stock prefixes such as ``N``/``C`` and risk-warning prefixes such as
``ST``/``*ST`` are preserved; they should be refreshed by the next
stock-list update.
"""
normalized = unicodedata.normalize('NFKC', str(name or '')).strip()
if market in {'CN', 'BSE'}:
normalized = re.sub(r'^(?:XD|XR|DR)\s*', '', normalized, flags=re.IGNORECASE)
return normalized.strip()
def extract_symbol_from_ts_code(ts_code: str, market: str) -> Optional[str]:
"""
从 ts_code 提取 displayCode
@@ -301,6 +333,7 @@ def get_stock_name(row: Dict[str, str], market: str) -> Optional[str]:
else:
# A股和港股使用中文名称
name = row.get('name', '').strip()
name = normalize_stock_name_for_index(name, market)
return name if name else None
@@ -568,6 +601,9 @@ def main():
print("=" * 60)
print(f"数据源:{args.source}")
if not require_pypinyin():
return 1
# 加载数据
print("\n[1/5] 读取 CSV 数据...")
if args.source == 'tushare':
@@ -586,11 +622,6 @@ def main():
print(f" 共读取 {len(stocks)} 只股票")
# 生成拼音提示
if not PYPINYIN_AVAILABLE:
print("\n[提示] 安装 pypinyin 可获得拼音搜索功能:")
print(" pip install pypinyin")
print("\n[2/5] 生成索引数据...")
index = build_stock_index(stocks)
@@ -618,7 +649,7 @@ def main():
for i, item in enumerate(compressed[:5]):
print(f" {i + 1}. {item}")
else:
print("\n[4/5] 写入文件:{output_path}")
print(f"\n[4/5] 写入文件:{output_path}")
with open(output_path, 'w', encoding='utf-8') as f:
f.write('[\n')
for i, item in enumerate(compressed):

View File

@@ -0,0 +1,97 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Refresh local stock autocomplete index assets.
Default flow:
1. Fetch Tushare stock lists into ``data/`` with ``--a-rk`` for A-share name correction.
2. Generate ``apps/dsa-web/public/stocks.index.json`` from CSV.
3. Copy the generated index to ``static/stocks.index.json`` for backend use.
"""
from __future__ import annotations
import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Sequence
REPO_ROOT = Path(__file__).resolve().parent.parent
WEB_INDEX_PATH = REPO_ROOT / "apps" / "dsa-web" / "public" / "stocks.index.json"
STATIC_INDEX_PATH = REPO_ROOT / "static" / "stocks.index.json"
def _run(command: Sequence[str]) -> None:
print(f"[refresh_stock_index] $ {' '.join(command)}", flush=True)
env = os.environ.copy()
env.setdefault("PYTHONUNBUFFERED", "1")
subprocess.run(command, cwd=REPO_ROOT, check=True, env=env)
def _has_tushare_token() -> bool:
env_path = REPO_ROOT / ".env"
try:
from dotenv import load_dotenv # type: ignore
except ImportError:
if env_path.is_file():
for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines():
key, sep, value = line.partition("=")
if sep and key.strip() == "TUSHARE_TOKEN" and value.strip().strip("'\""):
return True
return bool(os.getenv("TUSHARE_TOKEN", "").strip())
load_dotenv(env_path)
return bool(os.getenv("TUSHARE_TOKEN", "").strip())
def _sync_static_index() -> None:
if not WEB_INDEX_PATH.is_file():
raise FileNotFoundError(f"generated Web index not found: {WEB_INDEX_PATH}")
STATIC_INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
shutil.copyfile(WEB_INDEX_PATH, STATIC_INDEX_PATH)
print(f"[refresh_stock_index] synced {WEB_INDEX_PATH} -> {STATIC_INDEX_PATH}", flush=True)
def main(argv: Sequence[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="刷新股票自动补全索引")
parser.add_argument(
"--skip-fetch",
action="store_true",
help="跳过 Tushare 抓取,仅用现有 data/stock_list_*.csv 重新生成索引",
)
args = parser.parse_args(argv)
try:
if args.skip_fetch:
print("[refresh_stock_index] skip Tushare fetch; using existing CSV files")
else:
if not _has_tushare_token():
print(
"[refresh_stock_index] ERROR: missing TUSHARE_TOKEN. "
"Set it in .env or environment, or rerun with --skip-fetch.",
file=sys.stderr,
)
return 2
_run([sys.executable, "scripts/fetch_tushare_stock_list.py", "--a-rk"])
_run([sys.executable, "scripts/generate_index_from_csv.py", "--source", "tushare"])
_sync_static_index()
except subprocess.CalledProcessError as exc:
print(
f"[refresh_stock_index] ERROR: command failed with exit code {exc.returncode}",
file=sys.stderr,
)
return exc.returncode or 1
except (OSError, RuntimeError) as exc:
print(f"[refresh_stock_index] ERROR: {exc}", file=sys.stderr)
return 1
print("[refresh_stock_index] done")
return 0
if __name__ == "__main__":
sys.exit(main())

View File

@@ -0,0 +1,152 @@
# -*- coding: utf-8 -*-
"""Tests for scripts.fetch_tushare_stock_list A-share rt_k fix flow."""
from __future__ import annotations
import importlib
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock, patch
import pandas as pd
ROOT = Path(__file__).resolve().parent.parent
SCRIPTS_DIR = ROOT / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
sys.path.insert(0, str(SCRIPTS_DIR))
if "dotenv" not in sys.modules:
dotenv_stub = types.ModuleType("dotenv")
dotenv_stub.load_dotenv = lambda *args, **kwargs: None
sys.modules["dotenv"] = dotenv_stub
if "tushare" not in sys.modules:
tushare_stub = types.ModuleType("tushare")
tushare_stub.pro_api = lambda *args, **kwargs: MagicMock()
sys.modules["tushare"] = tushare_stub
fetch_tushare_stock_list = importlib.import_module("fetch_tushare_stock_list")
def test_should_fix_a_stock_name_matches_status_prefixes():
assert fetch_tushare_stock_list.should_fix_a_stock_name("XD西藏药")
assert fetch_tushare_stock_list.should_fix_a_stock_name("XR浦东建")
assert fetch_tushare_stock_list.should_fix_a_stock_name("DR罗曼股")
assert fetch_tushare_stock_list.should_fix_a_stock_name("N惠康")
assert fetch_tushare_stock_list.should_fix_a_stock_name("C天海")
assert not fetch_tushare_stock_list.should_fix_a_stock_name("平安银行")
assert not fetch_tushare_stock_list.should_fix_a_stock_name("ST罗顿")
assert not fetch_tushare_stock_list.should_fix_a_stock_name("*ST铖昌")
def test_fix_a_stock_names_with_rt_k_replaces_candidate_names():
api = MagicMock()
source_df = pd.DataFrame(
{
"ts_code": ["000001.SZ", "600848.SH", "300001.SZ"],
"name": ["XD西藏药", "平安银行", "N惠康"],
}
)
with patch.object(
fetch_tushare_stock_list,
"fetch_rt_k_names",
return_value={"000001.SZ": "西藏药", "300001.SZ": "惠康"},
) as fetch_rt_k_names:
fixed_df = fetch_tushare_stock_list.fix_a_stock_names_with_rt_k(api, source_df)
assert fixed_df.loc[fixed_df["ts_code"] == "000001.SZ", "name"].iloc[0] == "西藏药"
assert fixed_df.loc[fixed_df["ts_code"] == "600848.SH", "name"].iloc[0] == "平安银行"
assert fixed_df.loc[fixed_df["ts_code"] == "300001.SZ", "name"].iloc[0] == "惠康"
fetch_rt_k_names.assert_called_once_with(api, ["000001.SZ", "300001.SZ"])
def test_fetch_rt_k_names_batches_and_collects_results():
api = MagicMock()
api.rt_k.return_value = pd.DataFrame(
{
"ts_code": ["000001.SZ", "300001.SZ"],
"name": ["西藏药", "惠康"],
"close": [1.0, 2.0],
"pre_close": [1.0, 2.0],
"trade_time": ["10:00:00", "10:00:00"],
}
)
with patch.object(fetch_tushare_stock_list, "random_sleep") as random_sleep:
name_map = fetch_tushare_stock_list.fetch_rt_k_names(api, ["000001.SZ", "300001.SZ"])
assert name_map == {"000001.SZ": "西藏药", "300001.SZ": "惠康"}
api.rt_k.assert_called_once_with(
ts_code="000001.SZ,300001.SZ",
fields="ts_code,name,close,pre_close,trade_time",
)
random_sleep.assert_not_called()
def test_main_default_flow_keeps_original_filename():
api = MagicMock()
a_df = pd.DataFrame({"ts_code": ["000001.SZ"], "name": ["平安银行"]})
hk_df = pd.DataFrame({"ts_code": ["00001.HK"], "name": ["长和"]})
us_df = pd.DataFrame({"ts_code": ["AAPL"], "name": ["苹果"]})
with (
patch.object(fetch_tushare_stock_list, "get_tushare_api", return_value=api),
patch.object(fetch_tushare_stock_list, "fetch_a_stock_list", return_value=a_df) as fetch_a,
patch.object(fetch_tushare_stock_list, "save_to_csv") as save_to_csv,
patch.object(fetch_tushare_stock_list, "fetch_hk_stock_list", return_value=hk_df) as fetch_hk,
patch.object(fetch_tushare_stock_list, "fetch_us_stock_list", return_value=us_df) as fetch_us,
patch.object(fetch_tushare_stock_list, "generate_data_documentation") as generate_doc,
patch.object(fetch_tushare_stock_list, "random_sleep") as random_sleep,
patch.object(fetch_tushare_stock_list, "fix_a_stock_names_with_rt_k") as fix_a_stock_names,
):
exit_code = fetch_tushare_stock_list.main([])
assert exit_code == 0
fetch_a.assert_called_once_with(api)
save_to_csv.assert_any_call(a_df, "stock_list_a.csv", "A股")
fetch_hk.assert_called_once_with(api)
fetch_us.assert_called_once_with(api)
fix_a_stock_names.assert_not_called()
generate_doc.assert_called_once_with(a_df, hk_df, us_df, a_filename="stock_list_a.csv", a_title="A股列表")
assert random_sleep.call_count == 2
def test_main_a_rk_flow_overwrites_a_filename_and_rt_k():
api = MagicMock()
a_df = pd.DataFrame({"ts_code": ["000001.SZ"], "name": ["XD西藏药"]})
fixed_df = pd.DataFrame({"ts_code": ["000001.SZ"], "name": ["西藏药"]})
hk_df = pd.DataFrame({"ts_code": ["00001.HK"], "name": ["长和"]})
us_df = pd.DataFrame({"ts_code": ["AAPL"], "name": ["苹果"]})
with (
patch.object(fetch_tushare_stock_list, "get_tushare_api", return_value=api),
patch.object(fetch_tushare_stock_list, "fetch_a_stock_list", return_value=a_df) as fetch_a,
patch.object(fetch_tushare_stock_list, "fix_a_stock_names_with_rt_k", return_value=fixed_df) as fix_a_stock_names,
patch.object(fetch_tushare_stock_list, "save_to_csv") as save_to_csv,
patch.object(fetch_tushare_stock_list, "fetch_hk_stock_list", return_value=hk_df) as fetch_hk,
patch.object(fetch_tushare_stock_list, "fetch_us_stock_list", return_value=us_df) as fetch_us,
patch.object(fetch_tushare_stock_list, "generate_data_documentation") as generate_doc,
patch.object(fetch_tushare_stock_list, "random_sleep") as random_sleep,
):
exit_code = fetch_tushare_stock_list.main(["--a-rk"])
assert exit_code == 0
fetch_a.assert_called_once_with(api)
fix_a_stock_names.assert_called_once_with(api, a_df)
fixed_save_call = next(
call for call in save_to_csv.call_args_list if call.args[1] == "stock_list_a.csv"
)
pd.testing.assert_frame_equal(fixed_save_call.args[0], fixed_df)
assert fixed_save_call.args[2] == "A股"
fetch_hk.assert_called_once_with(api)
fetch_us.assert_called_once_with(api)
generate_doc.assert_called_once_with(
fixed_df,
hk_df,
us_df,
a_filename="stock_list_a.csv",
a_title="A股列表修正后",
)
assert random_sleep.call_count == 2

View File

@@ -22,7 +22,9 @@ from generate_index_from_csv import (
determine_market,
generate_aliases,
normalize_name_for_pinyin,
normalize_stock_name_for_index,
generate_pinyin,
main,
compress_index,
build_stock_index,
load_tushare_data,
@@ -140,6 +142,18 @@ class TestGetStockName:
result = get_stock_name(row, 'CN')
assert result is None
def test_cn_stock_name_strips_ex_rights_prefix(self):
"""测试 A股除权除息短期前缀不会写入长期索引名称"""
row = {'name': 'XD西藏药', 'enname': ''}
result = get_stock_name(row, 'CN')
assert result == '西藏药'
def test_cn_stock_name_preserves_new_stock_prefix(self):
"""测试 A股新股前缀保留等待后续数据包刷新自然消失"""
row = {'name': 'N惠康', 'enname': ''}
result = get_stock_name(row, 'CN')
assert result == 'N惠康'
class TestDataCleaning:
"""测试数据清洗逻辑"""
@@ -257,6 +271,26 @@ class TestDataCleaning:
assert get_us_delist_priority({'delist_date': '20250131'}) == 0
class TestNormalizeStockNameForIndex:
"""测试索引名称归一化"""
def test_strips_a_share_ex_rights_prefixes(self):
assert normalize_stock_name_for_index('XD西藏药', 'CN') == '西藏药'
assert normalize_stock_name_for_index('XR示例股', 'CN') == '示例股'
assert normalize_stock_name_for_index('DR罗曼股', 'CN') == '罗曼股'
assert normalize_stock_name_for_index('XD朱老六', 'BSE') == '朱老六'
def test_preserves_a_share_new_stock_and_st_prefixes(self):
assert normalize_stock_name_for_index('N惠康', 'CN') == 'N惠康'
assert normalize_stock_name_for_index('C天海', 'CN') == 'C天海'
assert normalize_stock_name_for_index('ST海王', 'CN') == 'ST海王'
assert normalize_stock_name_for_index('*ST美丽', 'CN') == '*ST美丽'
def test_does_not_strip_other_markets(self):
assert normalize_stock_name_for_index('DRAGONFLY ENERGY', 'US') == 'DRAGONFLY ENERGY'
assert normalize_stock_name_for_index('XD港股示例', 'HK') == 'XD港股示例'
class TestAliases:
"""测试别名生成函数"""
@@ -518,9 +552,24 @@ class TestPinyin:
def test_generate_pinyin(self):
"""测试拼音生成"""
# 注意:这个测试需要 pypinyin 可用
pinyin_full, pinyin_abbr = generate_pinyin('平安银行')
if pinyin_full:
assert isinstance(pinyin_full, str)
if pinyin_abbr:
assert isinstance(pinyin_abbr, str)
assert pinyin_full == 'pinganyinhang'
assert pinyin_abbr == 'payh'
def test_generate_pinyin_requires_dependency(self, monkeypatch):
"""测试缺少 pypinyin 时不会生成降级拼音字段"""
import generate_index_from_csv
monkeypatch.setattr(generate_index_from_csv, 'PYPINYIN_AVAILABLE', False)
with pytest.raises(RuntimeError, match='pypinyin is required'):
generate_index_from_csv.generate_pinyin('平安银行')
def test_main_fails_without_pypinyin(self, monkeypatch):
"""测试正式生成索引前必须具备 pypinyin"""
import generate_index_from_csv
monkeypatch.setattr(generate_index_from_csv, 'PYPINYIN_AVAILABLE', False)
monkeypatch.setattr(sys, 'argv', ['generate_index_from_csv.py'])
assert main() == 1

View File

@@ -0,0 +1,38 @@
# -*- coding: utf-8 -*-
"""Tests for scripts.refresh_stock_index default fetch behavior."""
from __future__ import annotations
import importlib
import sys
from pathlib import Path
from unittest.mock import patch
ROOT = Path(__file__).resolve().parent.parent
SCRIPTS_DIR = ROOT / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
sys.path.insert(0, str(SCRIPTS_DIR))
refresh_stock_index = importlib.import_module("refresh_stock_index")
def test_main_fetches_tushare_with_a_rk_by_default():
with (
patch.object(refresh_stock_index, "_has_tushare_token", return_value=True),
patch.object(refresh_stock_index, "_run") as run,
patch.object(refresh_stock_index, "_sync_static_index"),
):
exit_code = refresh_stock_index.main([])
assert exit_code == 0
assert run.call_args_list[0].args[0] == [
sys.executable,
"scripts/fetch_tushare_stock_list.py",
"--a-rk",
]
assert run.call_args_list[1].args[0] == [
sys.executable,
"scripts/generate_index_from_csv.py",
"--source",
"tushare",
]