mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 19:04:15 +08:00
feat(watchlist): batch import CSV / pasted codes into groups
自选页新增批量导入：支持 CSV/TXT 文件与粘贴证券代码两种来源，解析出 候选并在弹窗内确认后按 M:N 分组一次性写入（目标分组可多选或就地新建， 默认只勾新增标的，已在自选的可并入所选分组）。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
richard
@@ -5,6 +5,7 @@ import logging
|
||||
import math
|
||||
import time
|
||||
from datetime import date
|
||||
from typing import Callable
|
||||
|
||||
import anyio
|
||||
import polars as pl
|
||||
@@ -13,6 +14,7 @@ from pydantic import BaseModel
|
||||
|
||||
from app.db_safe import is_valid_ext_ident, quote_ident
|
||||
from app.services import watchlist
|
||||
from app.services.watchlist_csv import import_watchlist_codes, import_watchlist_csv
|
||||
from app.services.watchlist_ocr import import_watchlist_image
|
||||
from app.services.watchlist_ocr.provider import get_ocr_provider
|
||||
|
||||
@@ -31,6 +33,13 @@ _IMPORT_IMAGE_TYPES = {
|
||||
}
|
||||
# OCR 独立并发上限：避免多张大图同时解码 + 多 Tesseract 子进程
|
||||
_OCR_LIMITER = anyio.CapacityLimiter(2)
|
||||
# CSV/TXT 导入：文本远小于截图，上限 5MB 足够
|
||||
_MAX_IMPORT_CSV_BYTES = 5 * 1024 * 1024
|
||||
_IMPORT_CSV_TYPES = {
|
||||
"text/csv",
|
||||
"text/plain",
|
||||
"application/csv",
|
||||
}
|
||||
|
||||
|
||||
class AddRequest(BaseModel):
|
||||
@@ -43,6 +52,7 @@ class BatchAddRequest(BaseModel):
|
||||
symbols: list[str]
|
||||
note: str = ""
|
||||
group_id: str | None = None
|
||||
group_ids: list[str] | None = None
|
||||
|
||||
|
||||
class GroupNameRequest(BaseModel):
|
||||
@@ -58,6 +68,10 @@ class GroupAssignRequest(BaseModel):
|
||||
group_id: str | None = None
|
||||
|
||||
|
||||
class ImportCodesRequest(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
def _with_names(rows: list[dict], request: Request) -> list[dict]:
|
||||
if not rows:
|
||||
return rows
|
||||
@@ -89,7 +103,12 @@ def add_one(req: AddRequest, request: Request):
|
||||
@router.post("/batch")
|
||||
def add_batch(req: BatchAddRequest, request: Request):
|
||||
try:
|
||||
rows, added = watchlist.add_batch(req.symbols, req.note, req.group_id)
|
||||
rows, added = watchlist.add_batch(
|
||||
req.symbols,
|
||||
req.note,
|
||||
group_id=req.group_id,
|
||||
group_ids=req.group_ids,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(400, str(e)) from e
|
||||
return {"symbols": _with_names(rows, request), "added": added}
|
||||
@@ -194,6 +213,70 @@ async def import_from_image(request: Request, file: UploadFile = File(...)):
|
||||
return result
|
||||
|
||||
|
||||
def _run_candidate_import(parse: Callable[[], dict], empty_msg: str) -> dict:
|
||||
"""执行候选解析：ValueError→400、其他→500、空候选→400、剥离 raw_text。"""
|
||||
try:
|
||||
result = parse()
|
||||
except ValueError as e:
|
||||
raise HTTPException(400, str(e)) from e
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.exception("watchlist import failed")
|
||||
raise HTTPException(500, f"解析失败: {e}") from e
|
||||
if not result["candidates"]:
|
||||
raise HTTPException(400, empty_msg)
|
||||
result.pop("raw_text", None)
|
||||
return result
|
||||
|
||||
|
||||
@router.post("/import-csv")
|
||||
async def import_from_csv(request: Request, file: UploadFile = File(...)):
|
||||
"""从 CSV / TXT 导入自选候选列表（不自动写入自选）。
|
||||
|
||||
兼容同花顺/东财/通达信导出（逗号或 Tab 分隔、UTF-8 或 GBK 编码）。目标分组
|
||||
在候选确认时由前端传入 batch 接口，本端点只做解析与主数据校验。
|
||||
"""
|
||||
content_type = (file.content_type or "").split(";")[0].strip().lower()
|
||||
filename = (file.filename or "").lower()
|
||||
ok_type = content_type in _IMPORT_CSV_TYPES
|
||||
ok_ext = filename.endswith((".csv", ".txt"))
|
||||
if not ok_type and not ok_ext:
|
||||
raise HTTPException(400, "仅支持 CSV / TXT 文件")
|
||||
|
||||
data = await file.read()
|
||||
if not data:
|
||||
raise HTTPException(400, "空文件")
|
||||
if len(data) > _MAX_IMPORT_CSV_BYTES:
|
||||
raise HTTPException(400, "文件过大（上限 5MB）")
|
||||
|
||||
data_dir = request.app.state.repo.store.data_dir
|
||||
# 解码与自选/instruments parquet 读取为同步 CPU/IO，挪线程池避免卡事件循环
|
||||
return await anyio.to_thread.run_sync(
|
||||
lambda: _run_candidate_import(
|
||||
lambda: import_watchlist_csv(
|
||||
data,
|
||||
data_dir,
|
||||
existing_symbols={r["symbol"] for r in watchlist.list_symbols()},
|
||||
),
|
||||
"文件中未识别到股票代码或名称",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/import-codes")
|
||||
def import_from_codes(req: ImportCodesRequest, request: Request):
|
||||
"""从粘贴的证券代码导入自选候选列表（不自动写入自选）。"""
|
||||
text = req.text.strip()
|
||||
if not text:
|
||||
raise HTTPException(400, "请输入要导入的股票代码")
|
||||
|
||||
existing = {r["symbol"] for r in watchlist.list_symbols()}
|
||||
data_dir = request.app.state.repo.store.data_dir
|
||||
return _run_candidate_import(
|
||||
lambda: import_watchlist_codes(text, data_dir, existing_symbols=existing),
|
||||
"未识别到股票代码",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/{symbol}/top")
|
||||
def move_one_to_top(symbol: str, request: Request):
|
||||
rows = watchlist.move_to_top(symbol)
|
||||
|
||||
@@ -183,15 +183,23 @@ def add_batch(
|
||||
symbols: list[str],
|
||||
note: str = "",
|
||||
group_id: str | None = None,
|
||||
group_ids: list[str] | None = None,
|
||||
) -> tuple[list[dict], int]:
|
||||
"""批量添加并保持既有语义：每个新处理的标的移动到列表最前面。
|
||||
|
||||
group_id 为可选的初始分组(如从某分组页添加时); 重复添加的标的保留
|
||||
既有全部分组, 仅在显式传入 group_id 且尚未属于该组时并入。
|
||||
分组为可选的初始分组：``group_id`` 单组（如从某分组页添加）或 ``group_ids``
|
||||
多组（如批量导入同时并入多个分组）。重复添加的标的保留既有全部分组，
|
||||
仅把尚未属于的传入分组并入；二者可同时使用、内部去重。
|
||||
"""
|
||||
with _LOCK:
|
||||
groups = _read_groups()
|
||||
_validate_group_id(group_id, groups)
|
||||
# 合并单/多组参数并去重；逐组校验存在性
|
||||
apply_ids: list[str] = []
|
||||
for gid in (group_ids or []) + ([group_id] if group_id is not None else []):
|
||||
if gid in apply_ids:
|
||||
continue
|
||||
_validate_group_id(gid, groups)
|
||||
apply_ids.append(gid)
|
||||
rows = _read_entries().to_dicts()
|
||||
added = 0
|
||||
for symbol in symbols:
|
||||
@@ -200,8 +208,9 @@ def add_batch(
|
||||
added += 1
|
||||
rows = [row for row in rows if row["symbol"] != symbol]
|
||||
gids = list((existing or {}).get("group_ids") or [])
|
||||
if group_id is not None and group_id not in gids:
|
||||
gids.append(group_id)
|
||||
for gid in apply_ids:
|
||||
if gid not in gids:
|
||||
gids.append(gid)
|
||||
rows.insert(0, {
|
||||
"symbol": symbol,
|
||||
"added_at": datetime.utcnow().isoformat(timespec="seconds"),
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
"""自选股 CSV/TXT 与粘贴代码批量导入：解码 → 抽代码 → instruments 校验。
|
||||
|
||||
国内行情软件（同花顺/东财/通达信）导出的自选多为 CSV/TXT，且常为 GBK 系编码
|
||||
（参见 ext_data.ensure_utf8_csv 的说明）。本模块把上传字节 / 粘贴文本解析为与截图
|
||||
OCR 一致的候选结构，前端复用同一套勾选确认流程；写入目标统一走自选分组语义
|
||||
（watchlist.add_batch 的 group_ids M:N 并入），本模块不落盘、不建标签。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import io
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.services.watchlist_ocr.pipeline import (
|
||||
_CODE_RE,
|
||||
ImportCandidate,
|
||||
build_instrument_lookups,
|
||||
extract_codes,
|
||||
resolve_candidates,
|
||||
)
|
||||
|
||||
_CJK_RE = re.compile(f"[{chr(0x4E00)}-{chr(0x9FFF)}]") # CJK 统一表意文字块
|
||||
|
||||
# 编码回退链：UTF-8（含 BOM）→ GB18030（GB18030 是 GBK 超集，无需单独回退）
|
||||
_ENCODINGS = ("utf-8-sig", "gb18030")
|
||||
|
||||
|
||||
def _finalize(
|
||||
provider: str,
|
||||
text: str,
|
||||
codes: list[str],
|
||||
candidates: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""组装与截图 OCR 一致的候选响应，统一 matched/unmatched 计数口径。"""
|
||||
matched_count = sum(1 for c in candidates if c["matched"])
|
||||
return {
|
||||
"provider": provider,
|
||||
"raw_text": text,
|
||||
"codes": codes,
|
||||
"candidates": candidates,
|
||||
"matched_count": matched_count,
|
||||
"unmatched_count": len(candidates) - matched_count,
|
||||
}
|
||||
|
||||
|
||||
def decode_csv_bytes(raw: bytes) -> str:
|
||||
"""把上传字节解码为文本，兼容 UTF-8 / GBK 系编码。"""
|
||||
if not raw:
|
||||
raise ValueError("空文件")
|
||||
last_err: Exception | None = None
|
||||
for enc in _ENCODINGS:
|
||||
try:
|
||||
return raw.decode(enc)
|
||||
except (UnicodeDecodeError, LookupError) as e:
|
||||
last_err = e
|
||||
raise ValueError("无法识别文件编码，请另存为 UTF-8 或 GBK 后重试") from last_err
|
||||
|
||||
|
||||
def _is_code_cell(cell: str) -> bool:
|
||||
# 调用方（parse_csv_rows）已 strip 过单元格
|
||||
return bool(_CODE_RE.fullmatch(cell))
|
||||
|
||||
|
||||
def _pick_name(cells: list[str]) -> str | None:
|
||||
"""取行内首个含 ≥2 个汉字且非六位代码的单元格作为名称候选。"""
|
||||
for cell in cells:
|
||||
if not cell or _is_code_cell(cell):
|
||||
continue
|
||||
if len(_CJK_RE.findall(cell)) >= 2:
|
||||
return cell
|
||||
return None
|
||||
|
||||
|
||||
def parse_csv_rows(text: str) -> list[tuple[list[str], str | None]]:
|
||||
"""解析 CSV/TXT 文本为 [(行内代码列表, 名称候选), ...]。
|
||||
|
||||
- 自动识别逗号 / Tab 分隔（同花顺/通达信导出常见 Tab）。
|
||||
- 逐行取所有六位数字作为代码候选；无代码行保留给名称兜底（是否输出由
|
||||
import_watchlist_csv 决定：表头等名称命不中主数据的行会被忽略）。
|
||||
- 返回列表保持文件行序。
|
||||
"""
|
||||
if not text.strip():
|
||||
return []
|
||||
|
||||
first = next((ln for ln in text.splitlines() if ln.strip()), "")
|
||||
delimiter = "\t" if first.count("\t") > first.count(",") else ","
|
||||
reader = csv.reader(io.StringIO(text), delimiter=delimiter)
|
||||
|
||||
rows: list[tuple[list[str], str | None]] = []
|
||||
for raw_row in reader:
|
||||
cells = [c.strip() for c in raw_row if c is not None]
|
||||
if not cells:
|
||||
continue
|
||||
codes: list[str] = []
|
||||
for cell in cells:
|
||||
codes.extend(m.group(1) for m in _CODE_RE.finditer(cell))
|
||||
rows.append((codes, _pick_name(cells)))
|
||||
return rows
|
||||
|
||||
|
||||
def import_watchlist_csv(
|
||||
raw: bytes,
|
||||
data_dir: Path,
|
||||
*,
|
||||
existing_symbols: set[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""解析 CSV/TXT 字节并返回候选列表（不写入自选）。返回结构与 OCR 一致。"""
|
||||
text = decode_csv_bytes(raw)
|
||||
return _resolve_rows(text, data_dir, existing_symbols=existing_symbols)
|
||||
|
||||
|
||||
def import_watchlist_codes(
|
||||
text: str,
|
||||
data_dir: Path,
|
||||
*,
|
||||
existing_symbols: set[str] | None = None,
|
||||
max_codes: int = 1000,
|
||||
) -> dict[str, Any]:
|
||||
"""解析粘贴的证券代码并返回候选列表（不写入自选）。
|
||||
|
||||
与 CSV 行级解析不同：粘贴文本里的多个代码可能挤在同一行/同一段（逗号、空格、
|
||||
换行分隔），必须按 ``extract_codes`` 全量抽码、去重保序，逐码生成候选，
|
||||
否则会把同行多码压成单候选而静默丢码。仅与 CSV 路径共享 lookups/resolve。
|
||||
"""
|
||||
codes = extract_codes(text)
|
||||
if not codes:
|
||||
return _finalize("codes", text, [], [])
|
||||
if len(codes) > max_codes:
|
||||
raise ValueError(f"一次最多导入 {max_codes} 个股票代码，已识别 {len(codes)} 个")
|
||||
|
||||
code_to_symbol, symbol_to_name = build_instrument_lookups(data_dir)
|
||||
candidates = resolve_candidates(codes, code_to_symbol, symbol_to_name, existing_symbols)
|
||||
return _finalize("codes", text, codes, [c.to_dict() for c in candidates])
|
||||
|
||||
|
||||
def _resolve_rows(
|
||||
text: str,
|
||||
data_dir: Path,
|
||||
*,
|
||||
existing_symbols: set[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""逐行把文本解析为候选（CSV/TXT 用；行内多码取首个已匹配者）。"""
|
||||
rows = parse_csv_rows(text)
|
||||
|
||||
code_to_symbol, symbol_to_name = build_instrument_lookups(data_dir)
|
||||
# 名称兜底反向表：CSV 可能只有名称列（无代码）
|
||||
name_to_symbol: dict[str, str] = {}
|
||||
for symbol, name in symbol_to_name.items():
|
||||
name_to_symbol.setdefault(name, symbol)
|
||||
|
||||
existing = existing_symbols or set()
|
||||
# 全部唯一代码按出现顺序一次性构建候选（复用 OCR 的构造逻辑，单一来源）
|
||||
unique_codes: list[str] = []
|
||||
seen_codes: set[str] = set()
|
||||
for row_codes, _ in rows:
|
||||
for c in row_codes:
|
||||
if c not in seen_codes:
|
||||
seen_codes.add(c)
|
||||
unique_codes.append(c)
|
||||
cand_by_code = {
|
||||
c.code: c
|
||||
for c in resolve_candidates(unique_codes, code_to_symbol, symbol_to_name, existing)
|
||||
}
|
||||
|
||||
candidates: list[dict[str, Any]] = []
|
||||
seen_symbols: set[str] = set() # 已发出的已匹配 symbol
|
||||
emitted_unmatched: set[str] = set() # 已发出的未匹配 code
|
||||
for row_codes, row_name in rows:
|
||||
# 行内代码优先：取第一个已匹配主数据的（避免价格/成交量数字误报）
|
||||
matched = next((cand_by_code[c] for c in row_codes if cand_by_code[c].matched), None)
|
||||
symbol = matched.symbol if matched else (name_to_symbol.get(row_name) if row_name else None)
|
||||
if symbol:
|
||||
if symbol in seen_symbols:
|
||||
continue
|
||||
seen_symbols.add(symbol)
|
||||
cand = matched or ImportCandidate(
|
||||
code=row_codes[0] if row_codes else "",
|
||||
symbol=symbol,
|
||||
name=symbol_to_name.get(symbol),
|
||||
matched=True,
|
||||
already_in_watchlist=symbol in existing,
|
||||
)
|
||||
else:
|
||||
# 名称兜底失败且无代码 → 表头/杂项行，忽略
|
||||
if not row_codes:
|
||||
continue
|
||||
code = row_codes[0]
|
||||
if code in emitted_unmatched:
|
||||
continue
|
||||
emitted_unmatched.add(code)
|
||||
cand = ImportCandidate(
|
||||
code=code,
|
||||
symbol=None,
|
||||
name=row_name,
|
||||
matched=False,
|
||||
already_in_watchlist=False,
|
||||
)
|
||||
candidates.append(cand.to_dict())
|
||||
|
||||
return _finalize("csv", text, unique_codes, candidates)
|
||||
@@ -0,0 +1,377 @@
|
||||
"""自选 CSV/TXT 与粘贴代码导入：解码、逐行/整段解析、主数据匹配、API 门禁。
|
||||
|
||||
导入落点统一走自选分组语义（watchlist.add_batch 的 group_ids M:N 并入），
|
||||
本测试同时覆盖该并入语义，作为前端「并入目标分组」交互的服务端依据。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import polars as pl
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from app.api import watchlist as watchlist_api
|
||||
from app.api.watchlist import ImportCodesRequest, import_from_codes, import_from_csv
|
||||
from app.config import settings
|
||||
from app.services import watchlist
|
||||
from app.services.watchlist_csv import (
|
||||
decode_csv_bytes,
|
||||
import_watchlist_codes,
|
||||
import_watchlist_csv,
|
||||
parse_csv_rows,
|
||||
)
|
||||
|
||||
|
||||
def _write_instruments(data_dir: Path) -> None:
|
||||
inst = data_dir / "instruments"
|
||||
inst.mkdir()
|
||||
pl.DataFrame(
|
||||
{
|
||||
"code": ["600036", "515880"],
|
||||
"symbol": ["600036.SH", "515880.SH"],
|
||||
"name": ["招商银行", "通信ETF国泰"],
|
||||
}
|
||||
).write_parquet(inst / "instruments.parquet")
|
||||
|
||||
|
||||
def _mock_upload(*, content: bytes, content_type: str = "text/csv", filename: str = "watchlist.csv") -> MagicMock:
|
||||
file = MagicMock()
|
||||
file.content_type = content_type
|
||||
file.filename = filename
|
||||
file.read = AsyncMock(return_value=content)
|
||||
return file
|
||||
|
||||
|
||||
def _mock_request(data_dir: Path) -> MagicMock:
|
||||
request = MagicMock()
|
||||
request.app.state.repo.store.data_dir = data_dir
|
||||
request.app.state.repo.get_name_map = lambda _symbols: {}
|
||||
return request
|
||||
|
||||
|
||||
# ---- 解码 ----
|
||||
|
||||
def test_decode_utf8_bom_stripped():
|
||||
raw = "代码,名称\n600036,招商银行\n".encode("utf-8-sig")
|
||||
text = decode_csv_bytes(raw)
|
||||
assert not text.startswith("")
|
||||
assert "招商银行" in text
|
||||
|
||||
|
||||
def test_decode_gbk_fallback():
|
||||
raw = "600519,贵州茅台\n".encode("gbk")
|
||||
assert "贵州茅台" in decode_csv_bytes(raw)
|
||||
|
||||
|
||||
def test_decode_unknown_encoding_raises():
|
||||
# 同时破坏 UTF-8 与 GBK 系编码的字节序列
|
||||
raw = b"\x00\x80\x00\x80\x00\x80"
|
||||
with pytest.raises(ValueError, match="编码"):
|
||||
decode_csv_bytes(raw)
|
||||
|
||||
|
||||
def test_decode_empty_raises():
|
||||
with pytest.raises(ValueError, match="空文件"):
|
||||
decode_csv_bytes(b"")
|
||||
|
||||
|
||||
# ---- 解析 ----
|
||||
|
||||
def test_parse_comma_csv_with_header():
|
||||
text = "股票代码,股票名称,最新价\n600036,招商银行,42.50\n515880,通信ETF,1.34\n"
|
||||
rows = parse_csv_rows(text)
|
||||
# 表头行保留（无代码），由 import 层按名称是否命中主数据过滤
|
||||
assert rows[0] == ([], "股票代码")
|
||||
assert [(codes, name) for codes, name in rows[1:]] == [
|
||||
(["600036"], "招商银行"),
|
||||
(["515880"], "通信ETF"),
|
||||
]
|
||||
|
||||
|
||||
def test_parse_tab_separated():
|
||||
text = "600036\t招商银行\t42.50\n515880\t通信ETF\t1.34\n"
|
||||
rows = parse_csv_rows(text)
|
||||
assert [(codes, name) for codes, name in rows] == [
|
||||
(["600036"], "招商银行"),
|
||||
(["515880"], "通信ETF"),
|
||||
]
|
||||
|
||||
|
||||
def test_parse_plain_codes_one_per_line():
|
||||
text = "600036\n515880\n"
|
||||
assert [codes for codes, _ in parse_csv_rows(text)] == [["600036"], ["515880"]]
|
||||
|
||||
|
||||
def test_parse_concatenated_code_name_cell():
|
||||
# 通达信风格：代码+名称在同一字段
|
||||
text = "600036招商银行\n515880通信ETF国泰\n"
|
||||
rows = parse_csv_rows(text)
|
||||
assert [(codes, name) for codes, name in rows] == [
|
||||
(["600036"], "600036招商银行"),
|
||||
(["515880"], "515880通信ETF国泰"),
|
||||
]
|
||||
|
||||
|
||||
# ---- 主数据匹配 / 名称兜底（CSV 行级） ----
|
||||
|
||||
def test_import_matched_name_fallback_and_unmatched(tmp_path: Path):
|
||||
_write_instruments(tmp_path)
|
||||
raw = (
|
||||
"代码,名称,现价\n"
|
||||
"600036,招商银行,42.50\n"
|
||||
"通信ETF国泰,1.34\n"
|
||||
"000001,平安银行\n"
|
||||
"999999,不存在,1.00\n"
|
||||
).encode()
|
||||
res = import_watchlist_csv(raw, tmp_path, existing_symbols={"600036.SH"})
|
||||
|
||||
by_code = {c["code"]: c for c in res["candidates"]}
|
||||
assert by_code["600036"]["matched"] and by_code["600036"]["already_in_watchlist"]
|
||||
assert by_code["600036"]["name"] == "招商银行"
|
||||
# 名称兜底命中（无代码行）
|
||||
assert any(
|
||||
c["symbol"] == "515880.SH" and c["matched"] and c["name"] == "通信ETF国泰"
|
||||
for c in res["candidates"]
|
||||
)
|
||||
assert by_code["000001"]["matched"] is False
|
||||
assert by_code["999999"]["matched"] is False
|
||||
assert res["matched_count"] == 2
|
||||
assert res["unmatched_count"] == 2
|
||||
assert res["provider"] == "csv"
|
||||
assert res["codes"] == ["600036", "000001", "999999"]
|
||||
|
||||
|
||||
def test_import_dedupe_and_header_skipped(tmp_path: Path):
|
||||
_write_instruments(tmp_path)
|
||||
raw = "代码,名称\n600036,招商银行\n600036,招商银行\n515880,通信ETF国泰\n".encode()
|
||||
res = import_watchlist_csv(raw, tmp_path)
|
||||
symbols = [c["symbol"] for c in res["candidates"]]
|
||||
assert symbols == ["600036.SH", "515880.SH"]
|
||||
|
||||
|
||||
def test_import_junk_rows_ignored(tmp_path: Path):
|
||||
_write_instruments(tmp_path)
|
||||
raw = "hello,world\n600036,招商银行\n,,\n".encode()
|
||||
res = import_watchlist_csv(raw, tmp_path)
|
||||
assert [c["symbol"] for c in res["candidates"]] == ["600036.SH"]
|
||||
|
||||
|
||||
# ---- 粘贴代码（整段抽码，不压缩同行多码） ----
|
||||
|
||||
def test_codes_dedupe_order_and_counts(tmp_path: Path):
|
||||
_write_instruments(tmp_path)
|
||||
res = import_watchlist_codes("600036\n515880\n999999\n600036", tmp_path)
|
||||
assert res["provider"] == "codes"
|
||||
assert res["codes"] == ["600036", "515880", "999999"]
|
||||
assert res["matched_count"] == 2
|
||||
assert res["unmatched_count"] == 1
|
||||
by_code = {c["code"]: c for c in res["candidates"]}
|
||||
assert by_code["600036"]["symbol"] == "600036.SH"
|
||||
assert by_code["999999"]["matched"] is False
|
||||
|
||||
|
||||
def test_codes_same_line_multiple_codes_not_collapsed(tmp_path: Path):
|
||||
# 同行多码（逗号/空格分隔）不得被压成单候选
|
||||
_write_instruments(tmp_path)
|
||||
res = import_watchlist_codes("600036, 515880 601636", tmp_path)
|
||||
assert res["codes"] == ["600036", "515880", "601636"]
|
||||
assert res["matched_count"] == 2
|
||||
assert res["unmatched_count"] == 1
|
||||
|
||||
|
||||
def test_codes_blank_returns_empty_without_parquet(tmp_path: Path):
|
||||
# 无码时不读主数据直接返回空结果（数据目录甚至不存在）
|
||||
res = import_watchlist_codes(" \n ", tmp_path)
|
||||
assert res["codes"] == []
|
||||
assert res["candidates"] == []
|
||||
assert res["matched_count"] == 0 and res["unmatched_count"] == 0
|
||||
|
||||
|
||||
def test_codes_over_max_raises(tmp_path: Path):
|
||||
text = "\n".join(f"{600000 + i}" for i in range(1001))
|
||||
with pytest.raises(ValueError, match="最多"):
|
||||
import_watchlist_codes(text, tmp_path, max_codes=1000)
|
||||
|
||||
|
||||
def test_import_functions_do_not_write_watchlist(tmp_path: Path, monkeypatch):
|
||||
_write_instruments(tmp_path)
|
||||
monkeypatch.setattr(settings, "data_dir", tmp_path)
|
||||
before = watchlist.list_symbols()
|
||||
assert before == []
|
||||
import_watchlist_csv("600036,招商银行\n".encode(), tmp_path)
|
||||
import_watchlist_codes("515880", tmp_path)
|
||||
assert watchlist.list_symbols() == []
|
||||
|
||||
|
||||
# ---- 分组合并语义（导入落点，M:N 并入服务端依据） ----
|
||||
|
||||
def test_add_batch_merges_into_target_group_mn(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(settings, "data_dir", tmp_path)
|
||||
_, g1 = watchlist.create_group("组合一")
|
||||
_, g2 = watchlist.create_group("组合二")
|
||||
watchlist.add("600036.SH", group_id=g1["id"])
|
||||
|
||||
rows, added = watchlist.add_batch(["600036.SH", "515880.SH"], group_id=g2["id"])
|
||||
by_symbol = {r["symbol"]: r for r in rows}
|
||||
# 已在自选且属组合一：并入组合二，两组共存（M:N）
|
||||
assert g1["id"] in by_symbol["600036.SH"]["group_ids"]
|
||||
assert g2["id"] in by_symbol["600036.SH"]["group_ids"]
|
||||
# 真正新增的才计 added
|
||||
assert g2["id"] in by_symbol["515880.SH"]["group_ids"]
|
||||
assert added == 1
|
||||
|
||||
|
||||
def test_add_batch_merges_into_multiple_groups_mn(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(settings, "data_dir", tmp_path)
|
||||
_, g1 = watchlist.create_group("组合一")
|
||||
_, g2 = watchlist.create_group("组合二")
|
||||
_, g3 = watchlist.create_group("组合三")
|
||||
watchlist.add("600036.SH", group_id=g1["id"])
|
||||
|
||||
# 一次批量调用同时并入 g2/g3（多组导入，原子单写；含重复 id 验证去重）
|
||||
rows, added = watchlist.add_batch(
|
||||
["600036.SH", "515880.SH"],
|
||||
group_ids=[g2["id"], g3["id"], g2["id"]],
|
||||
)
|
||||
by_symbol = {r["symbol"]: r for r in rows}
|
||||
gids600 = by_symbol["600036.SH"]["group_ids"]
|
||||
# 保留原有 g1 且并入 g2/g3
|
||||
assert sorted(gids600) == sorted([g1["id"], g2["id"], g3["id"]])
|
||||
assert sorted(by_symbol["515880.SH"]["group_ids"]) == sorted([g2["id"], g3["id"]])
|
||||
assert added == 1
|
||||
|
||||
|
||||
# ---- API 门禁：CSV ----
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_csv_rejects_bad_extension(tmp_path: Path):
|
||||
request = _mock_request(tmp_path)
|
||||
# 类型与扩展名都非白名单 → 拒绝
|
||||
file = _mock_upload(
|
||||
content=b"x",
|
||||
content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
filename="watchlist.xlsx",
|
||||
)
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
await import_from_csv(request, file)
|
||||
assert ei.value.status_code == 400
|
||||
assert "仅支持" in str(ei.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_csv_rejects_oversized_bytes(tmp_path: Path):
|
||||
request = _mock_request(tmp_path)
|
||||
file = _mock_upload(content=b"x" * (5 * 1024 * 1024 + 1))
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
await import_from_csv(request, file)
|
||||
assert ei.value.status_code == 400
|
||||
assert "过大" in str(ei.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_csv_rejects_empty(tmp_path: Path):
|
||||
request = _mock_request(tmp_path)
|
||||
file = _mock_upload(content=b"")
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
await import_from_csv(request, file)
|
||||
assert ei.value.status_code == 400
|
||||
assert "空文件" in str(ei.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_csv_no_candidates_raises(tmp_path: Path, monkeypatch):
|
||||
_write_instruments(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
request = _mock_request(tmp_path)
|
||||
file = _mock_upload(content=b"hello\nworld\n")
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
await import_from_csv(request, file)
|
||||
assert ei.value.status_code == 400
|
||||
assert "未识别" in str(ei.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_csv_maps_value_error_to_400(tmp_path: Path, monkeypatch):
|
||||
request = _mock_request(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
|
||||
def boom(*_a, **_k):
|
||||
raise ValueError("无法识别文件编码")
|
||||
|
||||
monkeypatch.setattr(watchlist_api, "import_watchlist_csv", boom)
|
||||
file = _mock_upload(content=b"x")
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
await import_from_csv(request, file)
|
||||
assert ei.value.status_code == 400
|
||||
assert "编码" in str(ei.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_csv_success(tmp_path: Path, monkeypatch):
|
||||
_write_instruments(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
request = _mock_request(tmp_path)
|
||||
file = _mock_upload(content="代码,名称\n600036,招商银行\n515880,通信ETF国泰\n".encode())
|
||||
res = await import_from_csv(request, file)
|
||||
assert res["provider"] == "csv"
|
||||
assert res["matched_count"] == 2
|
||||
assert {c["symbol"] for c in res["candidates"]} == {"600036.SH", "515880.SH"}
|
||||
assert "raw_text" not in res
|
||||
|
||||
|
||||
# ---- API 门禁：粘贴代码 ----
|
||||
|
||||
def test_import_codes_rejects_empty(tmp_path: Path, monkeypatch):
|
||||
request = _mock_request(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
import_from_codes(ImportCodesRequest(text=" "), request)
|
||||
assert ei.value.status_code == 400
|
||||
assert "请输入" in str(ei.value.detail)
|
||||
|
||||
|
||||
def test_import_codes_no_codes_raises(tmp_path: Path, monkeypatch):
|
||||
request = _mock_request(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
import_from_codes(ImportCodesRequest(text="hello world"), request)
|
||||
assert ei.value.status_code == 400
|
||||
assert "未识别" in str(ei.value.detail)
|
||||
|
||||
|
||||
def test_import_codes_over_max_raises(tmp_path: Path, monkeypatch):
|
||||
request = _mock_request(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
text = "\n".join(f"{600000 + i}" for i in range(1001))
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
import_from_codes(ImportCodesRequest(text=text), request)
|
||||
assert ei.value.status_code == 400
|
||||
assert "最多" in str(ei.value.detail)
|
||||
|
||||
|
||||
def test_import_codes_maps_value_error_to_400(tmp_path: Path, monkeypatch):
|
||||
request = _mock_request(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
|
||||
def boom(*_a, **_k):
|
||||
raise ValueError("一次最多导入 1000 个")
|
||||
|
||||
monkeypatch.setattr(watchlist_api, "import_watchlist_codes", boom)
|
||||
with pytest.raises(HTTPException) as ei:
|
||||
import_from_codes(ImportCodesRequest(text="600036"), request)
|
||||
assert ei.value.status_code == 400
|
||||
assert "最多" in str(ei.value.detail)
|
||||
|
||||
|
||||
def test_import_codes_success(tmp_path: Path, monkeypatch):
|
||||
_write_instruments(tmp_path)
|
||||
monkeypatch.setattr(watchlist, "list_symbols", lambda: [])
|
||||
request = _mock_request(tmp_path)
|
||||
res = import_from_codes(ImportCodesRequest(text="600036\n515880\n999999"), request)
|
||||
assert res["provider"] == "codes"
|
||||
assert res["matched_count"] == 2
|
||||
assert res["unmatched_count"] == 1
|
||||
assert {c["symbol"] for c in res["candidates"] if c["matched"]} == {"600036.SH", "515880.SH"}
|
||||
assert "raw_text" not in res
|
||||
Reference in New Issue
Block a user