# ================================ # 版本切换开关 - 从配置文件导入 # ================================ from data_source_config import USE_YFINANCE_VERSION, API_TIMEOUT_SECONDS, MAX_RETRY_ATTEMPTS import logging import re import pandas as pd from datetime import datetime, timedelta import time # 导入标准库的 time 模块 import os import requests import threading import asyncio # 根据开关导入不同的模块 if USE_YFINANCE_VERSION: import yfinance as yf print("🔄 Using yfinance version (new)") else: import akshare as ak print("🔄 Using akshare version (old)") logging.basicConfig(level=logging.INFO) # 获取当前文件的目录 base_dir = os.path.dirname(os.path.abspath(__file__)) # 构建CSV文件的绝对路径 nasdaq_100_path = os.path.join(base_dir, './model/nasdaq100.csv') dow_jones_path = os.path.join(base_dir, './model/dji.csv') sp500_path = os.path.join(base_dir, './model/sp500.csv') nasdaq_composite_path = os.path.join(base_dir, './model/nasdaq_all.csv') # 从CSV文件加载成分股数据 nasdaq_100_stocks = pd.read_csv(nasdaq_100_path) dow_jones_stocks = pd.read_csv(dow_jones_path) sp500_stocks = pd.read_csv(sp500_path) nasdaq_composite_stocks = pd.read_csv(nasdaq_composite_path) def fetch_stock_us_spot_data_with_retries(): """根据开关选择不同的数据源获取股票列表""" if USE_YFINANCE_VERSION: return fetch_stock_us_spot_data_yfinance() else: return fetch_stock_us_spot_data_akshare() def fetch_stock_us_spot_data_akshare(): """原始的 akshare 实现""" if not USE_YFINANCE_VERSION: # 定义重试间隔时间序列（秒） retry_intervals = [10, 20, 60, 300, 600] retry_index = 0 # 初始重试序号 while True: try: # 尝试获取API数据 symbols = ak.stock_us_spot_em() return symbols # 成功获取数据后返回 except Exception as e: print(f"Error fetching data: {e}") # 获取当前重试等待时间 wait_time = retry_intervals[retry_index] print(f"Retrying in {wait_time} seconds...") time.sleep(wait_time) # 等待指定的秒数 # 更新重试索引，但不要超出重试时间列表的范围 retry_index = min(retry_index + 1, len(retry_intervals) - 1) else: print("Warning: akshare function called while using yfinance version") return pd.DataFrame() def fetch_stock_us_spot_data_yfinance(): """新的 yfinance 实现""" try: # 从本地CSV文件收集所有股票代码 all_symbols = set() # 从各个指数CSV文件中提取股票代码 for df, name in [ (nasdaq_100_stocks, "NASDAQ-100"), (dow_jones_stocks, "Dow Jones"), (sp500_stocks, "S&P 500"), (nasdaq_composite_stocks, "NASDAQ Composite") ]: if 'Symbol' in df.columns: symbols_from_csv = df['Symbol'].dropna().astype(str).tolist() all_symbols.update(symbols_from_csv) elif 'Code' in df.columns: symbols_from_csv = df['Code'].dropna().astype(str).tolist() all_symbols.update(symbols_from_csv) # 添加一些常见的ETF和热门股票 additional_symbols = [ # 主要ETF 'SPY', 'QQQ', 'IWM', 'VTI', 'ARKK', 'TQQQ', 'SQQQ', 'SPXL', # 热门科技股 'AAPL', 'MSFT', 'GOOGL', 'GOOG', 'AMZN', 'TSLA', 'META', 'NVDA', 'NFLX', 'AMD', 'INTC', 'ORCL', 'CRM', 'ADBE', 'PYPL', 'UBER', 'LYFT', # 中概股 'BABA', 'JD', 'PDD', 'NIO', 'XPEV', 'LI', 'DIDI', 'TME', # 其他热门股票 'COST', 'WMT', 'JPM', 'BAC', 'XOM', 'CVX', 'PFE', 'JNJ', 'KO', 'PEP' ] all_symbols.update(additional_symbols) # 创建DataFrame symbols_list = sorted(list(all_symbols)) symbols_df = pd.DataFrame({ '代码': symbols_list, '名称': [f'{symbol} Inc.' for symbol in symbols_list] # 简单的名称映射 }) print(f"Created symbols dataframe with {len(symbols_df)} symbols using yfinance version") return symbols_df except Exception as e: print(f"Error creating symbols dataframe: {e}") # 返回基本的fallback数据 fallback_symbols = [ 'AAPL', 'MSFT', 'GOOGL', 'AMZN', 'TSLA', 'META', 'NVDA', 'NFLX', 'SPY', 'QQQ', 'IWM', 'VTI' ] return pd.DataFrame({ '代码': fallback_symbols, '名称': [f'{symbol} Inc.' for symbol in fallback_symbols] }) async def fetch_stock_us_spot_data_with_retries_async(): """异步版本的股票数据获取，支持版本切换""" if USE_YFINANCE_VERSION: try: return await asyncio.to_thread(fetch_stock_us_spot_data_yfinance) except Exception as e: print(f"Error in async yfinance fetch: {e}") return pd.DataFrame() else: return await fetch_stock_us_spot_data_akshare_async() async def fetch_stock_us_spot_data_akshare_async(): """原始的 akshare 异步实现""" if not USE_YFINANCE_VERSION: retry_intervals = [10, 20] # 减少重试次数 retry_index = 0 max_retries = 2 # 最多重试2次 for attempt in range(max_retries + 1): try: # 添加30秒超时 symbols = await asyncio.wait_for( asyncio.to_thread(ak.stock_us_spot_em), timeout=30.0 ) return symbols except asyncio.TimeoutError: print(f"Timeout error fetching data (attempt {attempt + 1}/{max_retries + 1})") except Exception as e: print(f"Error fetching data (attempt {attempt + 1}/{max_retries + 1}): {e}") if attempt < max_retries: wait_time = retry_intervals[min(retry_index, len(retry_intervals) - 1)] print(f"Retrying in {wait_time} seconds...") await asyncio.sleep(wait_time) retry_index += 1 # 如果所有重试都失败，返回空数据 print("All retries failed, returning empty data") return pd.DataFrame() else: print("Warning: akshare async function called while using yfinance version") return pd.DataFrame() symbols = None def create_fallback_symbols(): """创建fallback符号数据，用于测试""" fallback_symbols = [ 'AAPL', 'MSFT', 'GOOGL', 'AMZN', 'TSLA', 'META', 'NVDA', 'NFLX', 'SPY', 'QQQ', 'IWM', 'VTI' ] return pd.DataFrame({ '代码': fallback_symbols, '名称': [f'{symbol} Inc.' for symbol in fallback_symbols] }) async def fetch_symbols(): global symbols try: print("Starting symbols initialization...") # 异步获取数据 symbols = await fetch_stock_us_spot_data_with_retries_async() if symbols is not None and not symbols.empty: print(f"Symbols initialized successfully: {len(symbols)} symbols loaded") else: print("Symbols initialization failed, using empty dataset") symbols = pd.DataFrame() except Exception as e: print(f"Error in fetch_symbols: {e}") symbols = pd.DataFrame() finally: print("Symbols initialization completed") # 全局变量 index_us_stock_index_INX = None index_us_stock_index_DJI = None index_us_stock_index_IXIC = None index_us_stock_index_NDX = None def update_stock_indices(): global index_us_stock_index_INX, index_us_stock_index_DJI, index_us_stock_index_IXIC, index_us_stock_index_NDX try: print("Starting stock indices update...") if USE_YFINANCE_VERSION: print("Updating indices using yfinance...") # 使用 yfinance 更新指数数据 from datetime import datetime, timedelta # 计算日期范围 end_date = datetime.now() start_date = end_date - timedelta(weeks=8) # 定义指数映射 indices = { '^GSPC': 'INX', # S&P 500 '^DJI': 'DJI', # Dow Jones '^IXIC': 'IXIC', # NASDAQ Composite '^NDX': 'NDX' # NASDAQ 100 } for yf_symbol, var_name in indices.items(): try: ticker = yf.Ticker(yf_symbol) hist_data = ticker.history(start=start_date, end=end_date) if not hist_data.empty: # 转换为与akshare相同的格式 formatted_data = pd.DataFrame({ 'date': hist_data.index.strftime('%Y-%m-%d'), '开盘': hist_data['Open'].values, '收盘': hist_data['Close'].values, '最高': hist_data['High'].values, '最低': hist_data['Low'].values, '成交量': hist_data['Volume'].values, '成交额': (hist_data['Close'] * hist_data['Volume']).values }) # 设置全局变量 if var_name == 'INX': index_us_stock_index_INX = formatted_data elif var_name == 'DJI': index_us_stock_index_DJI = formatted_data elif var_name == 'IXIC': index_us_stock_index_IXIC = formatted_data elif var_name == 'NDX': index_us_stock_index_NDX = formatted_data print(f"Successfully updated {var_name}: {len(formatted_data)} records") else: print(f"No data received for {yf_symbol}") except Exception as e: print(f"Error fetching {yf_symbol}: {e}") else: print("Updating indices using akshare...") # 使用 akshare 更新指数数据 index_us_stock_index_INX = ak.index_us_stock_sina(symbol=".INX") index_us_stock_index_DJI = ak.index_us_stock_sina(symbol=".DJI") index_us_stock_index_IXIC = ak.index_us_stock_sina(symbol=".IXIC") index_us_stock_index_NDX = ak.index_us_stock_sina(symbol=".NDX") print("Stock indices updated successfully") except Exception as e: print(f"Error updating stock indices: {e}") # 设置定时器，每隔12小时更新一次 threading.Timer(12 * 60 * 60, update_stock_indices).start() # 程序开始时不立即更新，而是延迟启动 def start_indices_update(): """延迟启动股票指数更新，避免阻塞应用启动""" threading.Timer(5, update_stock_indices).start() # 5秒后开始第一次更新 # 延迟启动股票指数更新 start_indices_update() # 创建列名转换的字典 column_mapping = { '日期': 'date', '开盘': 'open', '收盘': 'close', '最高': 'high', '最低': 'low', '成交量': 'volume', '成交额': 'amount', '振幅': 'amplitude', '涨跌幅': 'price_change_percentage', '涨跌额': 'price_change_amount', '换手率': 'turnover_rate' } # 定义一个标准的列顺序 standard_columns = ['date', 'open', 'close', 'high', 'low', 'volume', 'amount'] # 定义查找函数 def find_stock_entry(stock_code): # 使用 str.endswith 来匹配股票代码 matching_row = symbols[symbols['代码'].str.endswith(stock_code)] # print(symbols) if not matching_row.empty: # print(f"股票代码 {stock_code} 找到, 代码为 {matching_row['代码'].values[0]}") return matching_row['代码'].values[0] else: return "" ''' # 示例调用 # 测试函数 result = find_stock_entry('AAPL') if isinstance(result, pd.DataFrame) and not result.empty: # 如果找到的结果不为空，获取代码列的值 code_value = result['代码'].values[0] print(code_value) else: print(result) ''' def reduce_columns(df, columns_to_keep): return df[columns_to_keep] # 创建缓存字典 _price_cache = {} def get_last_minute_stock_price(symbol: str, max_retries=3) -> float: """获取股票最新价格，使用30分钟缓存，并包含重试机制""" if not symbol: return -1.0 if symbol == "NONE_SYMBOL_FOUND": return -1.0 current_time = datetime.now() # 检查缓存 if symbol in _price_cache: cached_price, cached_time = _price_cache[symbol] # 如果缓存时间在30分钟内，直接返回缓存的价格 if current_time - cached_time < timedelta(minutes=30): return cached_price # 重试机制 for attempt in range(max_retries): try: # 缓存无效或不存在，从yfinance获取新数据 if USE_YFINANCE_VERSION: stock_data = yf.download( symbol, period='1d', interval='5m', progress=False, # 禁用进度条 timeout=10 # 设置超时时间 ) else: # 使用akshare获取数据的逻辑 ticker = ak.stock_us_hist(symbol=symbol, period="daily", start_date="20240101", end_date="20240201") stock_data = ticker if not ticker.empty else pd.DataFrame() if stock_data.empty: print(f"Warning: Empty data received for {symbol}, attempt {attempt + 1}/{max_retries}") if attempt == max_retries - 1: return -1.0 time.sleep(1) # 等待1秒后重试 continue latest_price = float(stock_data['Close'].iloc[-1]) # 更新缓存 _price_cache[symbol] = (latest_price, current_time) return latest_price except Exception as e: print(f"Error fetching price for {symbol}, attempt {attempt + 1}/{max_retries}: {str(e)}") if attempt == max_retries - 1: return -1.0 time.sleep(1) # 等待1秒后重试 return -1.0 # 返回个股历史数据 def get_stock_history(symbol, news_date, retries=10): # 定义重试间隔时间序列（秒） retry_intervals = [10, 20, 60, 300, 600] retry_count = 0 # 如果传入的symbol不包含数字前缀，则通过 find_stock_entry 获取完整的symbol if not any(char.isdigit() for char in symbol): full_symbol = find_stock_entry(symbol) if len(symbol) != 0 and full_symbol: symbol = full_symbol else: symbol = "" # 将news_date转换为datetime对象 current_date = datetime.now() # 计算start_date和end_date start_date = (current_date - timedelta(days=60)).strftime("%Y%m%d") end_date = current_date.strftime("%Y%m%d") stock_hist_df = None retry_index = 0 # 初始化重试索引 while retry_count <= retries and len(symbol) != 0: # 无限循环重试 try: # 根据版本开关选择不同的API if USE_YFINANCE_VERSION: # 使用 yfinance 获取数据 ticker = yf.Ticker(symbol) # 将日期格式转换为 yfinance 期望的格式 (YYYY-MM-DD) yf_start_date = datetime.strptime(start_date, "%Y%m%d").strftime("%Y-%m-%d") yf_end_date = datetime.strptime(end_date, "%Y%m%d").strftime("%Y-%m-%d") stock_hist_df = ticker.history(start=yf_start_date, end=yf_end_date) if not stock_hist_df.empty: # 转换为与akshare相同的格式 stock_hist_df = stock_hist_df.reset_index() stock_hist_df = pd.DataFrame({ 'date': stock_hist_df['Date'].dt.strftime('%Y-%m-%d'), '开盘': stock_hist_df['Open'], '收盘': stock_hist_df['Close'], '最高': stock_hist_df['High'], '最低': stock_hist_df['Low'], '成交量': stock_hist_df['Volume'], '成交额': stock_hist_df['Close'] * stock_hist_df['Volume'], '振幅': 0, # yfinance没有直接提供，设为0 '涨跌幅': 0, # 可以计算，但这里简化为0 '涨跌额': 0, # 可以计算，但这里简化为0 '换手率': 0 # yfinance没有直接提供，设为0 }) else: stock_hist_df = None else: # 使用 akshare 获取数据 stock_hist_df = ak.stock_us_hist(symbol=symbol, period="daily", start_date=start_date, end_date=end_date, adjust="") if stock_hist_df is None or stock_hist_df.empty: # 检查是否为空数据 # print(f"No data for {symbol} on {news_date}.") stock_hist_df = None # 将 DataFrame 设置为 None break except (requests.exceptions.Timeout, ConnectionError) as e: print(f"Request timed out: {e}. Retrying...") retry_count += 1 # 增加重试次数 continue except (TypeError, ValueError, BaseException) as e: print(f"Error {e} scraping data for {symbol} on {news_date}. Break...") # 可能是没数据，直接Break break # 如果发生异常，等待一段时间再重试 wait_time = retry_intervals[retry_index] print(f"Waiting for {wait_time} seconds before retrying...") time.sleep(wait_time) retry_index = (retry_index + 1) if retry_index < len(retry_intervals) - 1 else retry_index # 更新重试索引，不超过列表长度 # 如果获取失败或数据为空，返回填充为0的 DataFrame if stock_hist_df is None or stock_hist_df.empty: # 构建一个空的 DataFrame，包含指定日期范围的空数据 date_range = pd.date_range(start=start_date, end=end_date) stock_hist_df = pd.DataFrame({ 'date': date_range, '开盘': 0, '收盘': 0, '最高': 0, '最低': 0, '成交量': 0, '成交额': 0, '振幅': 0, '涨跌幅': 0, '涨跌额': 0, '换手率': 0 }) # 使用rename方法转换列名 stock_hist_df = stock_hist_df.rename(columns=column_mapping) stock_hist_df = stock_hist_df.reindex(columns=standard_columns) # 处理个股数据，保留所需列 stock_hist_df = reduce_columns(stock_hist_df, standard_columns) return stock_hist_df # 统一列名 stock_hist_df = stock_hist_df.rename(columns=column_mapping) stock_hist_df = stock_hist_df.reindex(columns=standard_columns) # 处理个股数据，保留所需列 stock_hist_df = reduce_columns(stock_hist_df, standard_columns) return stock_hist_df ''' # 示例调用 result = get_stock_history('AAPL', '20240214') print(result) ''' # result = get_stock_history('ATMU', '20231218') # print(result) # 返回个股所属指数历史数据 def get_stock_index_history(symbol, news_date, force_index=0): # 检查股票所属的指数 if symbol in nasdaq_100_stocks['Symbol'].values or force_index == 1: index_code = ".NDX" index_data = index_us_stock_index_NDX elif symbol in dow_jones_stocks['Symbol'].values or force_index == 2: index_code = ".DJI" index_data = index_us_stock_index_DJI elif symbol in sp500_stocks['Symbol'].values or force_index == 3: index_code = ".INX" index_data = index_us_stock_index_INX elif symbol in nasdaq_composite_stocks["Symbol"].values or symbol is None or symbol == "" or force_index == 4: index_code = ".IXIC" index_data = index_us_stock_index_IXIC else: # print(f"股票代码 {symbol} 不属于纳斯达克100、道琼斯工业、标准普尔500或纳斯达克综合指数。") index_code = ".IXIC" index_data = index_us_stock_index_IXIC # 获取当前日期 current_date = datetime.now() # 计算 start_date 和 end_date start_date = (current_date - timedelta(weeks=8)).strftime("%Y-%m-%d") end_date = current_date.strftime("%Y-%m-%d") # 确保 index_data['date'] 是 datetime 类型 index_data['date'] = pd.to_datetime(index_data['date']) # 从指数历史数据中提取指定日期范围的数据 index_hist_df = index_data[(index_data['date'] >= start_date) & (index_data['date'] <= end_date)] # 统一列名 index_hist_df = index_hist_df.rename(columns=column_mapping) index_hist_df = index_hist_df.reindex(columns=standard_columns) # 处理个股数据，保留所需列 index_hist_df = reduce_columns(index_hist_df, standard_columns) return index_hist_df ''' # 示例调用 result = get_stock_index_history('AAPL', '20240214') print(result) ''' def find_stock_codes_or_names(entities): """ 从给定的实体列表中检索股票代码或公司名称。 :param entities: 命名实体识别结果列表，格式为 [('实体名称', '实体类型'), ...] :return: 相关的股票代码列表 """ stock_codes = set() # 合并所有股票字典并清理数据，确保都是字符串 all_symbols = pd.concat([nasdaq_100_stocks['Symbol'], dow_jones_stocks['Symbol'], sp500_stocks['Symbol'], nasdaq_composite_stocks['Symbol']]).dropna().astype(str).unique().tolist() all_names = pd.concat([nasdaq_100_stocks['Name'], nasdaq_composite_stocks['Name'], sp500_stocks['Security'], dow_jones_stocks['Company']]).dropna().astype(str).unique().tolist() # 创建一个 Name 到 Symbol 的映射 name_to_symbol = {} for idx, name in enumerate(all_names): if idx < len(all_symbols): symbol = all_symbols[idx] name_to_symbol[name.lower()] = symbol # 查找实体映射到的股票代码 for entity, entity_type in entities: entity_lower = entity.lower() entity_upper = entity.upper() # 检查 Symbol 列 if entity_upper in all_symbols: stock_codes.add(entity_upper) #print(f"Matched symbol: {entity_upper}") # 检查 Name 列，确保完整匹配而不是部分匹配 for name, symbol in name_to_symbol.items(): # 使用正则表达式进行严格匹配 pattern = rf'\b{re.escape(entity_lower)}\b' if re.search(pattern, name): stock_codes.add(symbol.upper()) #print(f"Matched name/company: '{entity_lower}' in '{name}' -> {symbol.upper()}") #print(f"Stock codes found: {stock_codes}") if not stock_codes: return ['NONE_SYMBOL_FOUND'] return list(stock_codes) def process_history(stock_history, target_date, history_days=30, following_days=3): # 检查数据是否为空 if stock_history.empty: return create_empty_data(history_days), create_empty_data(following_days) # 确保日期列存在并转换为datetime格式 if 'date' not in stock_history.columns: return create_empty_data(history_days), create_empty_data(following_days) stock_history['date'] = pd.to_datetime(stock_history['date']) target_date = pd.to_datetime(target_date) # 按日期升序排序 stock_history = stock_history.sort_values('date') # 找到目标日期对应的索引 target_row = stock_history[stock_history['date'] <= target_date] if target_row.empty: return create_empty_data(history_days), create_empty_data(following_days) # 获取目标日期最近的行 target_index = target_row.index[-1] target_pos = stock_history.index.get_loc(target_index) # 获取历史数据（包括目标日期） start_pos = max(0, target_pos - history_days + 1) previous_rows = stock_history.iloc[start_pos:target_pos + 1] # 获取后续数据 following_rows = stock_history.iloc[target_pos + 1:target_pos + following_days + 1] # 删除日期列并确保数据完整性 previous_rows = previous_rows.drop(columns=['date']) following_rows = following_rows.drop(columns=['date']) # 处理数据不足的情况 previous_rows = handle_insufficient_data(previous_rows, history_days) following_rows = handle_insufficient_data(following_rows, following_days) return previous_rows.iloc[:, :6], following_rows.iloc[:, :6] def create_empty_data(days): return pd.DataFrame({ '开盘': [-1] * days, '收盘': [-1] * days, '最高': [-1] * days, '最低': [-1] * days, '成交量': [-1] * days, '成交额': [-1] * days }) def handle_insufficient_data(data, required_days): current_rows = len(data) if current_rows < required_days: missing_rows = required_days - current_rows empty_data = create_empty_data(missing_rows) return pd.concat([empty_data, data]).reset_index(drop=True) return data if __name__ == "__main__": # 测试函数 result = find_stock_entry('AAPL') print(f"find_stock_entry: {result}") result = get_stock_history('AAPL', '20240214') print(f"get_stock_history: {result}") result = get_stock_index_history('AAPL', '20240214') print(f"get_stock_index_history: {result}") result = find_stock_codes_or_names([('苹果', 'ORG'), ('苹果公司', 'ORG')]) print(f"find_stock_codes_or_names: {result}") result = process_history(get_stock_history('AAPL', '20240214'), '20240214') print(f"process_history: {result}") result = process_history(get_stock_index_history('AAPL', '20240214'), '20240214') print(f"process_history: {result}") pass