292 lines
13 KiB
Python
292 lines
13 KiB
Python
import pandas as pd
|
||
from datetime import datetime, timedelta
|
||
import asyncio
|
||
from typing import Dict, List, Optional, Tuple, Any
|
||
from utils.logger import get_logger
|
||
|
||
# 获取日志器
|
||
logger = get_logger()
|
||
|
||
class StockDataProvider:
|
||
"""
|
||
异步股票数据提供服务
|
||
负责获取股票、基金等金融产品的历史数据
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""初始化数据提供者服务"""
|
||
logger.debug("初始化StockDataProvider")
|
||
|
||
async def get_stock_data(self, stock_code: str, market_type: str = 'A',
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None) -> pd.DataFrame:
|
||
"""
|
||
异步获取股票或基金数据
|
||
|
||
Args:
|
||
stock_code: 股票代码
|
||
market_type: 市场类型，默认为'A'股
|
||
start_date: 开始日期，格式YYYYMMDD，默认为一年前
|
||
end_date: 结束日期，格式YYYYMMDD，默认为今天
|
||
|
||
Returns:
|
||
包含历史数据的DataFrame
|
||
"""
|
||
# 使用线程池执行同步的akshare调用
|
||
return await asyncio.to_thread(
|
||
self._get_stock_data_sync,
|
||
stock_code,
|
||
market_type,
|
||
start_date,
|
||
end_date
|
||
)
|
||
|
||
def _get_stock_data_sync(self, stock_code: str, market_type: str = 'A',
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None) -> pd.DataFrame:
|
||
"""
|
||
同步获取股票数据的实现
|
||
将被异步方法调用
|
||
"""
|
||
import akshare as ak
|
||
|
||
if start_date is None:
|
||
start_date = (datetime.now() - timedelta(days=365)).strftime('%Y%m%d')
|
||
if end_date is None:
|
||
end_date = datetime.now().strftime('%Y%m%d')
|
||
|
||
# 确保日期格式统一（移除可能的'-'符号）
|
||
if isinstance(start_date, str) and '-' in start_date:
|
||
start_date = start_date.replace('-', '')
|
||
if isinstance(end_date, str) and '-' in end_date:
|
||
end_date = end_date.replace('-', '')
|
||
|
||
try:
|
||
if market_type == 'A':
|
||
logger.debug(f"获取A股数据: {stock_code}")
|
||
|
||
df = ak.stock_zh_a_hist(
|
||
symbol=stock_code,
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
adjust="qfq"
|
||
)
|
||
|
||
elif market_type in ['HK']:
|
||
logger.debug(f"获取港股数据: {stock_code}")
|
||
df = ak.stock_hk_daily(
|
||
symbol=stock_code,
|
||
adjust="qfq"
|
||
)
|
||
|
||
# 在获取数据后进行日期过滤
|
||
try:
|
||
if not isinstance(df.index, pd.DatetimeIndex):
|
||
# 如果存在命名为'date'的列，将其设为索引
|
||
if 'date' in df.columns:
|
||
df['date'] = pd.to_datetime(df['date'])
|
||
df.set_index('date', inplace=True)
|
||
else:
|
||
# 尝试将第一列转换为日期索引
|
||
date_col = df.columns[0]
|
||
df[date_col] = pd.to_datetime(df[date_col])
|
||
df.set_index(date_col, inplace=True)
|
||
|
||
# 转换日期字符串为日期对象
|
||
if start_date:
|
||
if start_date.isdigit() and len(start_date) == 8:
|
||
start_date_dt = pd.to_datetime(start_date, format='%Y%m%d')
|
||
else:
|
||
start_date_dt = pd.to_datetime(start_date)
|
||
else:
|
||
start_date_dt = pd.to_datetime((datetime.now() - timedelta(days=365)).strftime('%Y%m%d'))
|
||
|
||
if end_date:
|
||
if end_date.isdigit() and len(end_date) == 8:
|
||
end_date_dt = pd.to_datetime(end_date, format='%Y%m%d')
|
||
else:
|
||
end_date_dt = pd.to_datetime(end_date)
|
||
else:
|
||
end_date_dt = pd.to_datetime(datetime.now().strftime('%Y%m%d'))
|
||
|
||
# 过滤日期范围
|
||
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
|
||
logger.debug(f"港股日期过滤后数据点数: {len(df)}")
|
||
|
||
except Exception as e:
|
||
logger.warning(f"港股日期过滤出错: {str(e)}，使用原始数据")
|
||
|
||
elif market_type in ['US']:
|
||
logger.debug(f"获取美股数据: {stock_code}")
|
||
try:
|
||
df = ak.stock_us_daily(
|
||
symbol=stock_code,
|
||
adjust="qfq"
|
||
)
|
||
logger.debug(f"美股数据原始列: {df.columns.tolist()}")
|
||
logger.debug(f"美股数据形状: {df.shape}")
|
||
|
||
# 确保索引是日期时间类型
|
||
if not isinstance(df.index, pd.DatetimeIndex):
|
||
# 如果存在命名为'date'的列，将其设为索引
|
||
if 'date' in df.columns:
|
||
df['date'] = pd.to_datetime(df['date'])
|
||
df.set_index('date', inplace=True)
|
||
logger.debug("已将'date'列设置为索引")
|
||
else:
|
||
# 否则将当前索引转换为日期类型
|
||
df.index = pd.to_datetime(df.index)
|
||
logger.debug("已将索引转换为DatetimeIndex")
|
||
|
||
# 计算美股的成交额（Amount）= 成交量（Volume）× 收盘价（Close）
|
||
volume_col = next((col for col in df.columns if col.lower() == 'volume'), None)
|
||
close_col = next((col for col in df.columns if col.lower() == 'close'), None)
|
||
|
||
if volume_col and close_col:
|
||
df['amount'] = df[volume_col] * df[close_col]
|
||
logger.debug("已为美股数据计算成交额(amount)字段")
|
||
else:
|
||
logger.warning(f"美股数据缺少volume或close列，无法计算amount。当前列: {df.columns.tolist()}")
|
||
# 添加空的amount列，避免后续处理错误
|
||
df['amount'] = 0.0
|
||
|
||
# 将所有列名转为小写以进行统一处理
|
||
df.columns = [col.lower() for col in df.columns]
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取美股数据失败 {stock_code}: {str(e)}")
|
||
raise ValueError(f"获取美股数据失败 {stock_code}: {str(e)}")
|
||
|
||
# 将字符串日期转换为日期时间对象进行比较
|
||
try:
|
||
# 尝试多种格式解析日期
|
||
# 如果日期是数字格式（20220101），使用适当的格式
|
||
if start_date.isdigit() and len(start_date) == 8:
|
||
start_date_dt = pd.to_datetime(start_date, format='%Y%m%d')
|
||
else:
|
||
# 否则让pandas自动推断格式
|
||
start_date_dt = pd.to_datetime(start_date)
|
||
|
||
if end_date.isdigit() and len(end_date) == 8:
|
||
end_date_dt = pd.to_datetime(end_date, format='%Y%m%d')
|
||
else:
|
||
end_date_dt = pd.to_datetime(end_date)
|
||
except Exception as e:
|
||
logger.warning(f"日期转换出错: {str(e)}，使用默认值")
|
||
# 如果转换失败，使用合理的默认值
|
||
start_date_dt = pd.to_datetime('20000101', format='%Y%m%d')
|
||
end_date_dt = pd.to_datetime(datetime.now().strftime('%Y%m%d'), format='%Y%m%d')
|
||
|
||
# 过滤日期
|
||
try:
|
||
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
|
||
logger.debug(f"日期过滤后数据点数: {len(df)}")
|
||
except Exception as e:
|
||
logger.warning(f"日期过滤出错: {str(e)}，返回原始数据")
|
||
|
||
elif market_type in ['ETF', 'LOF']:
|
||
logger.debug(f"获取{market_type}基金数据: {stock_code}")
|
||
df = ak.fund_etf_hist_sina(
|
||
symbol=stock_code,
|
||
start_date=start_date.replace('-', ''),
|
||
end_date=end_date.replace('-', '')
|
||
)
|
||
|
||
else:
|
||
error_msg = f"不支持的市场类型: {market_type}"
|
||
logger.error(f"[市场类型错误] {error_msg}")
|
||
raise ValueError(error_msg)
|
||
|
||
# 标准化列名
|
||
if market_type == 'A':
|
||
# 根据实际数据结构调整列名映射
|
||
# 实际数据列：['日期', '股票代码', '开盘', '收盘', '最高', '最低', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
|
||
df.columns = ['Date', 'Code', 'Open', 'Close', 'High', 'Low', 'Volume', 'Amount', 'Amplitude', 'Change_pct', 'Change', 'Turnover']
|
||
elif market_type in ['HK', 'US']:
|
||
# 美股数据列可能不同，需要通过映射处理
|
||
columns_mapping = {
|
||
'open': 'Open',
|
||
'high': 'High',
|
||
'low': 'Low',
|
||
'close': 'Close',
|
||
'volume': 'Volume',
|
||
'amount': 'Amount'
|
||
}
|
||
|
||
# 创建新的DataFrame以确保列顺序和存在性
|
||
new_df = pd.DataFrame(index=df.index)
|
||
|
||
# 遍历映射，填充新DataFrame
|
||
for orig_col, new_col in columns_mapping.items():
|
||
if orig_col in df.columns:
|
||
new_df[new_col] = df[orig_col]
|
||
else:
|
||
# 如果原始列不存在，创建一个填充0的列
|
||
logger.warning(f"数据中缺少{orig_col}列，使用0值填充")
|
||
new_df[new_col] = 0.0
|
||
|
||
# 替换原始df
|
||
df = new_df
|
||
|
||
elif market_type in ['ETF', 'LOF']:
|
||
# 基金数据可能有不同的列
|
||
df.columns = ['Date', 'Open', 'High', 'Low', 'Close', 'Volume', 'Amount']
|
||
|
||
# 确保日期列是日期类型
|
||
if 'Date' in df.columns:
|
||
df['Date'] = pd.to_datetime(df['Date'])
|
||
df.set_index('Date', inplace=True)
|
||
|
||
# 确保按日期升序排序
|
||
df.sort_index(inplace=True)
|
||
|
||
logger.info(f"成功获取{market_type}数据 {stock_code}, 数据点数: {len(df)}")
|
||
return df
|
||
|
||
except Exception as e:
|
||
error_msg = f"获取{market_type}数据失败 {stock_code}: {str(e)}"
|
||
logger.error(error_msg)
|
||
logger.exception(e)
|
||
# 使用空的DataFrame并添加错误信息，而不是抛出异常
|
||
# 这样上层调用者可以检查是否有错误并适当处理
|
||
df = pd.DataFrame()
|
||
df.error = error_msg # 添加错误属性
|
||
return df
|
||
|
||
async def get_multiple_stocks_data(self, stock_codes: List[str],
|
||
market_type: str = 'A',
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
max_concurrency: int = 5) -> Dict[str, pd.DataFrame]:
|
||
"""
|
||
异步批量获取多只股票数据
|
||
|
||
Args:
|
||
stock_codes: 股票代码列表
|
||
market_type: 市场类型，默认为'A'股
|
||
start_date: 开始日期，格式YYYYMMDD
|
||
end_date: 结束日期，格式YYYYMMDD
|
||
max_concurrency: 最大并发数，默认为5
|
||
|
||
Returns:
|
||
字典，键为股票代码，值为对应的DataFrame
|
||
"""
|
||
# 使用信号量控制并发数
|
||
semaphore = asyncio.Semaphore(max_concurrency)
|
||
|
||
async def get_with_semaphore(code):
|
||
async with semaphore:
|
||
try:
|
||
return code, await self.get_stock_data(code, market_type, start_date, end_date)
|
||
except Exception as e:
|
||
logger.error(f"获取股票 {code} 数据时出错: {str(e)}")
|
||
return code, None
|
||
|
||
# 创建异步任务
|
||
tasks = [get_with_semaphore(code) for code in stock_codes]
|
||
|
||
# 等待所有任务完成
|
||
results = await asyncio.gather(*tasks)
|
||
|
||
# 构建结果字典，过滤掉失败的请求
|
||
return {code: df for code, df in results if df is not None} |