pinmap-to-pinlist/Releases/v1.0.1/source/xlsx_writer.py

"""XLSX writer — pure Python, zero dependencies.

Builds an OOXML .xlsx file using ``zipfile`` + ``xml.etree.ElementTree``.
"""

import zipfile
import xml.etree.ElementTree as ET
from io import BytesIO
from typing import Optional

from utils import col_to_letter, rc_to_cell_ref


def write_xlsx(data: dict[str, str], output_path: str):
    """Write a cell map to an .xlsx file.

    Parameters
    ----------
    data : dict[str, str]
        Mapping of Excel cell references to values.
        Example: {'A1': '封装信息', 'A2': 'PinName1', 'B2': '1'}
    output_path : str
        Path for the output .xlsx file.
    """
    writer = XLSXWriter()
    writer.write(data, output_path)


class XLSXWriter:
    """Build an OOXML .xlsx file from a cell map."""

    def __init__(self):
        self._strings: list[str] = []
        self._string_index: dict[str, int] = {}

    def write(self, data: dict[str, str], output_path: str):
        """Write *data* to *output_path* as an .xlsx file."""
        # Collect all unique strings for the shared strings table
        for value in data.values():
            self._add_string(value)

        with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            zf.writestr('[Content_Types].xml', self._content_types_xml())
            zf.writestr('_rels/.rels', self._rels_xml())
            zf.writestr('xl/workbook.xml', self._workbook_xml())
            zf.writestr('xl/_rels/workbook.xml.rels', self._workbook_rels_xml())
            zf.writestr('xl/sharedStrings.xml', self._shared_strings_xml())
            zf.writestr('xl/worksheets/sheet1.xml', self._sheet_xml(data))

    def _add_string(self, s: str) -> int:
        """Add a string to the SST and return its index."""
        if s in self._string_index:
            return self._string_index[s]
        idx = len(self._strings)
        self._strings.append(s)
        self._string_index[s] = idx
        return idx

    # ── XML templates ───────────────────────────────────────────────

    def _content_types_xml(self) -> str:
        return '''<?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/worksheets/sheet1.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>
  <Override PartName="/xl/sharedStrings.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sharedStrings+xml"/>
</Types>'''

    def _rels_xml(self) -> str:
        return '''<?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>'''

    def _workbook_xml(self) -> str:
        return '''<?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">
  <sheets>
    <sheet name="Sheet1" sheetId="1" r:id="rId1"/>
  </sheets>
</workbook>'''

    def _workbook_rels_xml(self) -> str:
        return '''<?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/worksheet" Target="worksheets/sheet1.xml"/>
  <Relationship Id="rId2" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/sharedStrings" Target="sharedStrings.xml"/>
</Relationships>'''

    def _shared_strings_xml(self) -> str:
        parts = ['<?xml version="1.0" encoding="UTF-8" standalone="yes"?>']
        parts.append(f'<sst xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" count="{len(self._strings)}" unique="{len(self._strings)}">')
        for s in self._strings:
            # Escape XML special characters
            escaped = s.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
            parts.append(f'  <si><t>{escaped}</t></si>')
        parts.append('</sst>')
        return '\n'.join(parts)

    def _sheet_xml(self, data: dict[str, str]) -> str:
        """Build sheet1.xml from the cell map.

        data keys are Excel cell references like 'A1', 'B2', etc.
        All values are treated as shared strings.
        """
        # Determine dimensions
        max_row = 0
        max_col = 0
        for ref in data:
            row, col = self._ref_to_rc(ref)
            max_row = max(max_row, row)
            max_col = max(max_col, col)

        parts = ['<?xml version="1.0" encoding="UTF-8" standalone="yes"?>']
        parts.append('<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">')
        parts.append(f'  <dimension ref="A1:{rc_to_cell_ref(max_row, max_col)}"/>')
        parts.append('  <sheetData>')

        # Group cells by row
        rows: dict[int, list[tuple[int, str]]] = {}
        for ref, value in data.items():
            row, col = self._ref_to_rc(ref)
            if row not in rows:
                rows[row] = []
            rows[row].append((col, value))

        for row_num in sorted(rows):
            parts.append(f'    <row r="{row_num + 1}">')
            for col, value in sorted(rows[row_num]):
                cell_ref = rc_to_cell_ref(row_num, col)
                si = self._add_string(value)
                parts.append(f'      <c r="{cell_ref}" t="s"><v>{si}</v></c>')
            parts.append('    </row>')

        parts.append('  </sheetData>')
        parts.append('</worksheet>')
        return '\n'.join(parts)

    @staticmethod
    def _ref_to_rc(ref: str) -> tuple[int, int]:
        """Convert cell reference to (row, col) 0-based."""
        col_letters = []
        row_digits = []
        for ch in ref:
            if ch.isalpha():
                col_letters.append(ch)
            else:
                row_digits.append(ch)
        col = 0
        for ch in ''.join(col_letters).upper():
            col = col * 26 + (ord(ch) - ord('A') + 1)
        col -= 1
        row = int(''.join(row_digits)) - 1
        return row, col