tushare_api.py

"""
QTrading Tushare 数据获取模块
"""

import pandas as pd
import tushare as ts
from typing import Optional, List
from config.config import config
from utils.logger import get_logger

logger = get_logger(__name__)

class TushareDataFetcher:
    """Tushare 数据获取器"""

    def __init__(self, token: str = None):
        """
        初始化 Tushare 数据获取器

        Args:
            token: Tushare 接口令牌（可选，默认从配置文件读取）
        """
        self.token = token or config.tushare.token
        self._pro = None
        self._init_tushare()

    def _init_tushare(self) -> None:
        """初始化 Tushare 接口"""
        try:
            ts.set_token(self.token)
            self._pro = ts.pro_api()
            logger.info("Tushare 接口初始化成功")
        except Exception as e:
            logger.error(f"Tushare 接口初始化失败: {e}")
            self._pro = None

    def get_daily_kline(self, ts_code: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """
        获取股票日线数据

        Args:
            ts_code: 股票代码（如 '000001.SZ'）
            start_date: 开始日期（YYYYMMDD）
            end_date: 结束日期（YYYYMMDD）

        Returns:
            pd.DataFrame: 日线数据（包含 ts_code, trade_date, open, high, low, close, vol, amount）
        """
        if self._pro is None:
            logger.error("Tushare 接口未初始化")
            return None

        try:
            logger.info(f"正在获取股票 {ts_code} 日线数据 ({start_date} - {end_date})")
            data = self._pro.daily(
                ts_code=ts_code,
                start_date=start_date,
                end_date=end_date
            )
            if data is not None and not data.empty:
                logger.info(f"成功获取 {len(data)} 条数据")
                return data
            else:
                logger.warning(f"未找到股票 {ts_code} 在指定日期范围内的数据")
                return None
        except Exception as e:
            logger.error(f"获取股票 {ts_code} 日线数据失败: {e}")
            return None

    def get_adj_factor(self, ts_code: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """
        获取股票复权因子

        Args:
            ts_code: 股票代码（如 '000001.SZ'）
            start_date: 开始日期（YYYYMMDD）
            end_date: 结束日期（YYYYMMDD）

        Returns:
            pd.DataFrame: 复权因子数据（包含 ts_code, trade_date, adj_factor）
        """
        if self._pro is None:
            logger.error("Tushare 接口未初始化")
            return None

        try:
            logger.info(f"正在获取股票 {ts_code} 复权因子 ({start_date} - {end_date})")
            data = self._pro.adj_factor(
                ts_code=ts_code,
                start_date=start_date,
                end_date=end_date
            )
            if data is not None and not data.empty:
                logger.info(f"成功获取 {len(data)} 条复权因子数据")
                return data
            else:
                logger.warning(f"未找到股票 {ts_code} 在指定日期范围内的复权因子")
                return None
        except Exception as e:
            logger.error(f"获取股票 {ts_code} 复权因子失败: {e}")
            return None

    def get_stock_basic(self, market: str = None) -> Optional[pd.DataFrame]:
        """
        获取股票基本信息

        Args:
            market: 市场类型（可选，如 'SZ' 或 'SH'）

        Returns:
            pd.DataFrame: 股票基本信息
        """
        if self._pro is None:
            logger.error("Tushare 接口未初始化")
            return None

        try:
            logger.info("正在获取股票基本信息")
            data = self._pro.stock_basic(exchange='', list_status='L')
            if market:
                data = data[data['market'] == market]
            logger.info(f"成功获取 {len(data)} 只股票基本信息")
            return data
        except Exception as e:
            logger.error(f"获取股票基本信息失败: {e}")
            return None

    def get_index_daily(self, ts_code: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """
        获取指数日线数据

        Args:
            ts_code: 指数代码（如 '000001.SH'）
            start_date: 开始日期（YYYYMMDD）
            end_date: 结束日期（YYYYMMDD）

        Returns:
            pd.DataFrame: 指数日线数据
        """
        if self._pro is None:
            logger.error("Tushare 接口未初始化")
            return None

        try:
            logger.info(f"正在获取指数 {ts_code} 日线数据 ({start_date} - {end_date})")
            data = self._pro.index_daily(
                ts_code=ts_code,
                start_date=start_date,
                end_date=end_date
            )
            if data is not None and not data.empty:
                logger.info(f"成功获取 {len(data)} 条指数数据")
                return data
            else:
                logger.warning(f"未找到指数 {ts_code} 在指定日期范围内的数据")
                return None
        except Exception as e:
            logger.error(f"获取指数 {ts_code} 日线数据失败: {e}")
            return None

    def get_trade_cal(self, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """
        获取交易日历

        Args:
            start_date: 开始日期（YYYYMMDD）
            end_date: 结束日期（YYYYMMDD）

        Returns:
            pd.DataFrame: 交易日历
        """
        if self._pro is None:
            logger.error("Tushare 接口未初始化")
            return None

        try:
            logger.info(f"正在获取交易日历 ({start_date} - {end_date})")
            data = self._pro.trade_cal(
                start_date=start_date,
                end_date=end_date
            )
            logger.info(f"成功获取 {len(data)} 个交易日")
            return data
        except Exception as e:
            logger.error(f"获取交易日历失败: {e}")
            return None