messyscript/小分表/python版/main.py

#拉小分表工具（即分类汇总）
import sys
import pandas as pd
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QPushButton,
                            QFileDialog, QTableWidget, QTableWidgetItem, QSpinBox,
                            QLabel, QProgressBar, QCheckBox, QLineEdit, QHBoxLayout, QComboBox)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QColor
import matplotlib.pyplot as plt
import ast
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class ExcelProcessor(QWidget):
    def __init__(self):
        super().__init__()
        # 先定义属性，再初始化UI
        self.df = pd.DataFrame()
        self.threshold = 10
        self.ignore_cols = [1,2,3,8,9]  # 默认忽略列索引
        self.current_file_path = None  # 添加当前文件路径属性
        self.color_groups = []  # 存储班级颜色分组
        self.current_result_df = None  # 存储当前处理结果
        self.current_group_col = None  # 存储当前分组列名
        self.initUI()

    def initUI(self):
        self.setWindowTitle('Excel数据处理工具')
        self.setGeometry(300, 300, 800, 600)

        layout = QVBoxLayout()

        # 文件选择
        self.btn_open = QPushButton('选择Excel文件', self)
        self.btn_open.clicked.connect(self.open_file)
        layout.addWidget(self.btn_open)

        # 文件路径显示
        self.lbl_file_path = QLabel('未选择文件', self)
        layout.addWidget(self.lbl_file_path)

        # 忽略列设置
        self.ignore_layout = QHBoxLayout()
        self.lbl_ignore = QLabel('忽略列索引:', self)
        self.txt_ignore = QLineEdit('1,2,3,8,9', self)
        self.ignore_layout.addWidget(self.lbl_ignore)
        self.ignore_layout.addWidget(self.txt_ignore)
        layout.addLayout(self.ignore_layout)

        # 分组列选择
        self.group_layout = QHBoxLayout()
        self.lbl_group = QLabel('分组列:', self)
        self.cmb_group = QComboBox(self)
        self.group_layout.addWidget(self.lbl_group)
        self.group_layout.addWidget(self.cmb_group)
        layout.addLayout(self.group_layout)

        # 阈值设置
        self.threshold_layout = QVBoxLayout()
        self.lbl_threshold = QLabel('总分阈值:', self)
        self.spin_threshold = QSpinBox(self)
        self.spin_threshold.setValue(self.threshold)  # 现在self.threshold已经定义了
        self.threshold_layout.addWidget(self.lbl_threshold)
        self.threshold_layout.addWidget(self.spin_threshold)
        layout.addLayout(self.threshold_layout)
        
        # 班级颜色分组设置
        self.color_layout = QVBoxLayout()
        self.lbl_color = QLabel('班级颜色分组 (格式: [1,2,30,31],[3,4,5,6,32,33,34,35],...):', self)
        self.txt_color = QLineEdit('[1,2,30,31],[3,4,5,6,32,33,34,35],[7,8,9,10,36,37,38,39],[12,13,14,42,43,44],[23,24,41]', self)
        self.color_layout.addWidget(self.lbl_color)
        self.color_layout.addWidget(self.txt_color)
        layout.addLayout(self.color_layout)

        # 处理选项
        self.chk_header = QCheckBox('包含表头', self)
        self.chk_header.setChecked(True)
        self.chk_header.stateChanged.connect(self.on_header_changed)  # 添加状态改变信号连接
        layout.addWidget(self.chk_header)

        # 处理按钮
        self.btn_process = QPushButton('开始处理', self)
        self.btn_process.clicked.connect(self.process_data)
        self.btn_process.setEnabled(False)  # 初始时禁用处理按钮
        layout.addWidget(self.btn_process)

        # 进度条
        self.progress = QProgressBar(self)
        self.progress.setValue(0)  # 初始化为0
        layout.addWidget(self.progress)

        # 结果显示
        self.table_result = QTableWidget(self)
        layout.addWidget(self.table_result)

        self.setLayout(layout)

    def open_file(self):
        path, _ = QFileDialog.getOpenFileName(self, '选择Excel文件', '', 'Excel Files (*.xlsx)')
        if path:
            self.current_file_path = path  # 保存文件路径
            self.load_file()

    def load_file(self):
        """加载文件的方法，供open_file和on_header_changed调用"""
        if self.current_file_path:
            try:
                # 根据复选框状态决定读取方式
                if self.chk_header.isChecked():
                    # 包含表头：忽略第一行，将第二行视为标题行，第三行开始为数据
                    self.df = pd.read_excel(self.current_file_path, header=1)  # header=1 表示第二行作为标题
                else:
                    # 不包含表头：第一行视为标题行，第二行开始为数据
                    self.df = pd.read_excel(self.current_file_path, header=0)  # header=0 表示第一行作为标题
                    
                self.lbl_file_path.setText(f'已选择文件: {self.current_file_path.split("/")[-1]}')  # 只显示文件名
                self.progress.setValue(0)  # 重置进度条
                
                # 更新分组列下拉框
                self.update_group_combo()
                
                # 启用处理按钮
                self.btn_process.setEnabled(True)
                
            except Exception as e:
                print(f"文件读取错误: {str(e)}")

    def on_header_changed(self):
        """当包含表头复选框状态改变时调用"""
        if self.current_file_path:  # 如果已经选择了文件
            self.load_file()  # 重新加载文件

    def update_group_combo(self):
        """更新分组列下拉框选项"""
        self.cmb_group.clear()
        if not self.df.empty:
            # 添加所有列名
            for col in self.df.columns:
                self.cmb_group.addItem(str(col))
            
            # 尝试自动选择"班级"列
            group_keywords = ['班级', '班', 'class', 'Class', 'CLASSS']
            for keyword in group_keywords:
                for i in range(self.cmb_group.count()):
                    if keyword in self.cmb_group.itemText(i):
                        self.cmb_group.setCurrentIndex(i)
                        return

    def parse_color_groups(self):
        """解析班级颜色分组设置"""
        color_text = self.txt_color.text().strip()
        if not color_text:
            return []
        
        try:
            # 使用ast.literal_eval安全地解析字符串
            # 首先将整个字符串包装在列表中，然后解析
            parsed = ast.literal_eval(f'[{color_text}]')
            
            # 确保结果是列表的列表
            color_groups = []
            for item in parsed:
                if isinstance(item, list):
                    color_groups.append(item)
                elif isinstance(item, int):
                    color_groups.append([item])
                else:
                    print(f"忽略无效的分组项: {item}")
            
            return color_groups
        except Exception as e:
            print(f"解析班级颜色分组时出错: {str(e)}")
            return []

    def get_class_color(self, class_value):
        """根据班级值获取对应的颜色索引"""
        # 尝试将班级值转换为整数进行比较
        try:
            class_num = int(class_value)
        except (ValueError, TypeError):
            # 如果无法转换为整数，尝试提取数字
            try:
                import re
                numbers = re.findall(r'\d+', str(class_value))
                if numbers:
                    class_num = int(numbers[0])
                else:
                    class_num = None
            except:
                class_num = None
        
        # 查找班级所属的颜色组
        color_index = -1
        if class_num is not None:
            for i, group in enumerate(self.color_groups):
                if class_num in group:
                    color_index = i
                    break
        
        return color_index

    def process_data(self):
        if self.df.empty:
            return

        # 重置进度条
        self.progress.setValue(0)

        # 解析忽略列
        try:
            self.ignore_cols = list(map(int, self.txt_ignore.text().split(',')))
        except:
            self.ignore_cols = [1,2,3,8,9]

        # 解析班级颜色分组
        self.color_groups = self.parse_color_groups()

        # 数据清洗 - 先找到用户选择的分组列的位置
        selected_group_col = self.cmb_group.currentText()
        group_col_index = None
        for i, col in enumerate(self.df.columns):
            if col == selected_group_col:
                group_col_index = i
                break
        
        # 确保分组列不被忽略
        if group_col_index is not None and group_col_index in self.ignore_cols:
            self.ignore_cols.remove(group_col_index)

        valid_cols = [i for i in range(self.df.shape[1]) if i not in self.ignore_cols]
        self.df = self.df.iloc[:, valid_cols]

        # 更新分组列下拉框（因为列索引可能已改变）
        self.update_group_combo()
        
        # 重新获取分组列名（因为列索引可能已改变）
        selected_group_col = self.cmb_group.currentText()

        # 删除低分行
        score_col = [col for col in self.df.columns if '总分' in str(col)]
        if score_col:
            self.df = self.df[pd.to_numeric(self.df[score_col[0]], errors='coerce').notnull()]
            self.df = self.df[self.df[score_col[0]] >= self.spin_threshold.value()]
            self.progress.setValue(40)

        # 分类汇总
        if len(valid_cols) == 0:
            return

        # 使用用户选择的分组列
        group_col_name = selected_group_col
        
        # 如果用户选择的列不存在，使用备选方案
        if group_col_name not in self.df.columns:
            # 查找可能的分组列（班级、班、class等）
            group_keywords = ['班级', '班', 'class', 'Class', 'CLASSS']
            group_col_name = None
            
            for keyword in group_keywords:
                for col in self.df.columns:
                    if keyword in str(col):
                        group_col_name = col
                        break
                if group_col_name:
                    break
            
            # 如果还是找不到，使用第一列作为分组列
            if group_col_name is None:
                group_col_name = self.df.columns[0]
                print(f"警告：未找到班级列，使用第一列 '{group_col_name}' 进行分组")

        numeric_cols = self.df.select_dtypes(include='number').columns.tolist()
        
        # 确保分组列不在数值列中
        if group_col_name in numeric_cols:
            numeric_cols.remove(group_col_name)
        
        if len(numeric_cols) == 0:
            print("错误：没有找到数值列进行汇总")
            return

        grouped = self.df.groupby(group_col_name)[numeric_cols].mean().reset_index()
        
        # 保留两位小数
        grouped[numeric_cols] = grouped[numeric_cols].round(2)
        
        grouped = grouped.sort_values(by=group_col_name)
        self.progress.setValue(60)

        # 保存处理结果
        self.current_result_df = grouped
        self.current_group_col = group_col_name

        # 显示结果
        self.display_result(grouped, group_col_name)

    def display_result(self, df, group_col_name):
        self.table_result.setRowCount(len(df))
        self.table_result.setColumnCount(len(df.columns))
        
        # 设置表头
        self.table_result.setHorizontalHeaderLabels(df.columns)
        
        # 预定义颜色列表
        colors = [
            QColor(255, 200, 200),  # 浅红色
            QColor(200, 255, 200),  # 浅绿色
            QColor(200, 200, 255),  # 浅蓝色
            QColor(255, 255, 200),  # 浅黄色
            QColor(255, 200, 255),  # 浅紫色
            QColor(200, 255, 255),  # 浅青色
            QColor(255, 220, 200),  # 浅橙色
            QColor(220, 200, 255),  # 浅紫罗兰色
        ]
        
        # 填充数据并设置颜色
        for row in range(len(df)):
            class_value = df.iloc[row][group_col_name]
            color_index = self.get_class_color(class_value)
            
            for col in range(len(df.columns)):
                item = QTableWidgetItem(str(df.iloc[row, col]))
                
                # 如果班级在颜色分组中，设置背景色
                if color_index >= 0 and color_index < len(colors):
                    item.setBackground(colors[color_index % len(colors)])
                
                self.table_result.setItem(row, col, item)
        
        # 完成进度
        self.progress.setValue(100)
        
        # 导出按钮
        if hasattr(self, 'btn_export'):
            self.layout().removeWidget(self.btn_export)
            self.btn_export.deleteLater()
            
        self.btn_export = QPushButton('导出结果', self)
        self.btn_export.clicked.connect(lambda: self.export_data(df))
        self.layout().addWidget(self.btn_export)

    def export_data(self, df):
        path, _ = QFileDialog.getSaveFileName(self, '保存结果', '', 
                                           'Excel Files (*.xlsx);;图片文件 (*.png)')
        if path:
            if path.endswith('.xlsx'):
                self.export_to_excel(df, path)
            elif path.endswith('.png'):
                self.save_as_image(df, path)

    def export_to_excel(self, df, path):
        """导出到Excel并应用颜色和边框"""
        try:
            # 检查是否安装了openpyxl
            try:
                import openpyxl
                from openpyxl.styles import PatternFill, Border, Side
            except ImportError:
                print("未安装openpyxl，无法导出带颜色的Excel文件")
                df.to_excel(path, index=False)
                print(f"已导出无颜色的Excel文件: {path}")
                return
            
            # 使用openpyxl引擎，以便设置单元格样式
            with pd.ExcelWriter(path, engine='openpyxl') as writer:
                df.to_excel(writer, index=False, sheet_name='Sheet1')
                
                # 获取工作表
                workbook = writer.book
                worksheet = writer.sheets['Sheet1']
                
                # 定义颜色映射（RGB值）
                colors_rgb = [
                    'FFFFC8C8',  # 浅红色
                    'FFC8FFC8',  # 浅绿色
                    'FFC8C8FF',  # 浅蓝色
                    'FFFFFFC8',  # 浅黄色
                    'FFFFC8FF',  # 浅紫色
                    'FFC8FFFF',  # 浅青色
                    'FFFFDCC8',  # 浅橙色
                    'FFDCC8FF',  # 浅紫罗兰色
                ]
                
                # 定义边框样式
                thin_border = Border(
                    left=Side(style='thin'),
                    right=Side(style='thin'),
                    top=Side(style='thin'),
                    bottom=Side(style='thin')
                )
                
                # 为每一行设置背景色和边框
                for row_idx in range(1, len(df) + 2):  # 从第1行（表头）开始到数据行结束
                    # 设置背景色（仅对数据行，不包括表头）
                    if row_idx >= 2:  # 数据行从第2行开始
                        class_value = df.iloc[row_idx-2][self.current_group_col]
                        color_index = self.get_class_color(class_value)
                        
                        if color_index >= 0 and color_index < len(colors_rgb):
                            fill_color = colors_rgb[color_index % len(colors_rgb)]
                            
                            # 为这一行的所有单元格设置背景色
                            for col_idx in range(1, len(df.columns) + 1):
                                cell = worksheet.cell(row=row_idx, column=col_idx)
                                cell.fill = PatternFill(start_color=fill_color, 
                                                       end_color=fill_color, 
                                                       fill_type='solid')
                    
                    # 为所有单元格设置边框（包括表头和数据行）
                    for col_idx in range(1, len(df.columns) + 1):
                        cell = worksheet.cell(row=row_idx, column=col_idx)
                        cell.border = thin_border
                
                # 自动调整列宽
                for column in worksheet.columns:
                    max_length = 0
                    column_letter = column[0].column_letter
                    for cell in column:
                        try:
                            if len(str(cell.value)) > max_length:
                                max_length = len(str(cell.value))
                        except:
                            pass
                    adjusted_width = (max_length + 2)
                    worksheet.column_dimensions[column_letter].width = adjusted_width
                
            print(f"带颜色和边框的Excel文件已保存到: {path}")
        except Exception as e:
            print(f"导出带颜色的Excel时出错: {str(e)}")
            # 如果带颜色的导出失败，尝试普通导出
            try:
                df.to_excel(path, index=False)
                print(f"Excel文件已保存到: {path} (无颜色和边框)")
            except Exception as e2:
                print(f"普通Excel导出也失败: {str(e2)}")

    def save_as_image(self, df, path):
        import matplotlib.pyplot as plt
        import numpy as np

        fig, ax = plt.subplots(figsize=(12, 8))
        ax.axis('off')
        
        # 获取数据
        cell_text = []
        for i in range(len(df)):
            row_data = []
            for j in range(len(df.columns)):
                row_data.append(str(df.iloc[i, j]))
            cell_text.append(row_data)
        
        col_labels = [str(col) for col in df.columns]
        
        # 预定义颜色
        colors_mpl = [
            [1.0, 0.8, 0.8],    # 浅红色
            [0.8, 1.0, 0.8],    # 浅绿色
            [0.8, 0.8, 1.0],    # 浅蓝色
            [1.0, 1.0, 0.8],    # 浅黄色
            [1.0, 0.8, 1.0],    # 浅紫色
            [0.8, 1.0, 1.0],    # 浅青色
            [1.0, 0.86, 0.8],   # 浅橙色
            [0.86, 0.8, 1.0],   # 浅紫罗兰色
        ]
        
        # 构建颜色数组 - 关键修复：只包含数据行，不包括表头
        cell_colors = []
        
        # 只添加数据行颜色（不包括表头）
        for i in range(len(df)):
            class_value = df.iloc[i][self.current_group_col]
            color_index = self.get_class_color(class_value)
            
            if 0 <= color_index < len(colors_mpl):
                row_color = colors_mpl[color_index]
            else:
                row_color = [1, 1, 1]  # 默认白色
                
            row_colors = [row_color for _ in range(len(df.columns))]
            cell_colors.append(row_colors)
        
        # 调试信息
        print(f"数据行数: {len(df)}")
        print(f"cell_text 行数: {len(cell_text)}")
        print(f"cell_colors 行数: {len(cell_colors)}")
        
        # 确保行数匹配 - 现在应该都是7行
        if len(cell_colors) != len(cell_text):
            print(f"警告: 颜色行数 ({len(cell_colors)}) 与数据行数 ({len(cell_text)}) 不匹配")
            # 强制调整颜色数组大小
            if len(cell_colors) > len(cell_text):
                cell_colors = cell_colors[:len(cell_text)]
            else:
                while len(cell_colors) < len(cell_text):
                    cell_colors.append([[1, 1, 1] for _ in range(len(df.columns))])
        
        # 创建表格 - 关键修复：使用colColours参数单独设置表头颜色
        table = ax.table(
            cellText=cell_text,
            colLabels=col_labels,
            cellLoc='center',
            loc='center',
            cellColours=cell_colors,
            colColours=[[1, 1, 1] for _ in range(len(df.columns))]  # 表头颜色设为白色
        )
        
        table.auto_set_font_size(False)
        table.set_fontsize(10)
        table.scale(1, 1.5)
        
        # 设置单元格边框
        for key, cell in table.get_celld().items():
            cell.set_linewidth(0.5)
            cell.set_edgecolor('black')
        
        plt.savefig(path, bbox_inches='tight', dpi=300)
        plt.close()
        print(f"图片已保存到: {path}")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = ExcelProcessor()
    ex.show()
    sys.exit(app.exec_())