195 lines
6.6 KiB
Python
195 lines
6.6 KiB
Python
|
"""
|
|||
|
|
美股数据源
|
|||
|
|
使用 yfinance 和 finnhub 获取数据
|
|||
|
|
"""
|
|||
|
|
from typing import Dict, List, Any, Optional
|
|||
|
|
from datetime import datetime, timedelta
|
|||
|
|
|
|||
|
|
import yfinance as yf
|
|||
|
|
|
|||
|
|
from app.data_sources.base import BaseDataSource
|
|||
|
|
from app.utils.logger import get_logger
|
|||
|
|
from app.config import APIKeys, YFinanceConfig
|
|||
|
|
|
|||
|
|
logger = get_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class USStockDataSource(BaseDataSource):
|
|||
|
|
"""美股数据源"""
|
|||
|
|
|
|||
|
|
name = "USStock/yfinance"
|
|||
|
|
|
|||
|
|
# yfinance 时间周期映射
|
|||
|
|
INTERVAL_MAP = {
|
|||
|
|
'1m': '1m',
|
|||
|
|
'5m': '5m',
|
|||
|
|
'15m': '15m',
|
|||
|
|
'30m': '30m',
|
|||
|
|
'1H': '1h',
|
|||
|
|
'4H': '4h',
|
|||
|
|
'1D': '1d',
|
|||
|
|
'1W': '1wk'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 不同周期获取数据的天数范围
|
|||
|
|
DAYS_MAP = {
|
|||
|
|
'1m': lambda limit: min(7, max(1, (limit // 390) + 2)),
|
|||
|
|
'5m': lambda limit: min(60, max(1, (limit // 78) + 2)),
|
|||
|
|
'15m': lambda limit: min(60, max(1, (limit // 26) + 2)),
|
|||
|
|
'30m': lambda limit: min(60, max(1, (limit // 13) + 2)),
|
|||
|
|
'1H': lambda limit: min(730, max(1, (limit // 24) + 2)),
|
|||
|
|
'4H': lambda limit: min(730, max(1, (limit // 6) + 2)),
|
|||
|
|
'1D': lambda limit: min(3650, limit + 1),
|
|||
|
|
'1W': lambda limit: min(3650, (limit * 7) + 7)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
# 初始化 finnhub 作为备选
|
|||
|
|
self.finnhub_client = None
|
|||
|
|
try:
|
|||
|
|
import finnhub
|
|||
|
|
if APIKeys.is_configured('FINNHUB_API_KEY'):
|
|||
|
|
self.finnhub_client = finnhub.Client(api_key=APIKeys.FINNHUB_API_KEY)
|
|||
|
|
logger.info("Finnhub client initialized")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Finnhub init failed: {e}")
|
|||
|
|
|
|||
|
|
def get_kline(
|
|||
|
|
self,
|
|||
|
|
symbol: str,
|
|||
|
|
timeframe: str,
|
|||
|
|
limit: int,
|
|||
|
|
before_time: Optional[int] = None
|
|||
|
|
) -> List[Dict[str, Any]]:
|
|||
|
|
"""获取美股K线数据"""
|
|||
|
|
klines = []
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
interval = self.INTERVAL_MAP.get(timeframe, '1d')
|
|||
|
|
days_func = self.DAYS_MAP.get(timeframe, lambda x: x + 1)
|
|||
|
|
days = days_func(limit)
|
|||
|
|
|
|||
|
|
# 计算日期范围
|
|||
|
|
if before_time:
|
|||
|
|
end_date = datetime.fromtimestamp(before_time)
|
|||
|
|
start_date = end_date - timedelta(days=days)
|
|||
|
|
else:
|
|||
|
|
end_date = datetime.now()
|
|||
|
|
start_date = end_date - timedelta(days=days)
|
|||
|
|
|
|||
|
|
# logger.info(f"使用 yfinance 获取 {symbol}, 周期: {interval}, 日期: {start_date.date()} ~ {end_date.date()}")
|
|||
|
|
|
|||
|
|
# 尝试 yfinance
|
|||
|
|
df = self._fetch_yfinance(symbol, interval, start_date, end_date)
|
|||
|
|
|
|||
|
|
if df is None or df.empty:
|
|||
|
|
# 尝试 finnhub
|
|||
|
|
if self.finnhub_client and timeframe == '1D':
|
|||
|
|
klines = self._fetch_finnhub(symbol, start_date, end_date, limit)
|
|||
|
|
if klines:
|
|||
|
|
return klines
|
|||
|
|
else:
|
|||
|
|
klines = self._convert_dataframe(df, limit)
|
|||
|
|
|
|||
|
|
# 过滤和限制
|
|||
|
|
klines = self.filter_and_limit(klines, limit, before_time)
|
|||
|
|
|
|||
|
|
# 记录结果
|
|||
|
|
self.log_result(symbol, klines, timeframe)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Failed to fetch US stock K-lines {symbol}: {str(e)}")
|
|||
|
|
import traceback
|
|||
|
|
logger.error(traceback.format_exc())
|
|||
|
|
|
|||
|
|
return klines
|
|||
|
|
|
|||
|
|
def _fetch_yfinance(self, symbol: str, interval: str, start_date: datetime, end_date: datetime):
|
|||
|
|
"""使用 yfinance 获取数据"""
|
|||
|
|
try:
|
|||
|
|
ticker = yf.Ticker(symbol)
|
|||
|
|
df = ticker.history(
|
|||
|
|
start=start_date.strftime('%Y-%m-%d'),
|
|||
|
|
end=end_date.strftime('%Y-%m-%d'),
|
|||
|
|
interval=interval
|
|||
|
|
)
|
|||
|
|
# logger.info(f"yfinance 返回 {len(df) if df is not None and not df.empty else 0} 条数据")
|
|||
|
|
return df
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"yfinance fetch failed: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def _fetch_finnhub(
|
|||
|
|
self,
|
|||
|
|
symbol: str,
|
|||
|
|
start_date: datetime,
|
|||
|
|
end_date: datetime,
|
|||
|
|
limit: int
|
|||
|
|
) -> List[Dict[str, Any]]:
|
|||
|
|
"""使用 finnhub 获取日线数据"""
|
|||
|
|
klines = []
|
|||
|
|
try:
|
|||
|
|
start_ts = int(start_date.timestamp())
|
|||
|
|
end_ts = int(end_date.timestamp())
|
|||
|
|
|
|||
|
|
# logger.info(f"使用 Finnhub 获取 {symbol} 日线数据")
|
|||
|
|
candles = self.finnhub_client.stock_candles(symbol, 'D', start_ts, end_ts)
|
|||
|
|
|
|||
|
|
if candles and candles.get('s') == 'ok':
|
|||
|
|
for i in range(len(candles['t'])):
|
|||
|
|
klines.append(self.format_kline(
|
|||
|
|
timestamp=candles['t'][i],
|
|||
|
|
open_price=candles['o'][i],
|
|||
|
|
high=candles['h'][i],
|
|||
|
|
low=candles['l'][i],
|
|||
|
|
close=candles['c'][i],
|
|||
|
|
volume=candles['v'][i]
|
|||
|
|
))
|
|||
|
|
# logger.info(f"Finnhub 返回 {len(klines)} 条数据")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Finnhub fetch failed: {e}")
|
|||
|
|
|
|||
|
|
return klines
|
|||
|
|
|
|||
|
|
def _convert_dataframe(self, df, limit: int) -> List[Dict[str, Any]]:
|
|||
|
|
"""转换 DataFrame 为K线列表"""
|
|||
|
|
klines = []
|
|||
|
|
df = df.tail(limit).reset_index()
|
|||
|
|
|
|||
|
|
# 确定时间列名（日线是 Date，分钟级是 Datetime）
|
|||
|
|
time_col = None
|
|||
|
|
if 'Datetime' in df.columns:
|
|||
|
|
time_col = 'Datetime'
|
|||
|
|
elif 'Date' in df.columns:
|
|||
|
|
time_col = 'Date'
|
|||
|
|
elif 'index' in df.columns:
|
|||
|
|
time_col = 'index'
|
|||
|
|
|
|||
|
|
if time_col is None:
|
|||
|
|
logger.warning(f"Unable to determine time column; available columns: {df.columns.tolist()}")
|
|||
|
|
return klines
|
|||
|
|
|
|||
|
|
for _, row in df.iterrows():
|
|||
|
|
try:
|
|||
|
|
# 处理时间戳
|
|||
|
|
time_value = row[time_col]
|
|||
|
|
if hasattr(time_value, 'timestamp'):
|
|||
|
|
ts = int(time_value.timestamp())
|
|||
|
|
else:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
klines.append(self.format_kline(
|
|||
|
|
timestamp=ts,
|
|||
|
|
open_price=row['Open'],
|
|||
|
|
high=row['High'],
|
|||
|
|
low=row['Low'],
|
|||
|
|
close=row['Close'],
|
|||
|
|
volume=row['Volume']
|
|||
|
|
))
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.debug(f"Failed to parse row data: {e}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
return klines
|
|||
|
|
|