schedule/tests/xlsx_builder.py

"""Минимальный сборщик xlsx в памяти для тестов (без openpyxl).

Поддерживает: shared strings (включая красные ранки), цифры, мержи,
скрытые столбцы. Формат достаточно строгий, чтобы покрывать реальный
парсер (он читает именно этот поднабор OOXML).
"""

from __future__ import annotations

import io
import zipfile
from dataclasses import dataclass, field
from xml.sax.saxutils import escape


def col_letter(idx: int) -> str:
    """0-based индекс столбца -> 'A', 'B', ..., 'AA'."""
    s = ""
    idx += 1
    while idx:
        idx, rem = divmod(idx - 1, 26)
        s = chr(ord("A") + rem) + s
    return s


@dataclass
class SheetSpec:
    rows: list[list]
    merges: list[str] = field(default_factory=list)
    hidden_cols: list[tuple[int, int]] = field(default_factory=list)  # 0-based, инклязивно
    hidden_rows: list[int] = field(default_factory=list)  # 0-based


def build_workbook(sheets: dict[str, SheetSpec]) -> bytes:
    shared: list[str] = []
    shared_index: dict[str, int] = {}
    red_runs: dict[int, list[tuple[str, bool]]] = {}

    def intern(value) -> tuple[str, str, object]:
        """-> (тип ячейки, атрибут t, значение v / None)"""
        if value is None or value == "":
            return None, "", None
        if isinstance(value, (int, float)):
            return "n", "", str(value)
        if isinstance(value, list):
            idx = len(shared)
            shared.append(None)
            red_runs[idx] = value
            return "s", "s", idx
        idx = shared_index.get(value)
        if idx is None:
            idx = len(shared)
            shared_index[value] = idx
            shared.append(value)
        return "s", "s", idx

    sheet_xml: dict[str, str] = {}
    rels = []
    sheet_defs = []
    for i, (name, spec) in enumerate(sheets.items(), start=1):
        rel_id = f"rId{i}"
        rels.append(
            f'<Relationship Id="{rel_id}" '
            'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" '
            f'Target="worksheets/sheet{i}.xml"/>'
        )
        sheet_defs.append(
            f'<sheet name="{escape(name)}" sheetId="{i}" r:id="{rel_id}"/>'
        )

        cols = "".join(
            f'<col min="{a + 1}" max="{b + 1}" hidden="1"/>'
            for a, b in spec.hidden_cols
        )
        cols_xml = f"<cols>{cols}</cols>" if cols else ""

        row_xml = []
        for r, row in enumerate(spec.rows, start=1):
            cells = []
            for c, value in enumerate(row):
                kind, t_attr, v = intern(value)
                if kind is None:
                    continue
                t = f' t="{t_attr}"' if t_attr else ""
                cells.append(f'<c r="{col_letter(c)}{r}"{t}><v>{escape(str(v))}</v></c>')
            if cells:
                hidden = ' hidden="1"' if (r - 1) in spec.hidden_rows else ""
                row_xml.append(f'<row r="{r}"{hidden}>{"".join(cells)}</row>')

        merges_xml = ""
        if spec.merges:
            mc = "".join(f'<mergeCell ref="{m}"/>' for m in spec.merges)
            merges_xml = f'<mergeCells count="{len(spec.merges)}">{mc}</mergeCells>'

        sheet_xml[f"sheet{i}.xml"] = (
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
            f"{cols_xml}<sheetData>{''.join(row_xml)}</sheetData>{merges_xml}"
            "</worksheet>"
        )

    sst_parts = []
    for i, value in enumerate(shared):
        if i in red_runs:
            runs = "".join(
                f'<r>{("<rPr><color rgb=\"FFFF0000\"/></rPr>" if red else "")}'
                f"<t xml:space='preserve'>{escape(text)}</t></r>"
                for text, red in red_runs[i]
            )
            sst_parts.append(f"<si>{runs}</si>")
        else:
            sst_parts.append(f"<si><t xml:space='preserve'>{escape(value)}</t></si>")
    sst = (
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
        '<sst xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
        f'count="{len(shared)}" uniqueCount="{len(shared)}">'
        f"{''.join(sst_parts)}</sst>"
    )

    content_types = ["<Default Extension=\"rels\" ContentType=\"application/vnd.openxmlformats-package.relationships+xml\"/>",
                     "<Default Extension=\"xml\" ContentType=\"application/xml\"/>",
                     '<Override PartName="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>']
    for i in range(1, len(sheets) + 1):
        content_types.append(
            f'<Override PartName="/xl/worksheets/sheet{i}.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
        )
    if shared:
        content_types.append(
            '<Override PartName="/xl/sharedStrings.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sharedStrings+xml"/>'
        )

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr(
            "[Content_Types].xml",
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
            + "".join(content_types)
            + "</Types>",
        )
        zf.writestr(
            "_rels/.rels",
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
            '<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>'
            "</Relationships>",
        )
        zf.writestr(
            "xl/workbook.xml",
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
            'xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">'
            f"<sheets>{''.join(sheet_defs)}</sheets></workbook>",
        )
        zf.writestr(
            "xl/_rels/workbook.xml.rels",
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
            + "".join(rels)
            + "</Relationships>",
        )
        for fname, xml in sheet_xml.items():
            zf.writestr(f"xl/worksheets/{fname}", xml)
        if shared:
            zf.writestr("xl/sharedStrings.xml", sst)
    return buf.getvalue()