shanxia2.0/backend/app/utils/excel_export.py

"""Excel 导出工具：集中处理字典 code → label（对齐 Java ExcelDictConvert / ExcelDictPreloadConvert）。"""
from __future__ import annotations

import io
from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence
from urllib.parse import quote

from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.models.models import SysDictData

CellFormatter = Callable[[Any], Any]


@dataclass(frozen=True)
class ExportColumn:
    """导出列定义，dict_type 非空时自动将字典 value 转为 label。"""

    title: str
    field: str
    dict_type: Optional[str] = None
    formatter: Optional[CellFormatter] = None


def load_dict_label_maps(
    db: Session,
    dict_types: Iterable[str],
    tenant_no: Optional[str] = None,
) -> Dict[str, Dict[str, str]]:
    """预加载多个字典类型的 value→label 映射（单次查询，避免 N+1）。"""
    types = sorted({t.strip() for t in dict_types if t and str(t).strip()})
    if not types:
        return {}

    q = db.query(SysDictData).filter(
        SysDictData.del_flag == 0,
        SysDictData.dict_type.in_(types),
    )
    if tenant_no:
        q = q.filter(SysDictData.tenant_no == tenant_no)

    maps: Dict[str, Dict[str, str]] = {t: {} for t in types}
    for row in q.all():
        dtype = row.dict_type
        if dtype not in maps:
            continue
        value = row.dict_value
        if value is None:
            continue
        maps[dtype][str(value)] = row.dict_label or str(value)
    return maps


class DictLabelResolver:
    """导出单元格解析：按列配置做字典翻译或自定义格式化。"""

    def __init__(
        self,
        db: Session,
        columns: Sequence[ExportColumn],
        tenant_no: Optional[str] = None,
        label_maps: Optional[Mapping[str, Mapping[str, str]]] = None,
    ):
        dict_types = [c.dict_type for c in columns if c.dict_type]
        self._maps = dict(label_maps) if label_maps is not None else load_dict_label_maps(
            db, dict_types, tenant_no
        )

    def resolve(self, row: Mapping[str, Any], column: ExportColumn) -> Any:
        raw = row.get(column.field, "")
        if column.formatter is not None:
            return column.formatter(raw)
        if column.dict_type and raw not in (None, ""):
            label = self._maps.get(column.dict_type, {}).get(str(raw))
            if label is not None:
                return label
        if raw is None:
            return ""
        return raw

    def row_to_cells(self, row: Mapping[str, Any], columns: Sequence[ExportColumn]) -> List[Any]:
        return [self.resolve(row, col) for col in columns]


def _build_workbook(
    sheet_title: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    resolver: DictLabelResolver,
):
    from openpyxl import Workbook

    wb = Workbook()
    ws = wb.active
    ws.title = sheet_title
    ws.append([c.title for c in columns])
    for row in rows:
        ws.append(resolver.row_to_cells(row, columns))
    return wb


def workbook_to_streaming_response(
    wb,
    *,
    filename_stem: str,
    include_legacy_filename: bool = True,
) -> StreamingResponse:
    output = io.BytesIO()
    wb.save(output)
    output.seek(0)
    encoded = quote(f"{filename_stem}.xlsx")
    disposition = f"attachment; filename*=UTF-8''{encoded}"
    if include_legacy_filename:
        disposition = f"attachment; filename={encoded}; {disposition}"
    return StreamingResponse(
        output,
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": disposition},
    )


def export_rows_to_xlsx_response(
    db: Session,
    *,
    tenant_no: Optional[str],
    sheet_title: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    filename_stem: str,
    include_legacy_filename: bool = True,
) -> StreamingResponse:
    """CRM 列表导出统一入口：表头 + 行数据 + 字典 label 转换 + 文件流响应。"""
    resolver = DictLabelResolver(db, columns, tenant_no=tenant_no)
    wb = _build_workbook(sheet_title, columns, rows, resolver)
    return workbook_to_streaming_response(
        wb,
        filename_stem=filename_stem,
        include_legacy_filename=include_legacy_filename,
    )