AlphaFutures/qihuo_analyzer/data/api_adapters/tqsdk_adapter.py

# TQSDK数据适配器
import os
import time
import pandas as pd
from typing import Dict, Optional, List
from qihuo_analyzer.data.api_adapters.base_adapter import BaseDataAdapter

# 尝试导入tqsdk

try:
    from tqsdk import TqApi, TqAuth
    TQSDK_AVAILABLE = True
except Exception as e:
    print(f"tqsdk导入失败：{e}，将使用模拟数据")
    TQSDK_AVAILABLE = False


class TqSdkAdapter(BaseDataAdapter):
    """TQSDK数据适配器
    
    使用天勤TQSDK获取期货数据。
    """
    
    def __init__(self):
        self.api = None
        # 交易所映射
        self.exchange_map = {
            'AU': 'SHFE',  # 黄金 - 上海期货交易所
            'AG': 'SHFE',  # 白银 - 上海期货交易所
            'CU': 'SHFE',  # 铜 - 上海期货交易所
            'NI': 'SHFE',  # 镍 - 上海期货交易所
            'SN': 'SHFE',  # 锡 - 上海期货交易所
        }
    
    def connect(self) -> bool:
        """连接API
        
        Returns:
            bool: 连接是否成功
        """
        try:
            if TQSDK_AVAILABLE:
                # 使用天勤TQSDK连接
                username = os.getenv('TQSDK_USERNAME', '')
                password = os.getenv('TQSDK_PASSWORD', '')
                
                if username and password:
                    self.api = TqApi(auth=TqAuth(username, password))
                    print("TQSDK API连接成功")
                    return True
                else:
                    print("TQSDK账号密码未配置，将使用模拟数据")
                    self.api = None
                    return False
            else:
                # 模拟API，用于测试
                print("TQSDK不可用，使用模拟API")
                self.api = None
                return False
        except Exception as e:
            print(f"TQSDK API连接失败：{e}")
            # 模拟API，用于测试
            self.api = None
            return False
    
    def disconnect(self):
        """断开连接"""
        if self.api:
            try:
                self.api.close()
                print("TQSDK API连接已断开")
            except:
                pass
    
    def _convert_duration(self, duration: str) -> int:
        """将时间周期字符串转换为分钟数
        
        Args:
            duration: 时间周期，如 '1m', '5m', '15m', '1h', '1d'
            
        Returns:
            分钟数
        """
        duration_map = {
            '1m': 1,
            '5m': 5,
            '15m': 15,
            '30m': 30,
            '1h': 60,
            '2h': 120,
            '4h': 240,
            '6h': 360,
            '12h': 720,
            '1d': 1440,
            '1w': 10080
        }
        return duration_map.get(duration, 60)  # 默认60分钟
    
    def _convert_symbol(self, symbol: str) -> str:
        """将合约代码转换为TQSDK格式
        
        Args:
            symbol: 合约代码，如 'CU2603'
            
        Returns:
            TQSDK格式的合约代码，如 'SHFE.cu2603'
        """
        # 提取品种代码和合约月份
        if len(symbol) >= 4:
            # 3字符品种代码
            if len(symbol) >= 5:
                product_code = symbol[:3].upper()
                if product_code in self.exchange_map:
                    contract_month = symbol[3:].lower()
                    exchange = self.exchange_map[product_code]
                    return f"{exchange}.{product_code.lower()}{contract_month}"
            
            # 2字符品种代码
            product_code = symbol[:2].upper()
            if product_code in self.exchange_map:
                contract_month = symbol[2:].lower()
                exchange = self.exchange_map[product_code]
                return f"{exchange}.{product_code.lower()}{contract_month}"
            
            # 1字符品种代码
            product_code = symbol[:1].upper()
            if product_code in self.exchange_map:
                contract_month = symbol[1:].lower()
                exchange = self.exchange_map[product_code]
                return f"{exchange}.{product_code.lower()}{contract_month}"
            
            # 无法识别的合约代码，返回原始代码
            return symbol
        else:
            return symbol
    
    def get_kline_data(self, symbol: str, duration: str, count: int = 200) -> Optional[pd.DataFrame]:
        """获取K线数据
        
        Args:
            symbol: 合约代码
            duration: 时间周期，如 '1m', '5m', '15m', '1h', '1d'
            count: 数据数量
            
        Returns:
            K线数据DataFrame，如果无法获取真实数据则返回None
        """
        print(f"[TqSdkAdapter]获取K线数据: {symbol}, {duration}, {count}")
        try:
            if TQSDK_AVAILABLE and self.api:
                # 转换合约代码为TQSDK格式
                tq_symbol = self._convert_symbol(symbol)
                print(f"使用TQSDK格式合约代码: {tq_symbol}")
                
                # 转换时间周期为分钟数
                duration_minutes = self._convert_duration(duration)
                # 使用真实API获取数据
                klines = self.api.get_kline_serial(tq_symbol, duration_minutes, data_length=count)
                
                # 等待数据准备就绪
                import time
                start_time = time.time()
                timeout = 5  # 5秒超时
                
                while True:
                    if hasattr(klines, 'datetime') and len(klines.datetime) > 0:
                        break
                    if time.time() - start_time > timeout:
                        print("获取K线数据超时")
                        return None
                    time.sleep(0.1)
                
                # 转换为DataFrame
                data = {
                    'datetime': klines.datetime,
                    'open': klines.open,
                    'high': klines.high,
                    'low': klines.low,
                    'close': klines.close,
                    'volume': klines.volume,
                    'open_interest': klines.open_oi
                }
                df = pd.DataFrame(data)
                df['datetime'] = pd.to_datetime(df['datetime'], unit='ns')
                df.set_index('datetime', inplace=True)
                
                print(f"成功获取K线数据，数据长度: {len(df)}")
                return df
            else:
                # 不再自动返回模拟数据，返回None
                print(f"无法获取真实数据：{'API未连接' if not self.api else 'TQSDK不可用'}")
                return None
        except Exception as e:
            print(f"获取K线数据失败：{e}")
            # 不再自动返回模拟数据，返回None
            return None
    
    def get_tick_data(self, symbol: str, count: int = 1000) -> Optional[pd.DataFrame]:
        """获取Tick数据"""
        try:
            if TQSDK_AVAILABLE and self.api:
                # 使用真实API获取数据
                ticks = self.api.get_tick_serial(symbol, data_length=count)
                self.api.wait_update()
                
                # 转换为DataFrame
                data = {
                    'datetime': ticks.datetime,
                    'last_price': ticks.last_price,
                    'volume': ticks.volume,
                    'open_interest': ticks.open_interest,
                    'bid_price1': ticks.bid_price1,
                    'bid_volume1': ticks.bid_volume1,
                    'ask_price1': ticks.ask_price1,
                    'ask_volume1': ticks.ask_volume1
                }
                df = pd.DataFrame(data)
                df['datetime'] = pd.to_datetime(df['datetime'], unit='ns')
                df.set_index('datetime', inplace=True)
                
                return df
            else:
                # 返回模拟数据
                print(f"无法获取真实数据：{'API未连接' if not self.api else 'TQSDK不可用'}")
                return None
        except Exception as e:
            print(f"获取Tick数据失败：{e}")
            return None
    
    def get_contract_info(self, symbol: str) -> Optional[Dict]:
        """获取合约信息"""
        try:
            if TQSDK_AVAILABLE and self.api:
                # 使用真实API获取数据
                quote = self.api.get_quote(symbol)
                self.api.wait_update()
                
                return {
                    'symbol': symbol,
                    'name': quote.instrument_name,
                    'exchange': quote.exchange_id,
                    'product': quote.product_id,
                    'price_tick': quote.price_tick,
                    'volume_multiple': quote.volume_multiple,
                    'margin_rate': quote.margin_rate,
                    'expire_datetime': quote.expire_datetime,
                    'create_datetime': quote.create_datetime
                }
            else:
                # 返回模拟数据
                print(f"无法获取真实数据：{'API未连接' if not self.api else 'TQSDK不可用'}")
                return None
        except Exception as e:
            print(f"获取合约信息失败：{e}")
            return None
    
    def get_market_data(self, symbols: List[str]) -> Dict[str, Dict]:
        """批量获取市场数据"""
        market_data = {}
        
        for symbol in symbols:
            try:
                if TQSDK_AVAILABLE and self.api:
                    quote = self.api.get_quote(symbol)
                    self.api.wait_update()
                    
                    market_data[symbol] = {
                        'latest_price': quote.last_price,
                        'open': quote.open,
                        'high': quote.high,
                        'low': quote.low,
                        'pre_close': quote.pre_close,
                        'volume': quote.volume,
                        'open_interest': quote.open_interest,
                        'bid_price1': quote.bid_price1,
                        'ask_price1': quote.ask_price1
                    }
                else:
                    # 模拟数据
                    market_data[symbol] = {
                        'latest_price': 0,
                        'open': 0,
                        'high': 0,
                        'low': 0,
                        'pre_close': 0,
                        'volume': 0,
                        'open_interest': 0,
                        'bid_price1': 0,
                        'ask_price1': 0
                    }
            except Exception as e:
                print(f"获取{symbol}市场数据失败：{e}")
                market_data[symbol] = {
                    'latest_price': 0,
                    'open': 0,
                    'high': 0,
                    'low': 0,
                    'pre_close': 0,
                    'volume': 0,
                    'open_interest': 0,
                    'bid_price1': 0,
                    'ask_price1': 0
                }
        
        return market_data
    
    def get_all_symbols(self) -> List[str]:
        """获取所有品种列表
        
        Returns:
            List[str]: 所有品种的合约代码列表
        """
        try:
            if TQSDK_AVAILABLE and self.api:
                # TQSDK 没有直接获取所有品种列表的方法，使用模拟数据
                print("TQSDK 不支持获取所有品种列表，使用模拟数据")
                return self._get_mock_all_symbols()
            else:
                # 返回模拟数据
                print("使用模拟品种列表")
                return self._get_mock_all_symbols()
        except Exception as e:
            print(f"获取所有品种列表失败：{e}")
            return self._get_mock_all_symbols()
    
    def _get_mock_all_symbols(self) -> List[str]:
        """获取模拟品种列表"""
        # 返回exchange_map中映射的所有品种
        symbols = []
        # 为每个品种生成一个合约代码（使用2603月份）
        for product_code in self.exchange_map:
            # 生成合约代码，格式：品种代码+2603
            contract_code = f"{product_code}2603"
            symbols.append(contract_code)
        print(f"模拟品种列表: {symbols}")
        return symbols