Excel (*.xlsx) data drawing

from openpyxl import load_workbook
from openpyxl.drawing.image import Image
import matplotlib.pyplot as plt
import pandas as pd
import os
import glob

plt.rcParams['font.sans-serif']=['SimHei'] #图像能显示中文
plt.rcParams["axes.unicode_minus"] = False #确保负坐标轴能显示负号


def excel_tu(fliename,sheetname,header,zhi,mainheadname,subheadname,colors):
    
    # 打开Excel文件并选择工作表
    wb = load_workbook(fliename)
    ws = wb[sheetname]
    df = pd.read_excel(fliename,sheetname)
    
    xx=[]
    for i in range(1,len(df[header[0]])+1):
        xx.append(i)
    
    # 绘制曲线统计图   
    fig, ax1 = plt.subplots(figsize=(6,4),dpi=300)
    
    # 设置坐标轴的值    
    ax1.set_ylim(bottom=zhi[0], top=zhi[1])

    for col_name in header:
        for i,mainhead in enumerate(mainheadname):
            if col_name== mainhead:                
                col_data = df[col_name]                
                # 绘制曲线
                ax1.plot(xx, col_data, color=colors[i], label=col_name)
            

    # 添加次坐标轴并绘制第二条曲线
    color_two=len(mainheadname)
    ax2 = ax1.twinx()
    
    # 设置坐标轴的最小值    
    ax2.set_ylim(bottom=zhi[2], top=zhi[3])

    
    for col_name in header:
        for i,subhead in enumerate(subheadname):
            if col_name== subhead:
                col_data = df[col_name]                
                # 绘制曲线
                ax2.plot(xx, col_data, color=colors[i+color_two], label=col_name)
    

    
    # 显示图例
    ax1.legend(loc='upper center')
    handles1, labels1 = ax1.get_legend_handles_labels()
    handles2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(handles1 + handles2, labels1 + labels2, loc='upper center',ncol=6,fontsize=7)

    # 设置标题
    ax1.set_title(fliename.split('.')[0].split('.')[0].split('\\')[-1])
    
    # 保存图像并将其插入到Excel工作表中
    fig.savefig('plot.png')
    img = Image('plot.png')
    ws.add_image(img, 'D1')

    # 保存Excel文件
    wb.save(fliename)
    

if __name__ == "__main__":
    
    #读取文件夹中所有的表格数据     
    folder_path = 'D:\临时资料\写报告\四水库单库调度\单库结果'

    # 使用glob模块查找所有的Excel文件
    excel_files = glob.glob(os.path.join(folder_path, '*.xlsx'))
    

    # 画图    
    sheetname="Sheet1"
    mainheadname=["入流","出流"]
    subheadname=["水位","汛限水位","防洪高水位","校核洪水位"]
    header=['时间', '入流','出流','水位','库容','汛限水位','防洪高水位','校核洪水位']
    colors=['blue','green','red','purple','brown','gray','black']
    
    # 坐标轴设置 zhi=[主坐标轴最小值，主坐标轴最大值，次坐标轴最小值，次坐标轴最大值]
    zhi=[0,900,226,245]
    for excel_file in excel_files:
        excel_tu(excel_file,sheetname,header,zhi,mainheadname,subheadname,colors)
        print('wanc')