selfgroup/group_system/xlsx_io.py

#!/usr/bin/env python3
"""Minimal XLSX reader/writer using only the Python standard library."""

from __future__ import annotations

import re
from datetime import datetime
from pathlib import Path
from xml.etree import ElementTree as ET
from zipfile import ZIP_DEFLATED, ZipFile


MAIN_NS = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"
REL_NS = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"
PKG_REL_NS = "http://schemas.openxmlformats.org/package/2006/relationships"
NS = {"m": MAIN_NS, "r": REL_NS, "p": PKG_REL_NS}


def _column_number(cell_ref: str) -> int:
    letters = re.match(r"([A-Z]+)", cell_ref.upper())
    if not letters:
        return 0
    value = 0
    for letter in letters.group(1):
        value = value * 26 + ord(letter) - 64
    return value


def _column_letter(number: int) -> str:
    result = ""
    while number:
        number, remainder = divmod(number - 1, 26)
        result = chr(65 + remainder) + result
    return result


def _shared_strings(archive: ZipFile) -> list[str]:
    if "xl/sharedStrings.xml" not in archive.namelist():
        return []
    root = ET.fromstring(archive.read("xl/sharedStrings.xml"))
    values = []
    for item in root.findall("m:si", NS):
        values.append("".join(node.text or "" for node in item.iter(f"{{{MAIN_NS}}}t")))
    return values


def _relationships(archive: ZipFile) -> dict[str, str]:
    root = ET.fromstring(archive.read("xl/_rels/workbook.xml.rels"))
    return {item.attrib["Id"]: item.attrib["Target"] for item in root}


def _cell_value(cell, shared: list[str]) -> str:
    cell_type = cell.attrib.get("t")
    inline = cell.find("m:is", NS)
    if cell_type == "inlineStr" and inline is not None:
        return "".join(node.text or "" for node in inline.iter(f"{{{MAIN_NS}}}t")).strip()
    value_node = cell.find("m:v", NS)
    raw = "" if value_node is None else (value_node.text or "")
    if cell_type == "s" and raw:
        return shared[int(raw)].strip()
    return raw.strip()


def read_xlsx_rows(path: str | Path) -> list[dict[str, str]]:
    """Read the first worksheet as dictionaries keyed by header text."""
    path = Path(path)
    with ZipFile(path, "r") as archive:
        shared = _shared_strings(archive)
        relationships = _relationships(archive)
        workbook = ET.fromstring(archive.read("xl/workbook.xml"))
        sheets = workbook.findall("m:sheets/m:sheet", NS)
        if not sheets:
            raise ValueError("Excel 文件中没有工作表")
        relationship_id = sheets[0].attrib[f"{{{REL_NS}}}id"]
        target = relationships[relationship_id].lstrip("/")
        if not target.startswith("xl/"):
            target = "xl/" + target
        worksheet = ET.fromstring(archive.read(target))

    rows: list[dict[int, str]] = []
    for row in worksheet.findall("m:sheetData/m:row", NS):
        values: dict[int, str] = {}
        for cell in row.findall("m:c", NS):
            number = _column_number(cell.attrib.get("r", "A1"))
            value = _cell_value(cell, shared)
            if value != "":
                values[number] = value
        rows.append(values)

    header_row_index = -1
    for index, values in enumerate(rows[:20]):
        texts = set(values.values())
        if "学号" in texts and "姓名" in texts:
            header_row_index = index
            break
    if header_row_index < 0:
        raise ValueError("Excel 中必须包含“学号”和“姓名”表头")

    headers = rows[header_row_index]
    result: list[dict[str, str]] = []
    for values in rows[header_row_index + 1:]:
        row = {headers.get(number, f"列{number}"): value for number, value in values.items()}
        if any(str(value).strip() for value in row.values()):
            result.append(row)
    return result


def _xml_escape(value: object) -> str:
    return (
        str(value)
        .replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


def _sheet_xml(headers: list[str], records: list[dict[str, object]]) -> str:
    all_headers = list(headers)
    all_records = list(records)
    row_xml: list[str] = []
    header_cells = []
    for column_number, header in enumerate(all_headers, 1):
        reference = f"{_column_letter(column_number)}1"
        header_cells.append(
            f'<c r="{reference}" t="inlineStr"><is><t>{_xml_escape(header)}</t></is></c>'
        )
    row_xml.append(f'<row r="1">{"".join(header_cells)}</row>')

    for row_number, record in enumerate(all_records, 2):
        cells = []
        for column_number, header in enumerate(all_headers, 1):
            value = record.get(header, "")
            reference = f"{_column_letter(column_number)}{row_number}"
            if isinstance(value, bool):
                cells.append(f'<c r="{reference}" t="inlineStr"><is><t>{_xml_escape("是" if value else "否")}</t></is></c>')
            elif isinstance(value, (int, float)) and not isinstance(value, bool):
                cells.append(f'<c r="{reference}"><v>{value}</v></c>')
            else:
                text = "" if value is None else str(value)
                preserve = ' xml:space="preserve"' if text != text.strip() else ""
                cells.append(
                    f'<c r="{reference}" t="inlineStr"><is><t{preserve}>{_xml_escape(text)}</t></is></c>'
                )
        row_xml.append(f'<row r="{row_number}">{"".join(cells)}</row>')

    column_xml = []
    for number in range(1, len(all_headers) + 1):
        column_xml.append(f'<col min="{number}" max="{number}" width="18" customWidth="1"/>')

    return f'''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<worksheet xmlns="{MAIN_NS}">
<sheetViews><sheetView workbookViewId="0"/></sheetViews>
<cols>{"".join(column_xml)}</cols>
<sheetData>{"".join(row_xml)}</sheetData>
</worksheet>'''


def write_xlsx(sheets: list[tuple[str, list[str], list[dict[str, object]]]]) -> bytes:
    """Create a valid XLSX in memory. Each item is (sheet name, headers, records)."""
    import posixpath

    used_names = set()
    safe_names = []
    for index, (name, _, _) in enumerate(sheets, 1):
        clean = re.sub(r"[\\/?*\\[\\]:]", "-", name).strip() or f"Sheet{index}"
        clean = clean[:31]
        candidate = clean
        suffix = 2
        while candidate.lower() in used_names:
            tail = f"-{suffix}"
            candidate = clean[: 31 - len(tail)] + tail
            suffix += 1
        used_names.add(candidate.lower())
        safe_names.append(candidate)

    sheet_entries = []
    workbook_entries = []
    worksheet_files = []
    for index, name in enumerate(safe_names, 1):
        filename = f"xl/worksheets/sheet{index}.xml"
        worksheet_files.append((filename, _sheet_xml(sheets[index - 1][1], sheets[index - 1][2])))
        sheet_entries.append(
            f'<sheet name="{_xml_escape(name)}" sheetId="{index}" '
            f'r:id="rId{index}"/>'
        )
        workbook_entries.append(
            f'<Relationship Id="rId{index}" '
            f'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" '
            f'Target="worksheets/sheet{index}.xml"/>'
        )
    style_relationship_id = len(sheets) + 1

    content_overrides = "".join(
        f'<Override PartName="/xl/worksheets/sheet{index}.xml" '
        'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
        for index in range(1, len(sheets) + 1)
    )

    files = {
        "[Content_Types].xml": f'''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/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"/>
<Override PartName="/xl/styles.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.styles+xml"/>{content_overrides}
</Types>''',
        "_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>''',
        "xl/workbook.xml": f'''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<workbook xmlns="{MAIN_NS}" xmlns:r="{REL_NS}">
<sheets>{"".join(sheet_entries)}</sheets>
</workbook>''',
        "xl/_rels/workbook.xml.rels": f'''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">{"".join(workbook_entries)}
<Relationship Id="rId{style_relationship_id}" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles" Target="styles.xml"/>
</Relationships>''',
        "xl/styles.xml": '''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<styleSheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">
<fonts count="2"><font><sz val="11"/><name val="Calibri"/></font><font><b/><sz val="11"/><name val="Calibri"/></font></fonts>
<fills count="2"><fill><patternFill patternType="none"/></fill><fill><patternFill patternType="gray125"/></fill></fills>
<borders count="1"><border/></borders>
<cellStyleXfs count="1"><xf numFmtId="0" fontId="0" fillId="0" borderId="0"/></cellStyleXfs>
<cellXfs count="2"><xf numFmtId="0" fontId="0" fillId="0" borderId="0" xfId="0"/><xf numFmtId="0" fontId="1" fillId="0" borderId="0" xfId="0" applyFont="1"/></cellXfs>
</styleSheet>''',
    }
    files.update(worksheet_files)

    import io
    output = io.BytesIO()
    with ZipFile(output, "w", ZIP_DEFLATED) as archive:
        for filename, content in files.items():
            archive.writestr(filename, content.encode("utf-8"))
    return output.getvalue()


def export_filename(prefix: str, class_name: str | None = None) -> str:
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    suffix = f"-{class_name}班" if class_name else ""
    return f"{prefix}{suffix}-{timestamp}.xlsx"