myquant/finance/data/sources/tushare_source.py

"""
Tushare 数据源封装。

与 AkShareSource 保持相同接口，作为平替/备份数据源。
"""

import os
import time

import pandas as pd
import tushare as ts

# 确保 .env 已加载（无论从哪个路径导入）
import config.settings  # noqa: F401


class TushareSource:
    """Tushare 数据源。"""

    def __init__(self, token: str | None = None):
        token = token or os.getenv("TUSHARE_TOKEN", "")
        if not token:
            self._pro = None
        else:
            ts.set_token(token)
            self._pro = ts.pro_api()

    @property
    def available(self) -> bool:
        return self._pro is not None

    def _ensure_pro(self):
        if self._pro is None:
            raise RuntimeError("Tushare token 未配置，请在 .env 中设置 TUSHARE_TOKEN")

    # ── 股票列表 ──────────────────────────────────────────

    def fetch_stock_list(self) -> pd.DataFrame:
        """获取 A 股股票列表。"""
        self._ensure_pro()
        fields = "ts_code,name,area,industry,market,list_date,is_hs"
        df = self._pro.stock_basic(
            exchange="", list_status="L",
            fields=fields,
        )
        if df is None or df.empty:
            return pd.DataFrame()
        return df[["ts_code", "name"]]

    # ── 日线数据 ──────────────────────────────────────────

    def fetch_daily(
        self, ts_code: str, start: str, end: str | None = None
    ) -> pd.DataFrame:
        """
        获取单只股票日线行情。

        参数:
            ts_code: 如 '000001.SZ'
            start: 'YYYYMMDD'
            end: 'YYYYMMDD'
        """
        self._ensure_pro()
        end = end or time.strftime("%Y%m%d")

        df = self._pro.daily(
            ts_code=ts_code,
            start_date=start,
            end_date=end,
            fields="ts_code,trade_date,open,high,low,close,pre_close,change,pct_chg,vol,amount,turnover_rate",
        )
        if df is None or df.empty:
            return pd.DataFrame()

        # 前复权因子
        try:
            adj = self._pro.adj_factor(ts_code=ts_code, start_date=start, end_date=end)
            if adj is not None and not adj.empty:
                df = df.merge(adj[["trade_date", "adj_factor"]], on="trade_date", how="left")
                for col in ["open", "high", "low", "close", "pre_close"]:
                    if col in df.columns and "adj_factor" in df.columns:
                        df[col] = (df[col] * df["adj_factor"]).round(2)
                df = df.drop(columns=["adj_factor"])
        except Exception:
            pass

        df["trade_date"] = df["trade_date"].astype(str)
        return df

    # ── 指数日线 ──────────────────────────────────────────

    def fetch_index_daily(
        self, ts_code: str, start: str, end: str | None = None
    ) -> pd.DataFrame:
        """
        获取指数日线行情。

        Tushare index_daily 接口。
        """
        self._ensure_pro()
        end = end or time.strftime("%Y%m%d")
        df = self._pro.index_daily(
            ts_code=ts_code,
            start_date=start, end_date=end,
            fields="ts_code,trade_date,open,high,low,close,pre_close,change,pct_chg,vol,amount",
        )
        if df is None or df.empty:
            return pd.DataFrame()
        df["trade_date"] = df["trade_date"].astype(str)
        return df

    # ── 财务数据 ──────────────────────────────────────────

    def fetch_financial(self, ts_code: str) -> pd.DataFrame:
        """获取单只股票核心财务指标。"""
        self._ensure_pro()
        fields = (
            "ts_code,end_date,roe,roa,grossprofit_margin,netprofit_margin,"
            "debt_to_assets,eps,dt_eps,bps,pe,pb,"
            "total_revenue,revenue_yoy,n_income,n_income_yoy"
        )
        df = self._pro.fina_indicator(
            ts_code=ts_code,
            fields=fields,
        )
        if df is None or df.empty:
            return pd.DataFrame()

        df = df.rename(columns={
            "grossprofit_margin": "gross_profit_margin",
            "netprofit_margin": "net_profit_margin",
            "dt_eps": "eps_diluted",
            "bps": "bvps",
            "revenue_yoy": "total_revenue_yoy",
            "n_income": "net_profit",
            "n_income_yoy": "net_profit_yoy",
        })
        df["end_date"] = df["end_date"].astype(str)
        return df