myquant/djapi/api/stock/divSearch.py

import pandas as pd

from config import START_DATE, END_DATE, PRECISION_CONFIG
from stock_utils import *
from getStockDiv2 import analyze_stock_dividend_and_price
from getStockParam import getStockParam
from .data_source import get_tushare_pro
import datetime

pro = get_tushare_pro()

# --- 数据输出配置 ---
OUTPUT_TO_CSV = True  # 开关：是否输出到CSV
OUTPUT_TO_MYSQL = False # 开关：是否输出到MySQL (需要额外配置数据库连接)
    

# --- 数据库输出函数 (伪代码) ---
def save_to_mysql(dataframe, table_name):
    """
    将DataFrame保存到MySQL数据库的伪代码函数。
    实际使用时需要配置数据库连接。
    """
    if not OUTPUT_TO_MYSQL:
        print("MySQL输出已禁用。")
        return

    print(f"正在将数据保存到MySQL表 {table_name}...")
    # --- 伪代码开始 ---
    # import pymysql # 或其他数据库连接库
    # 
    # connection = pymysql.connect(host='your_host',
    #                              user='your_user',
    #                              password='your_password',
    #                              database='your_database',
    #                              charset='utf8mb4')
    # try:
    #     with connection.cursor() as cursor:
    #         # 创建表 (如果不存在)
    #         # create_table_sql = "..."
    #         # cursor.execute(create_table_sql)
    #         
    #         # 遍历DataFrame并插入数据
    #         for index, row in dataframe.iterrows():
    #             # 注意：需要处理SQL注入风险，最好使用参数化查询
    #             insert_sql = f"INSERT INTO {table_name} (...) VALUES (...)"
    #             cursor.execute(insert_sql, tuple(row))
    #     connection.commit()
    #     print(f"成功保存 {len(dataframe)} 条记录到MySQL。")
    # except Exception as e:
    #     print(f"保存到MySQL时出错: {e}")
    #     connection.rollback()
    # finally:
    #     connection.close()
    # --- 伪代码结束 ---
    print("MySQL保存操作完成 (伪代码)。")


# --- CSV输出函数 ---
def save_to_csv(dataframe, filename, mode='w', header=True):
    """将DataFrame保存到CSV文件"""
    if not OUTPUT_TO_CSV:
        print("CSV输出已禁用。")
        return
        
    try:
        dataframe.to_csv(filename, index=False, encoding='utf-8', mode=mode, header=header)
        print(f"成功保存 {len(dataframe)} 条记录到 {filename}")
    except Exception as e:
        print(f"保存CSV文件 {filename} 时出错: {e}")

'''
写一个函数，调用上面的函数，传入股票代码和起止日期，计算div_yield的最大值，最小值，均值，最新值，以及标准差，返回pandas的Series
'''
def calculate_dividend_yield_stats(ts_code, start_date=START_DATE, end_date=END_DATE):
    """
    计算股息率的统计指标

    参数:
    ts_code (str): 股票代码
    start_date (str): 起始日期
    end_date (str): 结束日期

    返回:
    pd.Series: 包含股息率统计指标的Series
    """
    df = analyze_stock_dividend_and_price(ts_code, start_date, end_date)

    if df.empty or 'div_yield' not in df.columns:
        return pd.Series(dtype=float)

    div_yield_data = df['div_yield']

    stats = pd.Series({
        'ts_code': ts_code,
        'max': div_yield_data.max(),
        'min': div_yield_data.min(),
        'mean': div_yield_data.mean(),
        'latest': div_yield_data.iloc[-1] if len(div_yield_data) > 0 else 0,
        'std': div_yield_data.std(),
        'zero_count': (div_yield_data == 0).sum()
    })

    return stats

# --- 主处理函数 ---
def process_stock_dividend_batch(start_dt='20200101', end_dt='20251231', dataInterval=10, exchange='SSE'):
    """
    批量处理股票股息率数据，并分批输出。
    """
    print(f"开始处理 {exchange} 交易所的股票数据，时间范围: {start_dt} - {end_dt}")

    try:
        df_stocks = get_stock_basic(exchange=exchange)
        if df_stocks.empty:
            print("未获取到股票基础数据")
            return
    except Exception as e:
        print(f"获取股票基础数据时出错: {e}")
        return

    results_batch = []
    all_results = []
    processed_count = 0
    total_stocks = len(df_stocks)
    csv_filename = f"{exchange}_dividend_yield_stats_batch.csv"
    csv_first_write = True # 用于控制CSV文件头只写入一次
    
    for idx, row in df_stocks.iterrows():
        stock_code = row['ts_code']
        print(f"正在处理 {stock_code} ({idx+1}/{total_stocks})")

        try:
            raw_stats = calculate_dividend_yield_stats(stock_code, start_date=start_dt, end_date=end_dt)
            
            # 确保我们最终得到一个 Series (统计数据)
            stats_series = None
            if isinstance(raw_stats, pd.Series):
                # 如果直接返回了 Series
                stats_series = raw_stats
            elif isinstance(raw_stats, pd.DataFrame) and not raw_stats.empty:
                # 如果返回了 DataFrame，则取第一行
                stats_series = raw_stats.iloc[0]
            else:
                # 如果返回了空DataFrame, None, 或其他类型
                print(f"股票 {stock_code} 的 calculate_dividend_yield_stats 返回了空数据或无效类型 ({type(raw_stats)})。跳过...")
                continue
            
             
            yesterday = (datetime.datetime.now() - datetime.timedelta(days=1)).strftime("%Y%m%d")
            stockParam = getStockParam(stock_code, START_DATE=yesterday, END_DATE=yesterday)
            # 选择需要的字段与下方合并
            if not stockParam.empty:
                selected_params_dict = {
                    'ts_code': stockParam['ts_code'].iloc[0],
                    'trade_date': stockParam['trade_date'].iloc[0],
                    'total_mv': stockParam['total_mv'].iloc[0],
                    'circ_mv': stockParam['circ_mv'].iloc[0]
                }
            else:
                selected_params_dict = { 'ts_code': stock_code, 'trade_date': None, 'total_mv': None,'circ_mv': None}

            # --- 修改点2：简化合并逻辑 ---
            # 现在 stats_series 已确认是 Series，可以直接合并
            # 使用字典合并确保结果清晰且 stats 覆盖同名项
            combined_dict = {**row.to_dict(), **stats_series.to_dict(), **selected_params_dict}
            combined_series = pd.Series(combined_dict)
            print(f"合并结果: {combined_series}")
            print(f"{stock_code}-{start_dt} -- {end_dt} 处理完成。")
            results_batch.append(combined_series)
            all_results.append(combined_series)

        except Exception as e:
            print(f"处理股票 {stock_code} 时出错: {e}")
            # 可以选择在这里添加错误记录到 results_batch 或 all_results
            continue

        processed_count += 1

        # 检查是否达到批次大小
        if processed_count % dataInterval == 0 and results_batch:
            print(f"已处理 {processed_count} 支股票，达到批次大小 {dataInterval}，开始输出...")
            df_batch = pd.DataFrame(results_batch)
            
            # 输出到CSV (追加模式)
            save_to_csv(df_batch, csv_filename, mode='a', header=csv_first_write)
            if csv_first_write: csv_first_write = False # 之后的批次不再写入header
            
            # 输出到MySQL (伪代码)
            save_to_mysql(df_batch, 'stock_dividend_stats')
            
            # 清空批次缓存
            results_batch = []
            print("--- 批次处理完成 ---")

    # 处理最后一批不足dataInterval的数据
    if results_batch:
        print(f"处理剩余 {len(results_batch)} 支股票...")
        df_final_batch = pd.DataFrame(results_batch)
        save_to_csv(df_final_batch, csv_filename, mode='a', header=csv_first_write)
        save_to_mysql(df_final_batch, 'stock_dividend_stats')
        print("--- 最终批次处理完成 ---")

    # (可选) 将所有结果一次性保存到一个完整的CSV文件
    if all_results:
        final_csv_filename = f"{exchange}_dividend_yield_stats_final.csv"
        print(f"正在保存所有 {len(all_results)} 条结果到 {final_csv_filename}...")
        df_final = pd.DataFrame(all_results)
        save_to_csv(df_final, final_csv_filename, mode='w', header=True)
        save_to_mysql(df_final, 'stock_dividend_stats_final')
        print("所有数据处理并保存完成。")
    else:
        print("无有效结果数据可保存")

# --- 程序入口 ---
if __name__ == "__main__":
    # 可以通过修改这里的参数来调用函数
    process_stock_dividend_batch(start_dt='20150101', end_dt='20251231', dataInterval=10, exchange='SSE')

    process_stock_dividend_batch(start_dt='20150101', end_dt='20251231', dataInterval=10, exchange='SZSE')