path: root/src/FundLab.Api/akshare_collector.py

#!/usr/bin/env python3
"""Small raw-data adapter for the AKShare endpoints used by Fund Lab."""

import argparse
import datetime as dt
import json
import re
import sys
from decimal import Decimal, InvalidOperation

import akshare as ak
import pandas as pd


SCHEMA_VERSION = "fund-lab.akshare.v1"


def collected_at():
    return dt.datetime.now(dt.timezone.utc).isoformat().replace("+00:00", "Z")


def source_revision():
    return f"akshare-{getattr(ak, '__version__', 'unknown')}/eastmoney"


def is_missing(value):
    try:
        return bool(pd.isna(value))
    except (TypeError, ValueError):
        return False


def text(value):
    if is_missing(value):
        return None
    value = str(value).strip()
    return value or None


def fund_code(value):
    value = text(value)
    if value is None:
        return None
    if value.isdigit() and len(value) < 6:
        value = value.zfill(6)
    return value if re.fullmatch(r"\d{6}", value) else None


def decimal_text(value):
    value = text(value)
    if value is None:
        return None
    try:
        return format(Decimal(value), "f")
    except InvalidOperation:
        return None


def date_text(value):
    if is_missing(value):
        return None
    if isinstance(value, (dt.datetime, dt.date)):
        return value.strftime("%Y-%m-%d")
    parsed = pd.to_datetime(value, errors="coerce")
    if pd.isna(parsed):
        return None
    return parsed.strftime("%Y-%m-%d")


def search(query):
    query = query.strip()
    if not query:
        raise ValueError("search query cannot be empty")

    frame = ak.fund_name_em()
    query_lower = query.casefold()
    rows = []

    for _, row in frame.iterrows():
        code = fund_code(row.get("基金代码"))
        name = text(row.get("基金简称"))
        pinyin = text(row.get("拼音缩写"))
        full_pinyin = text(row.get("拼音全称"))

        if code is None or name is None:
            continue

        searchable = [code, name, pinyin or "", full_pinyin or ""]
        if not any(query_lower in value.casefold() for value in searchable):
            continue

        rows.append(
            {
                "code": code,
                "name": name,
                "fund_type": text(row.get("基金类型")),
                "rank": 0 if code == query else (1 if name == query else 2),
            }
        )

    rows.sort(key=lambda item: (item.pop("rank"), item["code"]))
    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "search",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instruments": rows[:100],
    }


def nav(code):
    if not re.fullmatch(r"\d{6}", code):
        raise ValueError("fund code must contain exactly six digits")

    frame = ak.fund_open_fund_info_em(
        symbol=code,
        indicator="单位净值走势",
        period="成立来",
    )
    observations = []

    for _, row in frame.iterrows():
        nav_date = date_text(row.get("净值日期"))
        nav_value = decimal_text(row.get("单位净值"))
        if nav_date is None or nav_value is None:
            continue

        observations.append(
            {
                "nav_date": nav_date,
                "published_at": None,
                "nav": nav_value,
                "accumulated_nav": None,
                "daily_return": decimal_text(row.get("日增长率")),
            }
        )

    if not observations:
        raise ValueError(f"AKShare returned no usable NAV observations for {code}")

    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "nav",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instrument": {"code": code},
        "observations": observations,
    }


def bond_maturity_date(code):
    # Maturity is a best-effort enrichment: the quote endpoint has no expiry
    # date, so a failure here only means maturity_date stays null.
    try:
        frame = ak.bond_zh_cov_info(symbol=code, indicator="基本信息")
        if frame is None or frame.empty:
            return None
        return date_text(frame.iloc[0].get("EXPIRE_DATE"))
    except Exception:
        return None


def bond_quote(code):
    code = text(code)
    if code is None or not re.fullmatch(r"\d{6}", code):
        raise ValueError("bond code must contain exactly six digits")

    frame = ak.bond_zh_hs_cov_spot()
    if frame is None or frame.empty:
        raise ValueError(f"no valid quote for {code}")

    prefix_rank = {"sh": 0, "sz": 1, "bj": 2}
    candidates = []

    for _, row in frame.iterrows():
        symbol = text(row.get("symbol")) or ""
        row_code = text(row.get("code")) or ""
        if row_code != code and not symbol.endswith(code):
            continue

        price = decimal_text(row.get("trade"))
        if price is None or Decimal(price) <= 0:
            continue

        prefix = symbol[:2].lower()
        candidates.append((prefix_rank.get(prefix, 3), symbol, text(row.get("name")), price))

    if not candidates:
        raise ValueError(f"no valid quote for {code}")

    _, _, name, price = sorted(candidates, key=lambda item: (item[0], item[1]))[0]

    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "bond-quote",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instrument": {"code": code},
        "quote": {
            "name": name,
            "price": price,
            "maturity_date": bond_maturity_date(code),
        },
    }


def first_positive_decimal(*values):
    for value in values:
        price = decimal_text(value)
        if price is not None and Decimal(price) > 0:
            return price
    return None


def stock_price_from_bid_ask(code):
    # Single-symbol realtime snapshot (eastmoney). Returns None when the venue
    # rejects the symbol or carries no valid last price.
    frame = ak.stock_bid_ask_em(symbol=code)
    if frame is None or frame.empty:
        return None

    for _, row in frame.iterrows():
        item = text(row.get("item"))
        if item in ("最新", "最新价"):
            return first_positive_decimal(row.get("value"))

    return None


def stock_from_snapshot(code):
    # Whole-market snapshots differ per provider; each entry is
    # (fetch, code column, name column, price column, exchange prefixes).
    providers = [
        (ak.stock_zh_a_spot_tx, "code", "name", "zxj", {"sh": 0, "sz": 1, "bj": 2}),
        (ak.stock_zh_a_spot_em, "代码", "名称", "最新价", None),
    ]

    for fetch, code_key, name_key, price_key, prefixes in providers:
        try:
            frame = fetch()
        except Exception:
            continue

        if frame is None or frame.empty:
            continue

        candidates = []

        for _, row in frame.iterrows():
            row_code = text(row.get(code_key)) or ""
            if row_code != code and not row_code.lower().endswith(code):
                continue

            price = first_positive_decimal(row.get(price_key))
            if price is None:
                continue

            prefix = row_code[:2].lower()
            rank = prefixes.get(prefix, 3) if prefixes is not None else 0
            candidates.append((rank, row_code, text(row.get(name_key)), price))

        if candidates:
            _, _, name, price = sorted(candidates, key=lambda item: (item[0], item[1]))[0]
            return name, price

    return None


def stock_quote(code):
    code = text(code)
    if code is None or not re.fullmatch(r"\d{6}", code):
        raise ValueError("stock code must contain exactly six digits")

    name = None
    price = None

    try:
        price = stock_price_from_bid_ask(code)
    except Exception:
        price = None

    if price is None:
        found = stock_from_snapshot(code)
        if found is None:
            raise ValueError(f"no valid quote for {code}")
        name, price = found

    return {
        "schema_version": SCHEMA_VERSION,
        "operation": "stock-quote",
        "source": "akshare",
        "source_revision": source_revision(),
        "collected_at": collected_at(),
        "instrument": {"code": code},
        "quote": {
            "name": name,
            "price": price,
            "currency": "CNY",
        },
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--operation", choices=("search", "nav", "bond-quote", "stock-quote"), required=True)
    parser.add_argument("--query")
    parser.add_argument("--code")
    args = parser.parse_args()

    try:
        if args.operation == "search":
            payload = search(args.query)
        elif args.operation == "nav":
            payload = nav(args.code)
        elif args.operation == "bond-quote":
            payload = bond_quote(args.code)
        else:
            payload = stock_quote(args.code)

        json.dump(payload, sys.stdout, ensure_ascii=False, separators=(",", ":"))
        sys.stdout.write("\n")
        return 0
    except Exception as error:
        print(f"AKShare collector failed: {error}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())