myquant/djapi/api/stock/data_source.py

"""
统一数据源模块 — djapi/api/stock/ 的单一数据入口。

归一化 Tushare / AkShare / MySQL 三种数据源。
所有模块通过此入口获取数据连接，不再各自创建 pro 实例。

用法:
    from .data_source import get_tushare_pro, get_mysql_db

    pro = get_tushare_pro()
    df = pro.daily(ts_code='000001.SZ', ...)

    db = get_mysql_db()
    rows = db.query("SELECT * FROM xwlb_daily WHERE ...")

扩展新数据源:
    class NewSource: ...
    _sources['new'] = NewSource()
    def get_new_source(): return _sources['new']
"""

import os
import threading

import tushare as ts

# 延迟加载 AkShare（避免不必要的导入开销）
_akshare = None


def _get_akshare():
    global _akshare
    if _akshare is None:
        import akshare as ak
        _akshare = ak
    return _akshare


# ═══════════════════════════════════════════════════════════
# Token 加载
# ═══════════════════════════════════════════════════════════

def get_ts_token() -> str:
    """获取 Tushare Token，优先级：环境变量 TUSHARE_TS_TOKEN > TUSHARE_TOKEN > config.py"""
    token = os.getenv("TUSHARE_TS_TOKEN", "") or os.getenv("TUSHARE_TOKEN", "")
    if not token:
        try:
            from .config import TS_TOKEN
            token = TS_TOKEN
        except ImportError:
            try:
                from config import TS_TOKEN
                token = TS_TOKEN
            except ImportError:
                pass
    return token


# ═══════════════════════════════════════════════════════════
# Tushare Pro 连接池（线程安全单例）
# ═══════════════════════════════════════════════════════════

_pro_lock = threading.Lock()
_pro = None


def get_tushare_pro():
    """获取 Tushare pro_api 实例（全局单例，线程安全）。"""
    global _pro
    if _pro is not None:
        return _pro
    with _pro_lock:
        if _pro is not None:
            return _pro
        token = get_ts_token()
        ts.set_token(token)
        _pro = ts.pro_api()
    return _pro


def reset_tushare_pro():
    """重置 Tushare 连接（token 变更时调用）。"""
    global _pro
    with _pro_lock:
        _pro = None


# ═══════════════════════════════════════════════════════════
# MySQL 连接
# ═══════════════════════════════════════════════════════════

_mysql_db = None


def get_mysql_db():
    """获取 MySQLDB 实例（全局单例）。"""
    global _mysql_db
    if _mysql_db is not None:
        return _mysql_db
    try:
        from ..utils.mysql_handler import MySQLDB
    except (ImportError, ValueError):
        try:
            from utils.mysql_handler import MySQLDB
        except ImportError:
            return None
    _mysql_db = MySQLDB()
    return _mysql_db


# ═══════════════════════════════════════════════════════════
# 日线行情 — 双源 fallback（Tushare → AkShare）
# ═══════════════════════════════════════════════════════════

def get_daily(ts_code: str, start_date: str, end_date: str, source: str = "tushare"):
    """
    获取个股日线行情。

    参数:
        ts_code: 如 '000001.SZ'
        start_date: YYYYMMDD
        end_date: YYYYMMDD
        source: 'tushare' | 'akshare' | 'auto' (tushare优先)

    返回:
        pd.DataFrame (trade_date, open, high, low, close, vol, amount, ...)
    """
    import pandas as pd

    if source == "auto":
        # Tushare 优先
        try:
            df = get_daily(ts_code, start_date, end_date, source="tushare")
            if df is not None and not df.empty:
                return df
        except Exception:
            pass
        return get_daily(ts_code, start_date, end_date, source="akshare")

    if source == "tushare":
        pro = get_tushare_pro()
        df = pro.daily(
            ts_code=ts_code, start_date=start_date, end_date=end_date,
            fields="ts_code,trade_date,open,high,low,close,pre_close,change,pct_chg,vol,amount"
        )
        if df is not None and not df.empty:
            df["trade_date"] = df["trade_date"].astype(str)
        return df

    if source == "akshare":
        symbol = ts_code.replace(".SZ", "").replace(".SH", "").replace(".BJ", "")
        ak = _get_akshare()
        df = ak.stock_zh_a_hist(
            symbol=symbol, period="daily",
            start_date=start_date, end_date=end_date, adjust="qfq"
        )
        if df is not None and not df.empty:
            df = df.rename(columns={
                "日期": "trade_date", "开盘": "open", "收盘": "close",
                "最高": "high", "最低": "low", "成交量": "vol", "成交额": "amount",
                "涨跌幅": "pct_chg", "涨跌额": "change",
            })
            df["ts_code"] = ts_code
            df["trade_date"] = df["trade_date"].astype(str)
        return df if df is not None else pd.DataFrame()

    raise ValueError("Unknown source: {}".format(source))


# ═══════════════════════════════════════════════════════════
# 扩展点：未来新增数据源
# ═══════════════════════════════════════════════════════════
#
# 1. 在 _sources dict 中注册新源
# 2. 实现与 get_daily() 相同签名的函数
# 3. 在 get_daily(source=...) 中添加路由
#
# _sources = {
#     "tushare": TushareDailySource(),
#     "akshare": AkShareDailySource(),
#     "wind": WindDailySource(),       # 未来
#     "joinquant": JoinQuantSource(),  # 未来
# }