diff --git a/src/easy_tdx/web/errors.py b/src/easy_tdx/web/errors.py new file mode 100644 index 0000000..a08737b --- /dev/null +++ b/src/easy_tdx/web/errors.py @@ -0,0 +1,46 @@ +"""Web API error handling.""" + +from __future__ import annotations + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from pydantic import BaseModel + +from easy_tdx.exceptions import TdxConnectionError + + +class ApiErrorResponse(BaseModel): + """标准错误响应格式。""" + + error: str + detail: str = "" + + +def register_exception_handlers(app: FastAPI) -> None: + """注册全局异常处理器到 FastAPI app。""" + + @app.exception_handler(TdxConnectionError) + async def tdx_connection_error_handler( + request: Request, exc: TdxConnectionError + ) -> JSONResponse: + del request # unused but required by FastAPI signature + return JSONResponse( + status_code=503, + content=ApiErrorResponse(error="TDX connection error", detail=str(exc)).model_dump(), + ) + + @app.exception_handler(ValueError) + async def value_error_handler(request: Request, exc: ValueError) -> JSONResponse: + del request # unused but required by FastAPI signature + return JSONResponse( + status_code=400, + content=ApiErrorResponse(error="Bad request", detail=str(exc)).model_dump(), + ) + + @app.exception_handler(Exception) + async def generic_error_handler(request: Request, exc: Exception) -> JSONResponse: + del request # unused but required by FastAPI signature + return JSONResponse( + status_code=500, + content=ApiErrorResponse(error="Internal server error", detail=str(exc)).model_dump(), + ) diff --git a/src/easy_tdx/web/schemas.py b/src/easy_tdx/web/schemas.py new file mode 100644 index 0000000..9fc7377 --- /dev/null +++ b/src/easy_tdx/web/schemas.py @@ -0,0 +1,105 @@ +"""Pydantic request/response schemas for the Web API.""" + +from __future__ import annotations + +from enum import IntEnum +from typing import Any + +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- +# Enums — mirror easy_tdx.models.enums but as string-based for REST clarity +# --------------------------------------------------------------------------- + + +class MarketEnum(IntEnum): + """Market identifier.""" + + SZ = 0 + SH = 1 + BJ = 2 + + +class KlineCategoryEnum(IntEnum): + """K-line period.""" + + MIN_5 = 0 + MIN_15 = 1 + MIN_30 = 2 + MIN_60 = 3 + DAY = 4 + WEEK = 5 + MONTH = 6 + MIN_1 = 7 + YEAR = 9 + SEASON = 10 + + +# --------------------------------------------------------------------------- +# Request models +# --------------------------------------------------------------------------- + + +class StockIdentifier(BaseModel): + """A single stock identified by market + code.""" + + market: str = Field(..., pattern=r"^(SZ|SH|BJ)$", description="市场代码") + code: str = Field(..., min_length=6, max_length=6, description="6位股票代码") + + +class QuoteRequest(BaseModel): + """Batch quote request.""" + + stocks: list[StockIdentifier] = Field( + ..., min_length=1, max_length=80, description="股票列表(最多80只)" + ) + + +class ChanlunRequest(BaseModel): + """缠论分析请求。""" + + market: str = Field(..., pattern=r"^(SZ|SH|BJ)$") + code: str = Field(..., min_length=6, max_length=6) + category: str = Field(default="DAY", description="K线周期") + count: int = Field(default=800, ge=1, le=800) + start: int = Field(default=0, ge=0) + + +# --------------------------------------------------------------------------- +# Response models +# --------------------------------------------------------------------------- + + +class DataFrameResponse(BaseModel): + """通用 DataFrame 响应(records 格式)。""" + + data: list[dict[str, Any]] + count: int + + @classmethod + def from_dataframe(cls, df: Any) -> DataFrameResponse: + """从 pandas DataFrame 构建响应。""" + import pandas as pd + + if isinstance(df, pd.DataFrame): + records = df.to_dict(orient="records") + cleaned: list[dict[str, Any]] = [] + for row in records: + clean_row: dict[str, Any] = {} + for k, v in row.items(): + if hasattr(v, "isoformat"): + clean_row[k] = v.isoformat() + elif hasattr(v, "item"): + # numpy scalar → Python native + clean_row[k] = v.item() + else: + clean_row[k] = v + cleaned.append(clean_row) + return cls(data=cleaned, count=len(cleaned)) + return cls(data=[], count=0) + + +class CountResponse(BaseModel): + """简单计数响应。""" + + count: int diff --git a/tests/unit/test_web_api.py b/tests/unit/test_web_api.py new file mode 100644 index 0000000..2766ee8 --- /dev/null +++ b/tests/unit/test_web_api.py @@ -0,0 +1,57 @@ +"""Web API schemas and error handling tests (offline, no network).""" + +from __future__ import annotations + +import pytest + + +def test_market_enum_values(): + """MarketEnum should map string names to int values matching Market enum.""" + pytest.importorskip("fastapi") + from easy_tdx.web.schemas import MarketEnum + + assert MarketEnum.SZ == 0 + assert MarketEnum.SH == 1 + assert MarketEnum.BJ == 2 + + +def test_kline_category_enum(): + """KlineCategoryEnum should map string names to int values.""" + pytest.importorskip("fastapi") + from easy_tdx.web.schemas import KlineCategoryEnum + + assert KlineCategoryEnum.MIN_5 == 0 + assert KlineCategoryEnum.DAY == 4 + assert KlineCategoryEnum.WEEK == 5 + + +def test_quote_request_validation(): + """QuoteRequest should validate stocks list.""" + pytest.importorskip("fastapi") + from easy_tdx.web.schemas import QuoteRequest + + req = QuoteRequest(stocks=[{"market": "SZ", "code": "000001"}]) + assert len(req.stocks) == 1 + assert req.stocks[0].market == "SZ" + assert req.stocks[0].code == "000001" + + +def test_chanlun_request_defaults(): + """ChanlunRequest should have sensible defaults.""" + pytest.importorskip("fastapi") + from easy_tdx.web.schemas import ChanlunRequest + + req = ChanlunRequest(market="SZ", code="000001") + assert req.category == "DAY" + assert req.count == 800 + + +def test_api_error_response(): + """ApiErrorResponse should serialize correctly.""" + pytest.importorskip("fastapi") + from easy_tdx.web.errors import ApiErrorResponse + + err = ApiErrorResponse(error="test error", detail="some detail") + d = err.model_dump() + assert d["error"] == "test error" + assert d["detail"] == "some detail"