diff --git a/data_manager.py b/data_manager.py index 886d20a..0068df5 100644 --- a/data_manager.py +++ b/data_manager.py @@ -258,4 +258,79 @@ class DataReader: except Exception as e: logger.error(f"获取已存在交易日期时出错: {e}") logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}") - return [] \ No newline at end of file + return [] + + @staticmethod + def get_main_board_stocks(): + """ + 筛选主板且名称不包含ST的股票代码列表 + + 返回: + list: 符合条件的股票代码(ts_code)列表 + """ + # 先检查表是否存在 + if not db_manager.table_exists('stock_basic'): + logger.debug("表 stock_basic 不存在") + # 可以在这里添加自动拉取股票基本信息的代码,类似于get_trade_cal方法 + return [] + + try: + # 查询条件:主板(market='主板')、名称不含ST + query = "SELECT ts_code FROM stock_basic WHERE market='主板' AND name NOT LIKE '%ST%'" + result = db_manager.query(query) + return result['ts_code'].tolist() if not result.empty else [] + except Exception as e: + logger.error(f"获取主板上市非ST股票列表时出错: {e}") + logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}") + return [] + + + @staticmethod + def get_stk_auction(start_date=None, end_date=None, filter_main_board=True): + """ + 从数据库获取盘前竞价数据(stk_auction),并截取指定日期范围内的数据 + + 参数: + start_date (str): 开始日期,格式'YYYYMMDD',默认为30天前 + end_date (str): 结束日期,格式'YYYYMMDD',默认为今天 + filter_main_board (bool): 是否只返回主板上市股票的数据,默认为True + + 返回: + pandas.DataFrame: 盘前竞价数据 + """ + # 先检查表是否存在 + if not db_manager.table_exists('stk_auction'): + logger.debug(f"表 stk_auction 不存在") + return pd.DataFrame() + + # 设置默认日期范围 + if start_date is None: + start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d') + if end_date is None: + end_date = datetime.now().strftime('%Y%m%d') + + try: + # 基础条件:日期过滤 + conditions = f"trade_date BETWEEN '{start_date}' AND '{end_date}'" + + # 如果需要过滤主板股票 + if filter_main_board: + # 获取主板上市非ST股票列表 + main_board_stocks = DataReader.get_main_board_stocks() + if main_board_stocks: + # 将股票代码列表转换为SQL安全的格式 + stock_codes = ", ".join([f"'{code}'" for code in main_board_stocks]) + # 添加股票代码过滤条件 + conditions += f" AND ts_code IN ({stock_codes})" + else: + logger.warning("未找到主板上市股票,返回空DataFrame") + return pd.DataFrame() + + # 查询数据 + df = db_manager.load_df_from_db('stk_auction', conditions=conditions) + return df + except Exception as e: + logger.error(f"获取盘前竞价数据时出错: {e}") + logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}") + return pd.DataFrame() +