|
|
import tushare as ts
|
|
|
import pandas as pd
|
|
|
from typing import List, Optional
|
|
|
from datetime import datetime
|
|
|
from app.services.datasource.base import DataSourceBase
|
|
|
from app.config import settings
|
|
|
|
|
|
|
|
|
class TushareSource(DataSourceBase):
|
|
|
"""Tushare 数据源适配器"""
|
|
|
|
|
|
def __init__(self, config: dict):
|
|
|
super().__init__(config)
|
|
|
self.name = "tushare"
|
|
|
self.pro = None
|
|
|
self._token = config.get("token", settings.TUSHARE_TOKEN)
|
|
|
|
|
|
def initialize(self) -> bool:
|
|
|
"""初始化 Tushare 连接"""
|
|
|
try:
|
|
|
ts.set_token(self._token)
|
|
|
self.pro = ts.pro_api()
|
|
|
# 简单测试连接
|
|
|
self.pro.trade_cal(exchange="DCE", start_date="20240101", end_date="20240105")
|
|
|
self._initialized = True
|
|
|
return True
|
|
|
except Exception as e:
|
|
|
self._initialized = False
|
|
|
raise e
|
|
|
|
|
|
def _format_date(self, date_str: str) -> str:
|
|
|
"""将 YYYY-MM-DD 转换为 YYYYMMDD"""
|
|
|
return date_str.replace("-", "")
|
|
|
|
|
|
def get_contract_list(self, exchange: Optional[str] = None) -> List[dict]:
|
|
|
"""获取期货合约列表"""
|
|
|
if not self._initialized:
|
|
|
self.initialize()
|
|
|
|
|
|
# Tushare 获取合约信息
|
|
|
df = self.pro.fut_basic(
|
|
|
exchange=exchange or "CFFEX", # 需要分别查询每个交易所
|
|
|
fut_type="1", # 1=标准合约
|
|
|
fut_series=""
|
|
|
)
|
|
|
|
|
|
results = []
|
|
|
if df is not None and not df.empty:
|
|
|
for _, row in df.iterrows():
|
|
|
results.append({
|
|
|
"symbol": row.get("symbol", ""),
|
|
|
"exchange": self._map_exchange(row.get("exchange", "")),
|
|
|
"name": row.get("name", ""),
|
|
|
"product": row.get("underlying_symbol", ""),
|
|
|
"multiplier": int(row.get("contract_multiplier", 10)) if row.get("contract_multiplier") else 10,
|
|
|
"price_tick": float(row.get("price_tick", 0)) if row.get("price_tick") else None,
|
|
|
"expire_date": self._parse_date(row.get("delist_date", "")),
|
|
|
"is_active": row.get("list_status", "") == "L",
|
|
|
})
|
|
|
|
|
|
return results
|
|
|
|
|
|
def _map_exchange(self, exchange: str) -> str:
|
|
|
"""交易所代码映射"""
|
|
|
mapping = {
|
|
|
"CFFEX": "CFFEX",
|
|
|
"SHFE": "SHFE",
|
|
|
"DCE": "DCE",
|
|
|
"CZCE": "ZCE", # Tushare 用 CZCE,我们统一用 ZCE
|
|
|
"INE": "INE",
|
|
|
"GFEX": "GFEX",
|
|
|
}
|
|
|
return mapping.get(exchange, exchange)
|
|
|
|
|
|
def _parse_date(self, date_str: str) -> Optional[datetime]:
|
|
|
"""解析日期字符串"""
|
|
|
if not date_str:
|
|
|
return None
|
|
|
try:
|
|
|
return datetime.strptime(str(date_str), "%Y%m%d")
|
|
|
except Exception:
|
|
|
return None
|
|
|
|
|
|
def get_kline_daily(
|
|
|
self,
|
|
|
symbol: str,
|
|
|
start_date: str,
|
|
|
end_date: str
|
|
|
) -> pd.DataFrame:
|
|
|
"""获取日K线数据"""
|
|
|
if not self._initialized:
|
|
|
self.initialize()
|
|
|
|
|
|
start = self._format_date(start_date)
|
|
|
end = self._format_date(end_date)
|
|
|
|
|
|
df = self.pro.fut_daily(
|
|
|
ts_code=symbol,
|
|
|
start_date=start,
|
|
|
end_date=end
|
|
|
)
|
|
|
|
|
|
if df is None or df.empty:
|
|
|
return pd.DataFrame()
|
|
|
|
|
|
# 统一列名
|
|
|
df = df.rename(columns={
|
|
|
"trade_date": "trade_date",
|
|
|
"open": "open",
|
|
|
"high": "high",
|
|
|
"low": "low",
|
|
|
"close": "close",
|
|
|
"vol": "volume",
|
|
|
"amount": "turnover",
|
|
|
"oi": "open_interest",
|
|
|
"settle": "settle",
|
|
|
"pre_settle": "pre_settle",
|
|
|
})
|
|
|
|
|
|
df["trade_date"] = pd.to_datetime(df["trade_date"])
|
|
|
return df[["trade_date", "open", "high", "low", "close", "volume", "turnover", "open_interest", "settle", "pre_settle"]]
|
|
|
|
|
|
def get_kline_weekly(
|
|
|
self,
|
|
|
symbol: str,
|
|
|
start_date: str,
|
|
|
end_date: str
|
|
|
) -> pd.DataFrame:
|
|
|
"""
|
|
|
获取周K线数据
|
|
|
Tushare 没有直接的周K接口,通过日K聚合
|
|
|
"""
|
|
|
daily_df = self.get_kline_daily(symbol, start_date, end_date)
|
|
|
if daily_df.empty:
|
|
|
return pd.DataFrame()
|
|
|
|
|
|
# 按周聚合
|
|
|
daily_df = daily_df.set_index("trade_date")
|
|
|
weekly = daily_df.resample("W-FRI").agg({
|
|
|
"open": "first",
|
|
|
"high": "max",
|
|
|
"low": "min",
|
|
|
"close": "last",
|
|
|
"volume": "sum",
|
|
|
"turnover": "sum",
|
|
|
"open_interest": "last",
|
|
|
}).dropna()
|
|
|
|
|
|
weekly = weekly.reset_index()
|
|
|
weekly = weekly.rename(columns={"trade_date": "trade_date"})
|
|
|
return weekly
|
|
|
|
|
|
def get_kline_intraday(
|
|
|
self,
|
|
|
symbol: str,
|
|
|
period: str,
|
|
|
start_date: str,
|
|
|
end_date: str
|
|
|
) -> pd.DataFrame:
|
|
|
"""
|
|
|
获取分钟级K线数据
|
|
|
Tushare 的 fut_mins 接口
|
|
|
"""
|
|
|
if not self._initialized:
|
|
|
self.initialize()
|
|
|
|
|
|
# Tushare fut_mins 接口
|
|
|
start = self._format_date(start_date)
|
|
|
end = self._format_date(end_date)
|
|
|
|
|
|
# 分钟数映射
|
|
|
freq_map = {"5m": "5", "15m": "15", "30m": "30", "60m": "60"}
|
|
|
freq = freq_map.get(period, "5")
|
|
|
|
|
|
try:
|
|
|
df = self.pro.fut_mins(
|
|
|
ts_code=symbol,
|
|
|
freq=freq,
|
|
|
start_date=start,
|
|
|
end_date=end
|
|
|
)
|
|
|
except Exception:
|
|
|
# 部分交易所可能不支持分钟数据
|
|
|
return pd.DataFrame()
|
|
|
|
|
|
if df is None or df.empty:
|
|
|
return pd.DataFrame()
|
|
|
|
|
|
df = df.rename(columns={
|
|
|
"trade_time": "trade_time",
|
|
|
"open": "open",
|
|
|
"high": "high",
|
|
|
"low": "low",
|
|
|
"close": "close",
|
|
|
"vol": "volume",
|
|
|
"amount": "turnover",
|
|
|
"oi": "open_interest",
|
|
|
})
|
|
|
|
|
|
df["trade_time"] = pd.to_datetime(df["trade_time"])
|
|
|
return df[["trade_time", "open", "high", "low", "close", "volume", "turnover", "open_interest"]]
|