feat: add Excel export functionality with support for metadata and customizable sheets
This commit is contained in:
@@ -1,13 +1,14 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import csv
|
import csv
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass, field
|
||||||
from datetime import date, datetime, timezone
|
from datetime import date, datetime, timezone
|
||||||
from decimal import Decimal, InvalidOperation, ROUND_HALF_UP
|
from decimal import Decimal, InvalidOperation, ROUND_HALF_UP
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from io import StringIO
|
from io import BytesIO, StringIO
|
||||||
from typing import Any, Callable, Iterable, Iterator, Sequence
|
from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence
|
||||||
|
|
||||||
|
from openpyxl import Workbook
|
||||||
CSVValueFormatter = Callable[[Any], str]
|
CSVValueFormatter = Callable[[Any], str]
|
||||||
Accessor = Callable[[Any], Any]
|
Accessor = Callable[[Any], Any]
|
||||||
|
|
||||||
@@ -18,6 +19,9 @@ __all__ = [
|
|||||||
"default_scenario_columns",
|
"default_scenario_columns",
|
||||||
"stream_projects_to_csv",
|
"stream_projects_to_csv",
|
||||||
"stream_scenarios_to_csv",
|
"stream_scenarios_to_csv",
|
||||||
|
"ExcelExporter",
|
||||||
|
"export_projects_to_excel",
|
||||||
|
"export_scenarios_to_excel",
|
||||||
"default_formatter",
|
"default_formatter",
|
||||||
"format_datetime_utc",
|
"format_datetime_utc",
|
||||||
"format_date_iso",
|
"format_date_iso",
|
||||||
@@ -34,13 +38,13 @@ class CSVExportColumn:
|
|||||||
formatter: CSVValueFormatter | None = None
|
formatter: CSVValueFormatter | None = None
|
||||||
required: bool = False
|
required: bool = False
|
||||||
|
|
||||||
def resolve_accessor(self) -> Accessor:
|
_accessor: Accessor = field(init=False, repr=False)
|
||||||
if isinstance(self.accessor, str):
|
|
||||||
return _coerce_accessor(self.accessor)
|
def __post_init__(self) -> None:
|
||||||
return self.accessor
|
object.__setattr__(self, "_accessor", _coerce_accessor(self.accessor))
|
||||||
|
|
||||||
def value_for(self, entity: Any) -> Any:
|
def value_for(self, entity: Any) -> Any:
|
||||||
accessor = self.resolve_accessor()
|
accessor = object.__getattribute__(self, "_accessor")
|
||||||
try:
|
try:
|
||||||
return accessor(entity)
|
return accessor(entity)
|
||||||
except Exception: # pragma: no cover - defensive safeguard
|
except Exception: # pragma: no cover - defensive safeguard
|
||||||
@@ -85,8 +89,7 @@ class CSVExporter:
|
|||||||
def _format_row(self, record: Any) -> list[str]:
|
def _format_row(self, record: Any) -> list[str]:
|
||||||
formatted: list[str] = []
|
formatted: list[str] = []
|
||||||
for column in self._columns:
|
for column in self._columns:
|
||||||
accessor = column.resolve_accessor()
|
raw_value = column.value_for(record)
|
||||||
raw_value = accessor(record)
|
|
||||||
formatter = column.formatter or default_formatter
|
formatter = column.formatter or default_formatter
|
||||||
formatted.append(formatter(raw_value))
|
formatted.append(formatter(raw_value))
|
||||||
return formatted
|
return formatted
|
||||||
@@ -216,6 +219,114 @@ def format_decimal(value: Any) -> str:
|
|||||||
return default_formatter(value)
|
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:
|
def _coerce_accessor(accessor: Accessor | str) -> Accessor:
|
||||||
if callable(accessor):
|
if callable(accessor):
|
||||||
return accessor
|
return accessor
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from datetime import date, datetime, timezone
|
from datetime import date, datetime, timezone
|
||||||
from decimal import Decimal
|
from decimal import Decimal
|
||||||
|
from io import BytesIO
|
||||||
from typing import Any, Iterable
|
from typing import Any, Iterable
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -10,15 +11,19 @@ import pytest
|
|||||||
from services.export_serializers import (
|
from services.export_serializers import (
|
||||||
CSVExportColumn,
|
CSVExportColumn,
|
||||||
CSVExporter,
|
CSVExporter,
|
||||||
|
ExcelExporter,
|
||||||
default_formatter,
|
default_formatter,
|
||||||
default_project_columns,
|
default_project_columns,
|
||||||
default_scenario_columns,
|
default_scenario_columns,
|
||||||
|
export_projects_to_excel,
|
||||||
|
export_scenarios_to_excel,
|
||||||
format_date_iso,
|
format_date_iso,
|
||||||
format_datetime_utc,
|
format_datetime_utc,
|
||||||
format_decimal,
|
format_decimal,
|
||||||
stream_projects_to_csv,
|
stream_projects_to_csv,
|
||||||
stream_scenarios_to_csv,
|
stream_scenarios_to_csv,
|
||||||
)
|
)
|
||||||
|
from openpyxl import load_workbook
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@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]
|
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:
|
def test_csv_exporter_writes_header_and_rows() -> None:
|
||||||
exporter = CSVExporter(
|
exporter = CSVExporter(
|
||||||
[
|
[
|
||||||
@@ -66,6 +76,64 @@ def test_csv_exporter_writes_header_and_rows() -> None:
|
|||||||
assert chunks[1] == "Alpha,Nevada\n"
|
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:
|
def test_csv_exporter_handles_optional_values_and_default_formatter() -> None:
|
||||||
exporter = CSVExporter(
|
exporter = CSVExporter(
|
||||||
[
|
[
|
||||||
|
|||||||
Reference in New Issue
Block a user