From 4b33a5dba3c86938823877d15b7f202322434422 Mon Sep 17 00:00:00 2001 From: zwitschi Date: Mon, 10 Nov 2025 18:32:09 +0100 Subject: [PATCH] feat: add Excel export functionality with support for metadata and customizable sheets --- services/export_serializers.py | 131 ++++++++++++++++++++++++++++--- tests/test_export_serializers.py | 68 ++++++++++++++++ 2 files changed, 189 insertions(+), 10 deletions(-) diff --git a/services/export_serializers.py b/services/export_serializers.py index 8e05da2..5242a15 100644 --- a/services/export_serializers.py +++ b/services/export_serializers.py @@ -1,13 +1,14 @@ from __future__ import annotations import csv -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import date, datetime, timezone from decimal import Decimal, InvalidOperation, ROUND_HALF_UP from enum import Enum -from io import StringIO -from typing import Any, Callable, Iterable, Iterator, Sequence +from io import BytesIO, StringIO +from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence +from openpyxl import Workbook CSVValueFormatter = Callable[[Any], str] Accessor = Callable[[Any], Any] @@ -18,6 +19,9 @@ __all__ = [ "default_scenario_columns", "stream_projects_to_csv", "stream_scenarios_to_csv", + "ExcelExporter", + "export_projects_to_excel", + "export_scenarios_to_excel", "default_formatter", "format_datetime_utc", "format_date_iso", @@ -34,13 +38,13 @@ class CSVExportColumn: formatter: CSVValueFormatter | None = None required: bool = False - def resolve_accessor(self) -> Accessor: - if isinstance(self.accessor, str): - return _coerce_accessor(self.accessor) - return self.accessor + _accessor: Accessor = field(init=False, repr=False) + + def __post_init__(self) -> None: + object.__setattr__(self, "_accessor", _coerce_accessor(self.accessor)) def value_for(self, entity: Any) -> Any: - accessor = self.resolve_accessor() + accessor = object.__getattribute__(self, "_accessor") try: return accessor(entity) except Exception: # pragma: no cover - defensive safeguard @@ -85,8 +89,7 @@ class CSVExporter: def _format_row(self, record: Any) -> list[str]: formatted: list[str] = [] for column in self._columns: - accessor = column.resolve_accessor() - raw_value = accessor(record) + raw_value = column.value_for(record) formatter = column.formatter or default_formatter formatted.append(formatter(raw_value)) return formatted @@ -216,6 +219,114 @@ def format_decimal(value: Any) -> str: return default_formatter(value) +class ExcelExporter: + """Produce Excel workbooks via write-only streaming.""" + + def __init__( + self, + columns: Sequence[CSVExportColumn], + *, + sheet_name: str = "Export", + workbook_title: str | None = None, + include_header: bool = True, + metadata: Mapping[str, Any] | None = None, + metadata_sheet_name: str = "Metadata", + ) -> None: + if not columns: + raise ValueError( + "At least one column is required for Excel export.") + self._columns: tuple[CSVExportColumn, ...] = tuple(columns) + self._sheet_name = sheet_name or "Export" + self._include_header = include_header + self._metadata = dict(metadata) if metadata else None + self._metadata_sheet_name = metadata_sheet_name or "Metadata" + self._workbook = Workbook(write_only=True) + if workbook_title: + self._workbook.properties.title = workbook_title + + def export(self, records: Iterable[Any]) -> bytes: + sheet = self._workbook.create_sheet(title=self._sheet_name) + if self._include_header: + sheet.append([column.header for column in self._columns]) + + for record in records: + sheet.append(self._format_row(record)) + + self._append_metadata_sheet() + return self._finalize() + + def _format_row(self, record: Any) -> list[Any]: + row: list[Any] = [] + for column in self._columns: + raw_value = column.value_for(record) + formatter = column.formatter or default_formatter + row.append(formatter(raw_value)) + return row + + def _append_metadata_sheet(self) -> None: + if not self._metadata: + return + + sheet_name = self._metadata_sheet_name + existing = set(self._workbook.sheetnames) + if sheet_name in existing: + index = 1 + while True: + candidate = f"{sheet_name}_{index}" + if candidate not in existing: + sheet_name = candidate + break + index += 1 + + meta_ws = self._workbook.create_sheet(title=sheet_name) + meta_ws.append(["Key", "Value"]) + for key, value in self._metadata.items(): + meta_ws.append([ + str(key), + "" if value is None else str(value), + ]) + + def _finalize(self) -> bytes: + buffer = BytesIO() + self._workbook.save(buffer) + buffer.seek(0) + return buffer.getvalue() + + +def export_projects_to_excel( + projects: Iterable[Any], + *, + columns: Sequence[CSVExportColumn] | None = None, + sheet_name: str = "Projects", + workbook_title: str | None = None, + metadata: Mapping[str, Any] | None = None, +) -> bytes: + exporter = ExcelExporter( + columns or default_project_columns(), + sheet_name=sheet_name, + workbook_title=workbook_title, + metadata=metadata, + ) + return exporter.export(projects) + + +def export_scenarios_to_excel( + scenarios: Iterable[Any], + *, + columns: Sequence[CSVExportColumn] | None = None, + sheet_name: str = "Scenarios", + workbook_title: str | None = None, + metadata: Mapping[str, Any] | None = None, +) -> bytes: + exporter = ExcelExporter( + columns or default_scenario_columns(), + sheet_name=sheet_name, + workbook_title=workbook_title, + metadata=metadata, + ) + return exporter.export(scenarios) + + def _coerce_accessor(accessor: Accessor | str) -> Accessor: if callable(accessor): return accessor diff --git a/tests/test_export_serializers.py b/tests/test_export_serializers.py index 7529d5d..b08413b 100644 --- a/tests/test_export_serializers.py +++ b/tests/test_export_serializers.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import date, datetime, timezone from decimal import Decimal +from io import BytesIO from typing import Any, Iterable import pytest @@ -10,15 +11,19 @@ import pytest from services.export_serializers import ( CSVExportColumn, CSVExporter, + ExcelExporter, default_formatter, default_project_columns, default_scenario_columns, + export_projects_to_excel, + export_scenarios_to_excel, format_date_iso, format_datetime_utc, format_decimal, stream_projects_to_csv, stream_scenarios_to_csv, ) +from openpyxl import load_workbook @dataclass(slots=True) @@ -50,6 +55,11 @@ def collect_csv_bytes(chunks: Iterable[bytes]) -> list[str]: return [chunk.decode("utf-8") for chunk in chunks] +def load_workbook_bytes(data: bytes): + buffer = BytesIO(data) + return load_workbook(buffer, read_only=True, data_only=True) + + def test_csv_exporter_writes_header_and_rows() -> None: exporter = CSVExporter( [ @@ -66,6 +76,64 @@ def test_csv_exporter_writes_header_and_rows() -> None: assert chunks[1] == "Alpha,Nevada\n" +def test_excel_exporter_basic_workbook() -> None: + exporter = ExcelExporter(default_project_columns(), sheet_name="Projects") + project = DummyProject(name="Alpha", location="Nevada") + + data = exporter.export([project]) + workbook = load_workbook_bytes(data) + sheet = workbook["Projects"] + rows = list(sheet.rows) + + assert [cell.value for cell in rows[0]] == [ + "name", + "location", + "operation_type", + "description", + "created_at", + "updated_at", + ] + assert rows[1][0].value == "Alpha" + + +def test_excel_export_projects_helper_with_metadata() -> None: + project = DummyProject(name="Alpha", location="Nevada") + data = export_projects_to_excel( + [project], metadata={"rows": 1}, workbook_title="Project Export") + + workbook = load_workbook_bytes(data) + assert workbook.properties.title == "Project Export" + assert "Projects" in workbook.sheetnames + assert any(sheet.title.startswith("Metadata") + for sheet in workbook.worksheets) + + +def test_excel_export_scenarios_helper_projects_resolved() -> None: + project = DummyProject(name="Alpha") + scenario = DummyScenario(project=project, name="Scenario 1") + data = export_scenarios_to_excel([scenario]) + + workbook = load_workbook_bytes(data) + sheet = workbook["Scenarios"] + rows = list(sheet.rows) + + assert rows[1][0].value == "Alpha" + assert rows[1][1].value == "Scenario 1" + exporter = CSVExporter( + [ + CSVExportColumn("Name", "name"), + CSVExportColumn("Location", "location"), + ] + ) + + project = DummyProject(name="Alpha", location="Nevada") + + chunks = collect_csv_bytes(exporter.iter_bytes([project])) + + assert chunks[0] == "Name,Location\n" + assert chunks[1] == "Alpha,Nevada\n" + + def test_csv_exporter_handles_optional_values_and_default_formatter() -> None: exporter = CSVExporter( [