xiaobai-review/xiaobai-datahub/datahub/normalize.py

"""Canonical field normalization for Tushare-native rows.

Units (architecture §7.1):
- price: 4 decimal REAL
- pct_chg: percent, 4 decimal REAL
- volume: shares (Tushare daily/index vol is 手 → ×100)
- amount: yuan (Tushare daily/index amount is 千元 → ×1000)
- moneyflow amounts: yuan (Tushare is 万元 → ×1e4)
- daily_basic total_mv / circ_mv: yuan (Tushare is 万元 → ×1e4)
- stk_auction.amount is already yuan in Tushare; volume 手 → ×100

Existing xiaobai-review stores Tushare native units and converts at display time.
Hub converts once at ingest. Golden tests compare hub output against applying
these same factors to review-native rows.
"""

from __future__ import annotations

from typing import Any

from datahub.numbers import finite_number, round4

AMOUNT_THOUSAND_YUAN = 1000.0
AMOUNT_WAN_YUAN = 10000.0
VOLUME_LOT = 100.0

DAILY_FIELDS = ("ts_code", "trade_date", "open", "high", "low", "close", "pct_chg", "vol", "amount")
VALUATION_FIELDS = (
    "ts_code", "trade_date", "turnover_rate", "volume_ratio",
    "total_mv", "circ_mv", "pe_ttm", "pb", "ps_ttm", "dv_ttm",
)
MONEYFLOW_FIELDS = (
    "ts_code", "trade_date",
    "buy_sm_amount", "sell_sm_amount", "buy_md_amount", "sell_md_amount",
    "buy_lg_amount", "sell_lg_amount", "buy_elg_amount", "sell_elg_amount",
    "net_mf_amount",
)
AUCTION_FIELDS = (
    "ts_code", "trade_date", "vol", "price", "amount", "pre_close",
    "turnover_rate", "volume_ratio", "float_share",
)
INDEX_FIELDS = ("ts_code", "trade_date", "open", "high", "low", "close", "pct_chg", "vol", "amount")
CALENDAR_FIELDS = ("exchange", "cal_date", "is_open", "pretrade_date")
STOCK_FIELDS = ("ts_code", "symbol", "name", "area", "industry", "market", "list_status", "list_date")


def _code(value: Any) -> str:
    return str(value or "").strip().upper()


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


def review_daily_to_canonical(row: dict[str, Any]) -> dict[str, Any]:
    """Convert a review-stored daily row (Tushare native units) to hub canonical."""
    return normalize_daily(row)


def normalize_daily(row: dict[str, Any], adj_factor: float | None = None) -> dict[str, Any]:
    return {
        "ts_code": _code(row.get("ts_code")),
        "trade_date": _date(row.get("trade_date")),
        "open": round4(finite_number(row.get("open"))),
        "high": round4(finite_number(row.get("high"))),
        "low": round4(finite_number(row.get("low"))),
        "close": round4(finite_number(row.get("close"))),
        "pct_chg": round4(finite_number(row.get("pct_chg"))),
        "volume": round4(_scale(row.get("vol"), VOLUME_LOT)),
        "amount": round4(_scale(row.get("amount"), AMOUNT_THOUSAND_YUAN)),
        "adj_factor": round4(finite_number(adj_factor if adj_factor is not None else row.get("adj_factor"))),
    }


def normalize_valuation(row: dict[str, Any]) -> dict[str, Any]:
    return {
        "ts_code": _code(row.get("ts_code")),
        "trade_date": _date(row.get("trade_date")),
        "turnover_rate": round4(finite_number(row.get("turnover_rate"))),
        "volume_ratio": round4(finite_number(row.get("volume_ratio"))),
        "total_mv": round4(_scale(row.get("total_mv"), AMOUNT_WAN_YUAN)),
        "circ_mv": round4(_scale(row.get("circ_mv"), AMOUNT_WAN_YUAN)),
        "pe_ttm": round4(finite_number(row.get("pe_ttm"))),
        "pb": round4(finite_number(row.get("pb"))),
        "ps_ttm": round4(finite_number(row.get("ps_ttm"))),
        "dv_ttm": round4(finite_number(row.get("dv_ttm"))),
    }


def normalize_moneyflow(row: dict[str, Any]) -> dict[str, Any]:
    converted = {
        "ts_code": _code(row.get("ts_code")),
        "trade_date": _date(row.get("trade_date")),
    }
    for field in MONEYFLOW_FIELDS[2:]:
        converted[field] = round4(_scale(row.get(field), AMOUNT_WAN_YUAN))
    return converted


def normalize_auction(row: dict[str, Any]) -> dict[str, Any]:
    return {
        "ts_code": _code(row.get("ts_code")),
        "trade_date": _date(row.get("trade_date")),
        "volume": round4(_scale(row.get("vol") if row.get("vol") is not None else row.get("volume"), VOLUME_LOT)),
        "price": round4(finite_number(row.get("price"))),
        "amount": round4(finite_number(row.get("amount"))),
        "pre_close": round4(finite_number(row.get("pre_close"))),
        "turnover_rate": round4(finite_number(row.get("turnover_rate"))),
        "volume_ratio": round4(finite_number(row.get("volume_ratio"))),
        "float_share": round4(_scale(row.get("float_share"), AMOUNT_WAN_YUAN) if row.get("float_share") is not None else None),
    }


def normalize_index_daily(row: dict[str, Any]) -> dict[str, Any]:
    return {
        "ts_code": _code(row.get("ts_code")),
        "trade_date": _date(row.get("trade_date")),
        "open": round4(finite_number(row.get("open"))),
        "high": round4(finite_number(row.get("high"))),
        "low": round4(finite_number(row.get("low"))),
        "close": round4(finite_number(row.get("close"))),
        "pct_chg": round4(finite_number(row.get("pct_chg"))),
        "volume": round4(_scale(row.get("vol"), VOLUME_LOT)),
        "amount": round4(_scale(row.get("amount"), AMOUNT_THOUSAND_YUAN)),
    }


def normalize_calendar(row: dict[str, Any]) -> dict[str, Any]:
    is_open = row.get("is_open")
    if is_open in (True, "1", 1, "Y", "y"):
        open_flag = 1
    elif is_open in (False, "0", 0, "N", "n", None, ""):
        open_flag = 0
    else:
        open_flag = int(is_open)
    return {
        "exchange": str(row.get("exchange") or "SSE"),
        "cal_date": _date(row.get("cal_date") or row.get("calDate")),
        "is_open": open_flag,
        "pretrade_date": _date(row.get("pretrade_date")) or None,
    }


def normalize_stock(row: dict[str, Any]) -> dict[str, Any]:
    ts_code = _code(row.get("ts_code"))
    symbol = str(row.get("symbol") or "").strip() or (ts_code.split(".")[0] if ts_code else "")
    return {
        "ts_code": ts_code,
        "symbol": symbol,
        "name": str(row.get("name") or "").strip(),
        "area": str(row.get("area") or "").strip() or None,
        "industry": str(row.get("industry") or "").strip() or None,
        "market": str(row.get("market") or "").strip() or None,
        "list_status": str(row.get("list_status") or "L").strip() or "L",
        "list_date": _date(row.get("list_date")) or None,
    }


def apply_qfq(price: float | None, factor: float | None, latest_factor: float | None) -> float | None:
    if price is None:
        return None
    current = factor if factor not in (None, 0) else 1.0
    latest = latest_factor if latest_factor not in (None, 0) else current
    return round4(price * current / latest)


def qfq_bar(row: dict[str, Any], latest_factor: float | None) -> dict[str, Any]:
    factor = finite_number(row.get("adj_factor"), 1.0) or 1.0
    out = dict(row)
    for field in ("open", "high", "low", "close"):
        out[field] = apply_qfq(finite_number(row.get(field)), factor, latest_factor)
    return out


NORMALIZERS = {
    "daily": normalize_daily,
    "valuation": normalize_valuation,
    "daily_basic": normalize_valuation,
    "moneyflow": normalize_moneyflow,
    "auction": normalize_auction,
    "stk_auction": normalize_auction,
    "index_daily": normalize_index_daily,
    "trade_cal": normalize_calendar,
    "calendar": normalize_calendar,
    "stock_basic": normalize_stock,
    "stocks": normalize_stock,
}


def normalize_rows(dataset: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    fn = NORMALIZERS.get(dataset)
    if fn is None:
        raise ValueError(f"unknown dataset: {dataset}")
    return [fn(row) for row in rows]


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