QingLong/AI/Text2Sql/Util/EchartsUtil.py

from pyecharts.charts import Bar, Pie
from pyecharts import options as opts
from pyecharts.commons.utils import JsCode
import os

def generate_chart(
    _data,
    chart_type: str,  # 图表类型：'bar' 或 'pie'
    title: str,  # 图表标题
    category_columns: list,  # 分类数据列（X 轴或饼图标签）
    value_column: str,  # 数值数据列（Y 轴或饼图值）
    output_file: str = "chart.html"  # 输出文件路径
):
    """
    统一生成柱形图或饼形图。

    参数:
        _data: 结果集，通常是一个列表（字典格式）。
        chart_type (str): 图表类型，'bar' 或 'pie'。
        title (str): 图表标题。
        category_columns (list): 分类数据列（X 轴或饼图标签）。
        value_column (str): 数值数据列（Y 轴或饼图值）。
        output_file (str): 输出 HTML 文件名（默认 "chart.html"）。
    """
    try:
        # 如果结果集为空，直接返回
        if not _data:
            raise ValueError("结果集为空，无法生成图表！")

        # 检查列名是否存在
        for col in category_columns + [value_column]:
            if col not in _data[0]:
                raise ValueError(f"列名 '{col}' 不存在！")

        # 提取分类数据（动态组合多列）
        categories = [" - ".join(str(row[col]) for col in category_columns) for row in _data]
        # 提取数值数据
        values = [row[value_column] for row in _data]

        # 颜色配置
        colors = [
            "#5470C6", "#91CC75", "#EE6666", "#73C0DE", "#3BA272", "#FC8452", "#9A60B4", "#EA7CCC"
        ]
        color_js = JsCode(f'''
            function(params) {{
                var colorList = {colors};
                return colorList[params.dataIndex % colorList.length];
            }}
        ''')

        # 生成柱形图
        if chart_type == 'bar':
            chart = Bar()
            chart.add_xaxis(categories)
            chart.add_yaxis(
                series_name=value_column,
                y_axis=values,
                bar_width="30%",
                itemstyle_opts=opts.ItemStyleOpts(color=color_js)
            )
            chart.set_global_opts(
                title_opts=opts.TitleOpts(title=title),
                xaxis_opts=opts.AxisOpts(
                    name=" - ".join(category_columns),
                    axislabel_opts=opts.LabelOpts(rotate=45),
                    boundary_gap=True,
                    splitline_opts=opts.SplitLineOpts(is_show=False)
                ),
                yaxis_opts=opts.AxisOpts(name=value_column),
                legend_opts=opts.LegendOpts(is_show=False)
            )

        # 生成饼形图
        elif chart_type == 'pie':
            data_pairs = list(zip(categories, values))
            chart = Pie()
            chart.add(
                series_name="占比",
                data_pair=data_pairs,
                radius=["30%", "55%"],
                label_opts=opts.LabelOpts(
                    formatter="{b}\n{d}%",
                    position="outside",
                    font_size=14
                ),
                itemstyle_opts=opts.ItemStyleOpts(color=color_js)
            )
            chart.set_global_opts(
                title_opts=opts.TitleOpts(
                    title=title,
                    subtitle="数据维度：" + " → ".join(category_columns)
                ),
                legend_opts=opts.LegendOpts(is_show=False)
            )

        else:
            raise ValueError("不支持的图表类型！")

        # 确保目标目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)

        # 保存为 HTML 文件
        # 保存为 HTML 文件，并替换默认标题
        chart.render(
            path=output_file,
            #template_name="simple_chart.html",  # 使用自定义模板
            title="长春云校数据统计"  # 替换默认标题
        )
        print(f"图表已保存为 {output_file}")

        # 检查文件是否存在
        if os.path.exists(output_file):
            print("文件生成成功！")
        else:
            print("文件生成失败！")

    except Exception as e:
        print(f"生成图表时发生错误: {e}")