mirror of
https://ghfast.top/https://github.com/aeroxw/easy_tdx_max.git
synced 2026-09-12 16:54:20 +08:00
feat(web): add FastAPI app factory, all routers, CLI serve command, and tests
- App factory with lifespan management and CORS middleware
- Market router: security list, quotes, market stat, fund-flow
- Bars router: kline, index kline, minute, transaction
- Finance router: xdxr, finance, company info, financial records
- Block router: block file parsing
- Chanlun router: POST /chanlun/analyze
- Realtime router: WebSocket /ws/realtime/{symbol}
- CLI: easy-tdx serve command
- 16 unit tests, all passing offline (no network)
This commit is contained in:
Justin Gu
@@ -29,6 +29,7 @@ from .cmd_quote import quote, quote_list
|
||||
from .cmd_run_all import run_all
|
||||
from .cmd_tick import tick
|
||||
from .cmd_transaction import transaction
|
||||
from .cmd_web import serve
|
||||
|
||||
|
||||
@click.group()
|
||||
@@ -84,3 +85,4 @@ cli.add_command(backtest)
|
||||
cli.add_command(portfolio)
|
||||
cli.add_command(run_all)
|
||||
cli.add_command(screen)
|
||||
cli.add_command(serve)
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
"""easy-tdx serve — 启动 Web API 服务器。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import click
|
||||
|
||||
|
||||
@click.command("serve")
|
||||
@click.option("--host", default="0.0.0.0", help="监听地址")
|
||||
@click.option("--port", default=8000, type=int, help="监听端口")
|
||||
@click.option("--tdx-host", default=None, help="TDX 服务器地址（默认自动选择最优）")
|
||||
@click.option("--tdx-port", default=None, type=int, help="TDX 服务器端口")
|
||||
@click.option("--reload", is_flag=True, help="开发模式（自动重载）")
|
||||
def serve(host: str, port: int, tdx_host: str | None, tdx_port: int | None, reload: bool) -> None:
|
||||
"""启动 Web API 服务器（需要安装 easy-tdx[web]）。"""
|
||||
try:
|
||||
import uvicorn
|
||||
except ImportError:
|
||||
click.echo(
|
||||
"错误：缺少 web 依赖。请运行: pip install easy-tdx[web]",
|
||||
err=True,
|
||||
)
|
||||
raise SystemExit(1) from None
|
||||
|
||||
if reload:
|
||||
uvicorn.run(
|
||||
"easy_tdx.web:app_factory",
|
||||
host=host,
|
||||
port=port,
|
||||
reload=True,
|
||||
factory=True,
|
||||
)
|
||||
else:
|
||||
from easy_tdx.web import create_app
|
||||
|
||||
app = create_app(host=tdx_host, port=tdx_port)
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""FastAPI application factory and lifespan management."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from easy_tdx.web.errors import register_exception_handlers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""管理 TDX 连接生命周期：启动时连接，关闭时断开。"""
|
||||
from easy_tdx.client import AsyncTdxClient
|
||||
|
||||
host = app.state.tdx_host
|
||||
port = app.state.tdx_port
|
||||
timeout = app.state.tdx_timeout
|
||||
|
||||
client = AsyncTdxClient(host=host, port=port, timeout=timeout)
|
||||
try:
|
||||
await client.connect()
|
||||
logger.info("TDX client connected to %s:%s", host, port)
|
||||
except Exception:
|
||||
logger.warning("TDX client connection failed — endpoints will return 503")
|
||||
|
||||
app.state.tdx_client = client
|
||||
yield
|
||||
|
||||
try:
|
||||
await client.close()
|
||||
logger.info("TDX client disconnected")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _create_app(
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> FastAPI:
|
||||
"""创建并配置 FastAPI 应用实例。"""
|
||||
from easy_tdx.config import get_best_host, get_port, get_timeout
|
||||
|
||||
if host is None:
|
||||
host = get_best_host()
|
||||
if port is None:
|
||||
port = get_port()
|
||||
if timeout is None:
|
||||
timeout = get_timeout()
|
||||
|
||||
app = FastAPI(
|
||||
title="easy-tdx API",
|
||||
description="通达信行情数据 REST + WebSocket API",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
# Store connection config in app.state for lifespan to use
|
||||
app.state.tdx_host = host
|
||||
app.state.tdx_port = port
|
||||
app.state.tdx_timeout = timeout
|
||||
app.state.tdx_client = None # will be set in lifespan
|
||||
|
||||
# CORS middleware (permissive for development)
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Register exception handlers
|
||||
register_exception_handlers(app)
|
||||
|
||||
# Mount routers
|
||||
from easy_tdx.web.routers.bars import router as bars_router
|
||||
from easy_tdx.web.routers.block import router as block_router
|
||||
from easy_tdx.web.routers.chanlun import router as chanlun_router
|
||||
from easy_tdx.web.routers.finance import router as finance_router
|
||||
from easy_tdx.web.routers.market import router as market_router
|
||||
from easy_tdx.web.routers.realtime import router as realtime_router
|
||||
|
||||
app.include_router(market_router, prefix="/api/v1")
|
||||
app.include_router(bars_router, prefix="/api/v1")
|
||||
app.include_router(finance_router, prefix="/api/v1")
|
||||
app.include_router(block_router, prefix="/api/v1")
|
||||
app.include_router(chanlun_router, prefix="/api/v1")
|
||||
app.include_router(realtime_router, prefix="/api/v1")
|
||||
|
||||
return app
|
||||
@@ -0,0 +1,16 @@
|
||||
"""Dependency injection for Web API routers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import typing
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from easy_tdx.client import AsyncTdxClient
|
||||
|
||||
|
||||
def get_client(request: Request) -> AsyncTdxClient:
|
||||
"""从 app.state 获取共享的 AsyncTdxClient 实例。"""
|
||||
return typing.cast(AsyncTdxClient, request.app.state.tdx_client)
|
||||
@@ -0,0 +1,113 @@
|
||||
"""K线 / 分时 / 逐笔成交路由。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
|
||||
from easy_tdx.web.deps import get_client
|
||||
from easy_tdx.web.schemas import DataFrameResponse, KlineCategoryEnum
|
||||
|
||||
router = APIRouter(tags=["bars"])
|
||||
|
||||
|
||||
def _market(market: str) -> Any:
|
||||
from easy_tdx.models.enums import Market
|
||||
|
||||
return Market[market.upper()]
|
||||
|
||||
|
||||
def _category(category: str) -> Any:
|
||||
from easy_tdx.models.enums import KlineCategory
|
||||
|
||||
# Support both int and string name
|
||||
try:
|
||||
return KlineCategory(int(category))
|
||||
except (ValueError, TypeError):
|
||||
return KlineCategory[KlineCategoryEnum[category.upper()].name]
|
||||
|
||||
|
||||
def _df_resp(df: Any) -> DataFrameResponse:
|
||||
return DataFrameResponse.from_dataframe(df)
|
||||
|
||||
|
||||
@router.get("/bars", response_model=DataFrameResponse)
|
||||
async def security_bars(
|
||||
market: str = Query(..., description="市场: SZ, SH, BJ"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
category: str = Query(
|
||||
"DAY",
|
||||
description="K线周期: MIN_1, MIN_5, MIN_15, MIN_30, MIN_60, DAY, WEEK, MONTH, YEAR",
|
||||
),
|
||||
start: int = Query(0, ge=0),
|
||||
count: int = Query(800, ge=1, le=800),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取股票K线数据。"""
|
||||
df = await client.get_security_bars(_market(market), code, _category(category), start, count)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/bars/index", response_model=DataFrameResponse)
|
||||
async def index_bars(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
category: str = Query("DAY", description="K线周期"),
|
||||
start: int = Query(0, ge=0),
|
||||
count: int = Query(800, ge=1, le=800),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取指数K线数据。"""
|
||||
df = await client.get_index_bars(_market(market), code, _category(category), start, count)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/minute", response_model=DataFrameResponse)
|
||||
async def minute_time(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取今日分时数据。"""
|
||||
df = await client.get_minute_time_data(_market(market), code)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/minute/history", response_model=DataFrameResponse)
|
||||
async def history_minute_time(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
date: int = Query(..., description="日期 YYYYMMDD"),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取历史某日分时数据。"""
|
||||
df = await client.get_history_minute_time_data(_market(market), code, date)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/transaction", response_model=DataFrameResponse)
|
||||
async def transaction_data(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
start: int = Query(0, ge=0),
|
||||
count: int = Query(800, ge=1, le=800),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取当日逐笔成交。"""
|
||||
df = await client.get_transaction_data(_market(market), code, start, count)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/transaction/history", response_model=DataFrameResponse)
|
||||
async def history_transaction_data(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
date: int = Query(..., description="日期 YYYYMMDD"),
|
||||
start: int = Query(0, ge=0),
|
||||
count: int = Query(800, ge=1, le=800),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取历史逐笔成交。"""
|
||||
df = await client.get_history_transaction_data(_market(market), code, date, start, count)
|
||||
return _df_resp(df)
|
||||
@@ -0,0 +1,29 @@
|
||||
"""板块信息路由。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
|
||||
from easy_tdx.web.deps import get_client
|
||||
from easy_tdx.web.schemas import DataFrameResponse
|
||||
|
||||
router = APIRouter(tags=["block"])
|
||||
|
||||
|
||||
def _df_resp(df: Any) -> DataFrameResponse:
|
||||
return DataFrameResponse.from_dataframe(df)
|
||||
|
||||
|
||||
@router.get("/block", response_model=DataFrameResponse)
|
||||
async def block_info(
|
||||
filename: str = Query(
|
||||
...,
|
||||
description=("板块文件名: block_zs.dat(行业指数), block_gn.dat(概念), block_fg.dat(风格)"),
|
||||
),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取并解析板块文件。"""
|
||||
df = await client.get_block_info(filename)
|
||||
return _df_resp(df)
|
||||
@@ -0,0 +1,64 @@
|
||||
"""缠论分析路由。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from easy_tdx.web.deps import get_client
|
||||
from easy_tdx.web.schemas import ChanlunRequest
|
||||
|
||||
router = APIRouter(tags=["chanlun"])
|
||||
|
||||
|
||||
def _market(market: str) -> Any:
|
||||
from easy_tdx.models.enums import Market
|
||||
|
||||
return Market[market.upper()]
|
||||
|
||||
|
||||
def _category(category: str) -> Any:
|
||||
from easy_tdx.models.enums import KlineCategory
|
||||
|
||||
try:
|
||||
return KlineCategory(int(category))
|
||||
except (ValueError, TypeError):
|
||||
return KlineCategory[category.upper()]
|
||||
|
||||
|
||||
@router.post("/chanlun/analyze")
|
||||
async def chanlun_analyze(
|
||||
req: ChanlunRequest,
|
||||
client: Any = Depends(get_client),
|
||||
) -> dict[str, Any]:
|
||||
"""执行缠论分析。
|
||||
|
||||
自动从 TDX 服务器获取 K 线数据，运行完整缠论计算管道，
|
||||
返回笔、中枢、线段、买卖点、背驰等分析结果。
|
||||
"""
|
||||
from easy_tdx.chanlun import ChanlunAnalyser
|
||||
|
||||
# 1. Fetch kline data
|
||||
df = await client.get_security_bars(
|
||||
_market(req.market), req.code, _category(req.category), req.start, req.count
|
||||
)
|
||||
|
||||
# 2. Run chanlun analysis
|
||||
symbol = f"{req.market}{req.code}"
|
||||
frequency_map: dict[str, str] = {
|
||||
"MIN_1": "1min",
|
||||
"MIN_5": "5min",
|
||||
"MIN_15": "15min",
|
||||
"MIN_30": "30min",
|
||||
"MIN_60": "60min",
|
||||
"DAY": "daily",
|
||||
"WEEK": "weekly",
|
||||
"MONTH": "monthly",
|
||||
"YEAR": "yearly",
|
||||
}
|
||||
freq = frequency_map.get(req.category.upper(), req.category)
|
||||
analyser = ChanlunAnalyser(code=symbol, frequency=freq)
|
||||
result = analyser.process_klines(df)
|
||||
|
||||
return result.to_dict()
|
||||
@@ -0,0 +1,88 @@
|
||||
"""财务 / 公司信息路由。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
|
||||
from easy_tdx.web.deps import get_client
|
||||
from easy_tdx.web.schemas import DataFrameResponse
|
||||
|
||||
router = APIRouter(tags=["finance"])
|
||||
|
||||
|
||||
def _market(market: str) -> Any:
|
||||
from easy_tdx.models.enums import Market
|
||||
|
||||
return Market[market.upper()]
|
||||
|
||||
|
||||
def _df_resp(df: Any) -> DataFrameResponse:
|
||||
return DataFrameResponse.from_dataframe(df)
|
||||
|
||||
|
||||
@router.get("/xdxr", response_model=DataFrameResponse)
|
||||
async def xdxr_info(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取除权除息历史记录。"""
|
||||
df = await client.get_xdxr_info(_market(market), code)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/finance", response_model=DataFrameResponse)
|
||||
async def finance_info(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取最新财务数据。"""
|
||||
df = await client.get_finance_info(_market(market), code)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/company/category", response_model=DataFrameResponse)
|
||||
async def company_info_category(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取公司信息文件目录。"""
|
||||
df = await client.get_company_info_category(_market(market), code)
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/company/content")
|
||||
async def company_info_content(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6),
|
||||
filename: str = Query(..., description="文件名"),
|
||||
offset: int = Query(0, ge=0),
|
||||
length: int = Query(1024, ge=1),
|
||||
client: Any = Depends(get_client),
|
||||
) -> dict[str, str]:
|
||||
"""读取公司信息文本。"""
|
||||
content = await client.get_company_info_content(_market(market), code, filename, offset, length)
|
||||
return {"content": content}
|
||||
|
||||
|
||||
@router.get("/financial/file-list", response_model=DataFrameResponse)
|
||||
async def financial_file_list(
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取可用的历史专业财报文件列表。"""
|
||||
df = await client.get_financial_file_list()
|
||||
return _df_resp(df)
|
||||
|
||||
|
||||
@router.get("/financial/records", response_model=DataFrameResponse)
|
||||
async def financial_records(
|
||||
filename: str = Query(..., description="财报文件名，如 tdxfin/gpcw20260331.zip"),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""下载财报 zip 并解析为记录列表。"""
|
||||
df = await client.get_financial_records(filename)
|
||||
return _df_resp(df)
|
||||
@@ -0,0 +1,109 @@
|
||||
"""市场信息路由：证券列表、实时行情、市场统计、资金流向。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
|
||||
from easy_tdx.web.deps import get_client
|
||||
from easy_tdx.web.schemas import (
|
||||
CountResponse,
|
||||
DataFrameResponse,
|
||||
MarketEnum,
|
||||
QuoteRequest,
|
||||
)
|
||||
|
||||
router = APIRouter(tags=["market"])
|
||||
|
||||
|
||||
def _market_from_str(s: str) -> Any:
|
||||
"""将字符串转为 Market 枚举。"""
|
||||
from easy_tdx.models.enums import Market
|
||||
|
||||
return Market[MarketEnum[s].name]
|
||||
|
||||
|
||||
def _df_response(df: Any) -> DataFrameResponse:
|
||||
"""将 DataFrame 转为 API 响应。"""
|
||||
return DataFrameResponse.from_dataframe(df)
|
||||
|
||||
|
||||
@router.get("/security/count", response_model=CountResponse)
|
||||
async def security_count(
|
||||
market: str = Query(..., description="市场: SZ, SH, BJ"),
|
||||
client: Any = Depends(get_client),
|
||||
) -> CountResponse:
|
||||
"""获取市场证券总数。"""
|
||||
count = await client.get_security_count(_market_from_str(market))
|
||||
return CountResponse(count=count)
|
||||
|
||||
|
||||
@router.get("/security/list", response_model=DataFrameResponse)
|
||||
async def security_list(
|
||||
market: str = Query(..., description="市场: SZ, SH, BJ"),
|
||||
start: int = Query(0, ge=0, description="分页起始位置"),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取证券列表（每页约1000条）。"""
|
||||
df = await client.get_security_list(_market_from_str(market), start)
|
||||
return _df_response(df)
|
||||
|
||||
|
||||
@router.get("/security/list-all", response_model=DataFrameResponse)
|
||||
async def security_list_all(
|
||||
pages: int = Query(1, ge=1, description="拉取页数（每个市场每页1000条）"),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取沪深 A 股完整列表。"""
|
||||
df = await client.get_security_list_all(pages=pages)
|
||||
return _df_response(df)
|
||||
|
||||
|
||||
@router.post("/quotes", response_model=DataFrameResponse)
|
||||
async def security_quotes(
|
||||
req: QuoteRequest,
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""批量获取实时五档行情（最多80只/次）。"""
|
||||
from easy_tdx.models.enums import Market
|
||||
|
||||
stocks_parsed: list[tuple[Any, str]] = []
|
||||
for s in req.stocks:
|
||||
m = Market[MarketEnum[s.market].name]
|
||||
stocks_parsed.append((m, s.code))
|
||||
df = await client.get_security_quotes(stocks_parsed)
|
||||
return _df_response(df)
|
||||
|
||||
|
||||
@router.get("/market/stat", response_model=DataFrameResponse)
|
||||
async def market_stat(
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取 A 股全市场涨跌统计。"""
|
||||
df = await client.get_market_stat()
|
||||
return _df_response(df)
|
||||
|
||||
|
||||
@router.get("/fund-flow", response_model=DataFrameResponse)
|
||||
async def fund_flow(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6, description="6位股票代码"),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取个股当日资金流向。"""
|
||||
df = await client.get_fund_flow(_market_from_str(market), code)
|
||||
return _df_response(df)
|
||||
|
||||
|
||||
@router.get("/fund-flow/history", response_model=DataFrameResponse)
|
||||
async def history_fund_flow(
|
||||
market: str = Query(..., description="市场: SZ, SH"),
|
||||
code: str = Query(..., min_length=6, max_length=6, description="6位股票代码"),
|
||||
start: int = Query(0, ge=0),
|
||||
count: int = Query(100, ge=1, le=800),
|
||||
client: Any = Depends(get_client),
|
||||
) -> DataFrameResponse:
|
||||
"""获取个股历史日线资金流向。"""
|
||||
df = await client.get_history_fund_flow(_market_from_str(market), code, start, count)
|
||||
return _df_response(df)
|
||||
@@ -0,0 +1,103 @@
|
||||
"""实时数据 WebSocket 路由。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(tags=["realtime"])
|
||||
|
||||
|
||||
@router.websocket("/ws/realtime/{symbol}")
|
||||
async def realtime_websocket(websocket: WebSocket, symbol: str) -> None:
|
||||
"""WebSocket 实时行情订阅。
|
||||
|
||||
连接后自动订阅指定标的的实时事件。
|
||||
symbol 格式: SZ000001, SH600000 等。
|
||||
|
||||
客户端可发送 JSON 消息来控制订阅：
|
||||
- {"action": "subscribe", "symbol": "SZ000001"}
|
||||
- {"action": "unsubscribe", "symbol": "SZ000001"}
|
||||
|
||||
服务端推送消息格式：
|
||||
- {"type": "tick", "market": "SZ", "code": "000001", "price": 10.5, ...}
|
||||
- {"type": "signal", "direction": "BUY", ...}
|
||||
"""
|
||||
await websocket.accept()
|
||||
logger.info("WebSocket client connected for symbol: %s", symbol)
|
||||
|
||||
# Try to get EventBus from app state
|
||||
event_bus = getattr(websocket.app.state, "event_bus", None)
|
||||
|
||||
subscribed_symbols: set[str] = {symbol.upper()}
|
||||
|
||||
async def _on_event(event: Any) -> None:
|
||||
"""EventBus 回调 → 推送 WebSocket 消息。"""
|
||||
event_symbol = f"{event.market}{event.code}"
|
||||
if event_symbol in subscribed_symbols:
|
||||
try:
|
||||
msg = {
|
||||
"type": event.event_type.value,
|
||||
"market": event.market,
|
||||
"code": event.code,
|
||||
"price": event.price,
|
||||
"volume": event.volume,
|
||||
"timestamp": event.timestamp,
|
||||
"data": event.data,
|
||||
}
|
||||
await websocket.send_json(msg)
|
||||
except Exception:
|
||||
logger.warning("Failed to send WebSocket message")
|
||||
|
||||
# Subscribe to event bus if available
|
||||
if event_bus is not None:
|
||||
event_bus.subscribe_all(_on_event)
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Receive client messages (subscribe/unsubscribe control)
|
||||
try:
|
||||
raw = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
|
||||
data = json.loads(raw)
|
||||
action = data.get("action", "")
|
||||
|
||||
if action == "subscribe":
|
||||
new_symbol = data.get("symbol", "").upper()
|
||||
if new_symbol:
|
||||
subscribed_symbols.add(new_symbol)
|
||||
await websocket.send_json(
|
||||
{"type": "status", "msg": f"subscribed {new_symbol}"}
|
||||
)
|
||||
|
||||
elif action == "unsubscribe":
|
||||
old_symbol = data.get("symbol", "").upper()
|
||||
subscribed_symbols.discard(old_symbol)
|
||||
await websocket.send_json(
|
||||
{"type": "status", "msg": f"unsubscribed {old_symbol}"}
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# Send heartbeat ping
|
||||
try:
|
||||
await websocket.send_json({"type": "ping"})
|
||||
except Exception:
|
||||
break
|
||||
except WebSocketDisconnect:
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
await websocket.send_json({"type": "error", "msg": "invalid JSON"})
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info("WebSocket client disconnected: %s", symbol)
|
||||
except Exception:
|
||||
logger.exception("WebSocket error for %s", symbol)
|
||||
finally:
|
||||
if event_bus is not None:
|
||||
event_bus.unsubscribe(symbol.upper(), _on_event)
|
||||
logger.info("WebSocket connection closed: %s", symbol)
|
||||
+193
-1
@@ -1,9 +1,17 @@
|
||||
"""Web API schemas and error handling tests (offline, no network)."""
|
||||
"""Web API tests (offline, no network).
|
||||
|
||||
Covers: schemas, error handling, app factory, DI, all routers,
|
||||
CLI serve command, OpenAPI schema generation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 2: Schemas & Error Handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_market_enum_values():
|
||||
"""MarketEnum should map string names to int values matching Market enum."""
|
||||
@@ -55,3 +63,187 @@ def test_api_error_response():
|
||||
d = err.model_dump()
|
||||
assert d["error"] == "test error"
|
||||
assert d["detail"] == "some detail"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 3: App Factory & Dependency Injection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_app_returns_fastapi_instance():
|
||||
"""create_app should return a FastAPI app with routers mounted."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web import create_app
|
||||
|
||||
app = create_app()
|
||||
assert app.title == "easy-tdx API"
|
||||
|
||||
# Check routers are mounted
|
||||
routes = [r.path for r in app.routes]
|
||||
assert any("/api/v1/security" in r for r in routes)
|
||||
assert any("/api/v1/bars" in r for r in routes)
|
||||
assert any("/api/v1/chanlun" in r for r in routes)
|
||||
assert any("/ws/realtime" in r for r in routes)
|
||||
|
||||
|
||||
def test_deps_get_client_type():
|
||||
"""get_client should be callable (actual client creation needs network)."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.deps import get_client
|
||||
|
||||
assert callable(get_client)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 4: Market Router
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_market_router_endpoints():
|
||||
"""Market router should define all expected endpoints."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.routers.market import router
|
||||
|
||||
paths = [r.path for r in router.routes]
|
||||
assert "/security/count" in paths
|
||||
assert "/security/list" in paths
|
||||
assert "/security/list-all" in paths
|
||||
assert "/quotes" in paths
|
||||
assert "/market/stat" in paths
|
||||
assert "/fund-flow" in paths
|
||||
assert "/fund-flow/history" in paths
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 5: Bars Router
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_bars_router_endpoints():
|
||||
"""Bars router should define all expected endpoints."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.routers.bars import router
|
||||
|
||||
paths = [r.path for r in router.routes]
|
||||
assert "/bars" in paths
|
||||
assert "/bars/index" in paths
|
||||
assert "/minute" in paths
|
||||
assert "/minute/history" in paths
|
||||
assert "/transaction" in paths
|
||||
assert "/transaction/history" in paths
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 6: Finance Router
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_finance_router_endpoints():
|
||||
"""Finance router should define all expected endpoints."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.routers.finance import router
|
||||
|
||||
paths = [r.path for r in router.routes]
|
||||
assert "/xdxr" in paths
|
||||
assert "/finance" in paths
|
||||
assert "/company/category" in paths
|
||||
assert "/company/content" in paths
|
||||
assert "/financial/file-list" in paths
|
||||
assert "/financial/records" in paths
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 7: Block Router
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_block_router_endpoints():
|
||||
"""Block router should define expected endpoints."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.routers.block import router
|
||||
|
||||
paths = [r.path for r in router.routes]
|
||||
assert "/block" in paths
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 8: Chanlun Router
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_chanlun_router_endpoints():
|
||||
"""Chanlun router should define the analyze endpoint."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.routers.chanlun import router
|
||||
|
||||
paths = [r.path for r in router.routes]
|
||||
assert "/chanlun/analyze" in paths
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 9: Realtime Router
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_realtime_router_endpoints():
|
||||
"""Realtime router should define the WebSocket endpoint."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web.routers.realtime import router
|
||||
|
||||
paths = [r.path for r in router.routes]
|
||||
assert any("realtime" in p for p in paths)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 10: CLI serve command
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_serve_command_exists():
|
||||
"""CLI should have a serve command registered."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.cli import cli
|
||||
|
||||
assert "serve" in cli.commands
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 11: Integration — route registration & OpenAPI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_full_app_routes_registered():
|
||||
"""All routers should be mounted and accessible."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web import create_app
|
||||
|
||||
app = create_app()
|
||||
all_paths = [r.path for r in app.routes]
|
||||
expected_prefixes = [
|
||||
"/api/v1/security",
|
||||
"/api/v1/bars",
|
||||
"/api/v1/xdxr",
|
||||
"/api/v1/block",
|
||||
"/api/v1/chanlun",
|
||||
"/ws/realtime",
|
||||
]
|
||||
for prefix in expected_prefixes:
|
||||
matched = any(prefix in p for p in all_paths)
|
||||
assert matched, f"Expected route with prefix '{prefix}' not found in {all_paths}"
|
||||
|
||||
|
||||
def test_openapi_schema_generated():
|
||||
"""OpenAPI schema should be auto-generated and contain key paths."""
|
||||
pytest.importorskip("fastapi")
|
||||
from easy_tdx.web import create_app
|
||||
|
||||
app = create_app()
|
||||
schema = app.openapi()
|
||||
assert schema["info"]["title"] == "easy-tdx API"
|
||||
assert "/api/v1/security/count" in schema["paths"]
|
||||
assert "/api/v1/bars" in schema["paths"]
|
||||
assert "/api/v1/chanlun/analyze" in schema["paths"]
|
||||
# WebSocket routes are NOT included in OpenAPI schema by default;
|
||||
# they are verified in test_full_app_routes_registered instead.
|
||||
# Just ensure REST paths are present.
|
||||
assert "/api/v1/fund-flow" in schema["paths"]
|
||||
|
||||
Reference in New Issue
Block a user