import pymysql import logging import sys import time import config import argparse # 配置不同市场的表名管理 tables_mapping = { 'sp500': { 'none_his_kline': 'sp500_his_kline_none', 'adjust_his_kline': 'sp500_ajust_kline_202410', 'rehab_table': 'futu_rehab' }, 'hs300': { 'none_his_kline': 'hs300_his_kline_none', 'adjust_his_kline': 'hs300_ajust_kline_202410', 'rehab_table': 'futu_rehab' } } # 日志配置 config.setup_logging("./log/stat_adjust_kline.log") logger = logging.getLogger() # MySQL数据库连接 def connect_to_db(): try: return pymysql.connect( **config.db_config, cursorclass=pymysql.cursors.DictCursor ) except pymysql.MySQLError as e: logger.error(f"Error connecting to the database: {e}", exc_info=True) return None # 从指定市场表中读取code和code_name字段 def fetch_codes_from_market_table(market, debug=False): db = connect_to_db() if db is None: logger.error("Failed to connect to database.") return [] try: with db.cursor() as cursor: query = f"SELECT code, code_name FROM {market}" if debug: query += " LIMIT 2" cursor.execute(query) codes = cursor.fetchall() return codes finally: db.close() # 读取复权因子数据 def fetch_rehab_data(db, code): try: with db.cursor() as cursor: cursor.execute(f"SELECT ex_div_date, forward_adj_factorA, forward_adj_factorB, backward_adj_factorA, backward_adj_factorB FROM futu_rehab WHERE code = %s ORDER BY ex_div_date DESC", (code,)) return cursor.fetchall() except pymysql.MySQLError as e: logger.error(f"Error fetching rehab data for {code}: {e}", exc_info=True) return [] # 读取不复权的股票价格数据 def fetch_kline_none_data(db, table_name, code): try: with db.cursor() as cursor: cursor.execute(f"SELECT code, time_key, open, close FROM {table_name} WHERE code = %s ORDER BY time_key ASC", (code,)) return cursor.fetchall() except pymysql.MySQLError as e: logger.error(f"Error fetching kline none data for {code}: {e}", exc_info=True) return [] # 插入前后复权价格到数据库 def insert_hfq_data(db, hfq_data, hfq_table): try: with db.cursor() as cursor: insert_query = f""" INSERT INTO {hfq_table} (code, name, time_key, hfq_open, hfq_close, qfq_open, qfq_close, none_open, none_close) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE hfq_open = VALUES(hfq_open), hfq_close = VALUES(hfq_close), qfq_open = VALUES(qfq_open), qfq_close = VALUES(qfq_close), none_open = VALUES(none_open), none_close = VALUES(none_close) """ cursor.executemany(insert_query, hfq_data) db.commit() except pymysql.MySQLError as e: logger.error(f"Error inserting hfq data: {e}", exc_info=True) # 计算后复权价格和前复权价格 def calculate_hfq_qfq_price(market, debug=False): db = connect_to_db() if db is None: return # 获取表名 table_names = tables_mapping[market] none_his_kline_table = table_names['none_his_kline'] adjust_kline_table = table_names['adjust_his_kline'] # 获取股票代码 codes = fetch_codes_from_market_table(market, debug) for row in codes: code = row['code'] name = row['code_name'] # 如果是 sp500 市场，拼接 'US.' + code rehab_code = code if market == 'sp500': rehab_code = 'US.' + code logger.info(f"Processing {code} ({name})...") # 获取复权因子数据 rehab_res = fetch_rehab_data(db, rehab_code) if not rehab_res: logger.warning(f"No rehab data found for {code}") continue # 反转复权因子行，为了计算前复权 rehab_res_asc = list(reversed(rehab_res)) # 获取不复权的价格数据 kline_none = fetch_kline_none_data(db, none_his_kline_table, code) if not kline_none: logger.warning(f"No kline none data found for {code}") continue hfq_data = [] # 遍历kline_none，计算前后复权价格 for kline_row in kline_none: none_open = kline_row['open'] none_close = kline_row['close'] time_key = kline_row['time_key'] # 将 time_key 转换为 date 格式 time_key_date = time_key.date() # 计算后复权价格 hfq_open = none_open hfq_close = none_close tmp_close = none_close tmp_open = none_open for rehab_row in rehab_res: if rehab_row['ex_div_date'] <= time_key_date: hfq_close = (tmp_close * rehab_row['backward_adj_factorA']) + rehab_row['backward_adj_factorB'] hfq_open = (tmp_open * rehab_row['backward_adj_factorA']) + rehab_row['backward_adj_factorB'] tmp_close = hfq_close tmp_open = hfq_open # 计算前复权价格 qfq_close = none_close qfq_open = none_open tmp_close = none_close tmp_open = none_open zeroFactorB = True if market == 'sp500' else False for rehab_row in rehab_res_asc: factorB = 0 if zeroFactorB else rehab_row['forward_adj_factorB'] if rehab_row['ex_div_date'] > time_key_date: qfq_close = (tmp_close * rehab_row['forward_adj_factorA']) + factorB qfq_open = (tmp_open * rehab_row['forward_adj_factorA']) + factorB tmp_close = qfq_close tmp_open = qfq_open # 保存计算后的复权价格 hfq_data.append((code, name, time_key, hfq_open, hfq_close, qfq_open, qfq_close, none_open, none_close)) # 插入后复权价格数据 insert_hfq_data(db, hfq_data, adjust_kline_table) logger.info(f"Inserted HFQ/QFQ data for {code} ({name})") time.sleep(1) db.close() if __name__ == "__main__": # 命令行参数处理 parser = argparse.ArgumentParser(description='Calculate HFQ and QFQ Prices for Market') parser.add_argument('--market', type=str, default='hs300', help='Market to process (sp500 or hs300)') parser.add_argument('--debug', action='store_true', help='Enable debug mode (limit records)') args = parser.parse_args() # 调用主函数 calculate_hfq_qfq_price(args.market, args.debug)