xiaobai-review/backend/data/datahub/native.py

from __future__ import annotations

from typing import Any

from backend.data.numbers import finite_number


AMOUNT_THOUSAND_YUAN = 1000.0
AMOUNT_WAN_YUAN = 10000.0
VOLUME_LOT = 100.0

API_TO_DATASET = {
    "trade_cal": "calendar",
    "stock_basic": "stocks",
    "daily": "daily",
    "daily_basic": "valuation",
    "index_daily": "index_daily",
    "moneyflow": "moneyflow",
    "stk_auction": "auction",
}

SCALE_FIELDS = {
    "daily": {"vol": VOLUME_LOT, "amount": AMOUNT_THOUSAND_YUAN},
    "index_daily": {"vol": VOLUME_LOT, "amount": AMOUNT_THOUSAND_YUAN},
    "valuation": {"total_mv": AMOUNT_WAN_YUAN, "circ_mv": AMOUNT_WAN_YUAN},
    "moneyflow": {
        "buy_sm_amount": AMOUNT_WAN_YUAN,
        "sell_sm_amount": AMOUNT_WAN_YUAN,
        "buy_md_amount": AMOUNT_WAN_YUAN,
        "sell_md_amount": AMOUNT_WAN_YUAN,
        "buy_lg_amount": AMOUNT_WAN_YUAN,
        "sell_lg_amount": AMOUNT_WAN_YUAN,
        "buy_elg_amount": AMOUNT_WAN_YUAN,
        "sell_elg_amount": AMOUNT_WAN_YUAN,
        "net_mf_amount": AMOUNT_WAN_YUAN,
    },
    "auction": {"vol": VOLUME_LOT, "float_share": AMOUNT_WAN_YUAN},
}


def yyyymmdd(value: Any) -> str:
    return str(value or "").replace("-", "")[:8]


def to_native_rows(dataset: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [to_native_row(dataset, row) for row in rows]


def to_native_row(dataset: str, row: dict[str, Any]) -> dict[str, Any]:
    if dataset == "calendar":
        is_open = row.get("is_open")
        return {
            "exchange": str(row.get("exchange") or "SSE"),
            "cal_date": yyyymmdd(row.get("cal_date")),
            "is_open": 1 if is_open in (True, 1, "1", "Y", "y") else 0,
            "pretrade_date": yyyymmdd(row.get("pretrade_date")) or None,
        }
    converted = dict(row)
    converted.pop("batch_id", None)
    if "volume" in converted and "vol" not in converted:
        converted["vol"] = converted.pop("volume")
    elif "volume" in converted:
        converted.pop("volume", None)
    scales = SCALE_FIELDS.get(dataset) or {}
    for field, factor in scales.items():
        if field in converted:
            converted[field] = _unscale(converted.get(field), factor)
    if dataset == "stocks":
        converted.pop("updated_at", None)
    return converted


def to_canonical_row(dataset: str, row: dict[str, Any]) -> dict[str, Any]:
    if dataset == "calendar":
        is_open = row.get("is_open")
        return {
            "exchange": str(row.get("exchange") or "SSE"),
            "cal_date": yyyymmdd(row.get("cal_date")),
            "is_open": 1 if is_open in (True, 1, "1", "Y", "y") else 0,
            "pretrade_date": yyyymmdd(row.get("pretrade_date")) or None,
        }
    converted = dict(row)
    if "volume" in converted and "vol" not in converted:
        converted["vol"] = converted.pop("volume")
    scales = SCALE_FIELDS.get(dataset) or {}
    for field, factor in scales.items():
        if field in converted:
            converted[field] = _scale(converted.get(field), factor)
    return converted


def row_key(dataset: str, row: dict[str, Any]) -> tuple[str, ...]:
    if dataset == "calendar":
        return (yyyymmdd(row.get("cal_date")),)
    if dataset == "stocks":
        return (str(row.get("ts_code") or "").upper(),)
    if dataset == "status":
        return (str(row.get("dataset") or ""), yyyymmdd(row.get("trade_date")))
    return (str(row.get("ts_code") or "").upper(), yyyymmdd(row.get("trade_date")))


def project_fields(rows: list[dict[str, Any]], fields: str) -> list[dict[str, Any]]:
    keys = [item.strip() for item in str(fields or "").split(",") if item.strip()]
    if not keys:
        return rows
    return [{key: row.get(key) for key in keys} for row in rows]


def filter_stock_rows(rows: list[dict[str, Any]], params: dict[str, Any] | None) -> list[dict[str, Any]]:
    payload = params or {}
    ts_code = str(payload.get("ts_code") or "").strip().upper()
    status = str(payload.get("list_status") or "").strip()
    name = str(payload.get("name") or "").strip()
    filtered = rows
    if ts_code:
        filtered = [row for row in filtered if str(row.get("ts_code") or "").upper() == ts_code]
    if status:
        filtered = [row for row in filtered if str(row.get("list_status") or status) == status]
    if name:
        filtered = [row for row in filtered if name.casefold() in str(row.get("name") or "").casefold()]
    return filtered


def filter_calendar_rows(rows: list[dict[str, Any]], params: dict[str, Any] | None) -> list[dict[str, Any]]:
    payload = params or {}
    if payload.get("is_open") in (1, "1", True):
        return [row for row in rows if int(row.get("is_open") or 0) == 1]
    if payload.get("is_open") in (0, "0", False):
        return [row for row in rows if int(row.get("is_open") or 0) == 0]
    return rows


def _scale(value: Any, factor: float) -> float | None:
    number = _optional_number(value)
    if number is None:
        return None
    return number * factor


def _unscale(value: Any, factor: float) -> float | None:
    number = _optional_number(value)
    if number is None or factor == 0:
        return None
    return number / factor


def _optional_number(value: Any) -> float | None:
    if value in (None, ""):
        return None
    number = finite_number(value, default=float("nan"))
    if number != number:
        return None
    return number