app/utils/serialization.py

"""JSON serialization helpers for pandas, numpy, and datetime values."""

from __future__ import annotations

import json
import math
from datetime import date, datetime
from decimal import Decimal
from typing import Any

import numpy as np
import pandas as pd


def _normalize_text(value: Any) -> str:
    """Normalize dict keys / text values to valid UTF-8 strings."""
    if isinstance(value, bytes):
        value = value.decode("utf-8-sig", errors="replace")
    text = str(value)
    # Drop any bytes that cannot round-trip through UTF-8; prevents invalid
    # characters from reaching JSON serialization.
    return text.encode("utf-8", errors="ignore").decode("utf-8")


# Intentionally expose the helper for use by other modules.
__all__ = ["to_jsonable", "dataframe_to_records", "dumps_json", "loads_json", "_normalize_text"]


def to_jsonable(value: Any) -> Any:
    """Convert common data-science objects into strict JSON-compatible values."""
    if value is None:
        return None
    if isinstance(value, (str, bool, int)):
        return value
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if isinstance(value, Decimal):
        return float(value)
    if isinstance(value, (datetime, date, pd.Timestamp)):
        if pd.isna(value):
            return None
        return value.isoformat()
    if isinstance(value, np.generic):
        return to_jsonable(value.item())
    if isinstance(value, dict):
        return {_normalize_text(k): to_jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [to_jsonable(item) for item in value]
    if isinstance(value, pd.DataFrame):
        return dataframe_to_records(value)
    if isinstance(value, pd.Series):
        return to_jsonable(value.to_dict())
    if pd.isna(value):
        return None
    return _normalize_text(value)


def dataframe_to_records(df: pd.DataFrame, limit: int | None = None) -> list[dict[str, Any]]:
    """Return JSON-safe records while preserving AKShare's Chinese column names."""
    if df is None:
        return []
    work = df.head(limit).copy() if limit else df.copy()
    work.columns = [_normalize_text(c) for c in work.columns]
    return [to_jsonable(record) for record in work.to_dict(orient="records")]


def dumps_json(value: Any) -> str:
    return json.dumps(to_jsonable(value), ensure_ascii=False, separators=(",", ":"))


def loads_json(value: str) -> Any:
    return json.loads(value)