large-order-fusion/test.py

import os
from typing import Dict, List
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, PatternFill
from openpyxl.drawing.image import Image
import logging

logger = logging.getLogger(__name__)

class ExcelExporter:
    """Excel导出器，用于将商品组合数据导出为Excel格式"""
    
    def __init__(self):
        self.wb = None
        self.ws = None
        self.styles = {
            'GREEN_FONT': Font(color="FF0000"),  # 注意：这里颜色定义可能有误，应该是00FF00为绿色
            'BLUE_FONT': Font(color="0000FF"),
            'RED_FONT': Font(color="00FF00"),    # 注意：这里应该是FF0000为红色
            'BOLD_FONT': Font(bold=True),
            'CENTER_ALIGN': Alignment(horizontal='center', vertical='center'),
            'RED_FILL': PatternFill(start_color="FF9999", end_color="FF9999", fill_type="solid"),
            'GREEN_FILL': PatternFill(start_color="99FF99", end_color="99FF99", fill_type="solid")
        }
    
    def export_to_excel(self, data: Dict, filename: str = None) -> str:
        """
        将商品数据导出到Excel文件
        
        Args:
            data: 包含商品信息的字典
            filename: 输出的Excel文件名
            
        Returns:
            导出的文件路径
        """
        try:
            self._initialize_workbook()
            self._extract_data(data)
            self._setup_worksheet()
            self._write_data()
            return self._save_workbook(filename)
        except Exception as e:
            logger.error(f"导出Excel失败: {e}")
            raise
    
    def _initialize_workbook(self):
        """初始化工作簿和工作表"""
        self.wb = Workbook()
        self.ws = self.wb.active
        self.ws.title = "商品数据"
    
    def _extract_data(self, data: Dict):
        """从数据中提取所需信息"""
        self.primary_good = data["primary_goods_info"]["goods_info"]
        self.combine_two = data["combine_two_info"]
        self.combine_three = data["combine_three_info"]
        self.outfit_orders = data["outfit_orders"]
        self.primary_code = self.primary_good['sku'] + self.primary_good['color_code']
        self.header_count = [len(order) for order in self.outfit_orders]
    
    def _setup_worksheet(self):
        """设置工作表基本属性"""
        self._setup_column_dimensions()
        self._write_headers()
        self.ws.freeze_panes = 'A2'
    
    def _setup_column_dimensions(self):
        """设置列宽"""
        # 图片列
        for col in ['B', 'C', 'D']:
            self.ws.column_dimensions[col].width = 13.6
        
        # 数据列
        for col in ['A', *[self._get_column_letter(i) for i in range(4, 58)]]:
            self.ws.column_dimensions[col].width = 14 if col != 'A' else 40
    
    def _write_headers(self):
        """写入表头"""
        self._write_main_headers()
        self._write_sub_headers()
    
    def _write_main_headers(self):
        """写入主表头"""
        headers = [
            "组合", "主推款图片", "连带款1图片", "连带款2图片", "连带排名", 
            "连带频次", "主商品频次", "连带率", "连带件数", "连带金额", ""
        ]
        
        for col_idx, header in enumerate(headers, start=1):
            cell = self.ws.cell(row=1, column=col_idx, value=header)
            cell.font = self.styles['BOLD_FONT']
            cell.alignment = self.styles['CENTER_ALIGN']
    
    def _write_sub_headers(self):
        """写入子表头（大单组合）"""
        # ... 原有的子表头逻辑，略作调整使用self.styles
    
    def _write_data(self):
        """写入所有数据"""
        self._write_combine_two_data()
        self._write_combine_three_data()
        self._write_outfit_orders_data()
        self._format_worksheet()
    
    def _write_combine_two_data(self):
        """写入两种商品组合数据"""
        for row_idx, item in enumerate(self.combine_two, start=2):
            self._write_combine_two_row(row_idx, item)
    
    def _write_combine_two_row(self, row_idx: int, item: Dict):
        """写入单行两种商品组合数据"""
        goods_info = item['goods_info']
        sub_code = goods_info['sku'] + goods_info['color_code']
        
        # 写入数据
        data_mapping = [
            (1, f"{self.primary_code},{sub_code}"), (4, "22组合"),
            (5, row_idx-1), (6, item['freq']), (7, item['main_freq']),
            (8, item['combine_rate']), (9, item['count']), (10, item['sales'])
        ]
        
        for col, value in data_mapping:
            cell = self.ws.cell(row=row_idx, column=col, value=value)
            cell.font = self.styles['GREEN_FONT']
        
        # 插入图片
        self._insert_image(self.primary_good['image_path'], f'B{row_idx}')
        self._insert_image(goods_info['image_path'], f'C{row_idx}')
    
    def _write_combine_three_data(self):
        """写入三种商品组合数据"""
        start_row = 2 + len(self.combine_two)
        for row_idx, item in enumerate(self.combine_three, start=start_row):
            self._write_combine_three_row(row_idx, item, start_row)
    
    def _write_combine_three_row(self, row_idx: int, item: Dict, start_row: int):
        """写入单行三种商品组合数据"""
        goods_info = item['goods_info']
        sub1_code = goods_info[0]['sku'] + goods_info[0]['color_code']
        sub2_code = goods_info[1]['sku'] + goods_info[1]['color_code']
        
        data_mapping = [
            (1, f"{self.primary_code},{sub1_code},{sub2_code}"),
            (5, row_idx - (1 + len(self.combine_two))), (6, item['freq']),
            (7, item['main_freq']), (8, item['combine_rate']),
            (9, item['count']), (10, item['sales'])
        ]
        
        for col, value in data_mapping:
            cell = self.ws.cell(row=row_idx, column=col, value=value)
            cell.font = self.styles['BLUE_FONT']
        
        # 插入图片
        self._insert_image(self.primary_good['image_path'], f'B{row_idx}')
        self._insert_image(goods_info[0]['image_path'], f'C{row_idx}')
        self._insert_image(goods_info[1]['image_path'], f'D{row_idx}')
    
    def _write_outfit_orders_data(self):
        """写入大单组合数据"""
        # ... 原有逻辑，略作调整
    
    def _insert_image(self, image_path: str, cell_location: str):
        """在指定单元格插入图片"""
        try:
            if os.path.exists(image_path):
                img = Image(image_path)
                img.width = img.height = 100
                self.ws.add_image(img, cell_location)
        except Exception as e:
            logger.warning(f"插入图片失败 {image_path}: {e}")
    
    def _format_worksheet(self):
        """格式化工作表"""
        # 设置百分比格式
        for cell in self.ws['H']:
            cell.number_format = '0.00%'
        
        # 设置行高
        data_rows_count = len(self.combine_two) + len(self.combine_three)
        for row in range(2, data_rows_count + 2):
            self.ws.row_dimensions[row].height = 80
        
        # 设置居中对齐
        for row in self.ws.iter_rows():
            for cell in row:
                cell.alignment = self.styles['CENTER_ALIGN']
    
    def _save_workbook(self, filename: str) -> str:
        """保存工作簿到文件"""
        if not filename:
            filename = f"{self.primary_code}.xlsx"
        
        os.makedirs('./output', exist_ok=True)
        file_path = os.path.join('./output', filename)
        self.wb.save(file_path)
        logger.info(f"数据已成功写入 {file_path} 文件")
        return file_path
    
    @staticmethod
    def _get_column_letter(idx: int) -> str:
        """获取列字母（简化版本）"""
        from openpyxl.utils import get_column_letter
        return get_column_letter(idx)

# 使用示例
def export_to_excel(data: Dict, filename: str = None) -> str:
    """保持原有函数接口的包装函数"""
    exporter = ExcelExporter()
    return exporter.export_to_excel(data, filename)