diff --git a/docling/backend/msexcel_backend.py b/docling/backend/msexcel_backend.py index 74838a5972..766f3355c2 100644 --- a/docling/backend/msexcel_backend.py +++ b/docling/backend/msexcel_backend.py @@ -79,7 +79,7 @@ _OPENPYXL_AVAILABLE: bool = False _OPENPYXL_IMPORT_ERROR: ImportError | None = None try: # pragma: no cover - import-time guard - from openpyxl import Workbook, load_workbook + from openpyxl import Workbook from openpyxl.cell.cell import Cell, MergedCell from openpyxl.chart import DoughnutChart, PieChart, PieChart3D, ProjectedPieChart from openpyxl.chart.data_source import AxDataSource, NumDataSource @@ -93,8 +93,10 @@ TwoCellAnchor, ) from openpyxl.packaging.relationship import get_dependents, get_rels_path + from openpyxl.reader.excel import ExcelReader from openpyxl.styles import PatternFill from openpyxl.utils.cell import range_boundaries + from openpyxl.worksheet.merge import MergedCellRange from openpyxl.worksheet.worksheet import Worksheet from openpyxl.xml.constants import IMAGE_NS @@ -275,6 +277,236 @@ def span_at(self, row: int, col: int) -> tuple[int, int]: return self._anchor_spans.get((row, col), (1, 1)) +# While openpyxl loads a worksheet, it creates one ``MergedCell`` object for +# each covered cell of a merged range, and one cell with its own ``Hyperlink`` +# copy for each cell of a hyperlink range. These limit the cells that the ranges +# of one workbook can make it create. A hyperlink also sets its target as the +# value of each empty cell, so these cells become table data: per cell, they +# cost about five times more than merged cells, and their limit is smaller. +_MAX_MERGED_CELLS: Final[int] = 100_000 +_MAX_HYPERLINK_CELLS: Final[int] = 10_000 + + +class _RangeBudget: + """Fit cell ranges of one workbook into a budget of created cells. + + Ranges are kept unchanged while the workbook stays within the budget. A + range that does not fit is clipped to the cells with values in its + worksheet, and to its top-left cell if it still does not fit. + """ + + def __init__(self, cells: int) -> None: + self._remaining = cells + self.clipped = 0 + + def fit(self, region: DataRegion, bounds: DataRegion | None) -> DataRegion: + """Return the part of a range that fits into the remaining budget.""" + covered = _covered_cells(region) + if covered <= self._remaining: + self._remaining -= covered + return region + self.clipped += 1 + if bounds is not None: + region = DataRegion( + region.min_row, + max(region.min_row, min(region.max_row, bounds.max_row)), + region.min_col, + max(region.min_col, min(region.max_col, bounds.max_col)), + ) + covered = _covered_cells(region) + if covered > self._remaining: + region = DataRegion( + region.min_row, region.min_row, region.min_col, region.min_col + ) + covered = 0 + self._remaining -= covered + return region + + +def _covered_cells(region: DataRegion) -> int: + """Return the number of cells of a range after its top-left cell.""" + # An empty range, such as rows past the end of a sheet, covers no cells. + if region.width() <= 0 or region.height() <= 0: + return 0 + return region.width() * region.height() - 1 + + +class _SheetRangeClipper: + """Clip the merged and hyperlink ranges of one worksheet while it loads. + + openpyxl has no option to limit how it expands these ranges. The clipper + shadows three methods on the worksheet instance and removes them after the + load; no class or module attribute of openpyxl changes, so concurrent loads + are not affected. + + - ``_clean_merge_range`` creates the ``MergedCell`` objects of a merged range. + - ``iter_rows`` and ``iter_cols`` create the cells of a range lookup such as + ``ws["A1:J10"]``. While a worksheet loads, openpyxl looks up ranges only to + bind hyperlinks; comments use single-cell lookups. + """ + + def __init__( + self, + sheet: Worksheet, + merged_budget: _RangeBudget, + hyperlink_budget: _RangeBudget, + ) -> None: + self._sheet = sheet + self._merged_budget = merged_budget + self._hyperlink_budget = hyperlink_budget + self._value_bounds: list[DataRegion | None] = [] + self._clean_merge_range = sheet._clean_merge_range # type: ignore[attr-defined] + self._iter_rows = sheet.iter_rows + self._iter_cols = sheet.iter_cols + + def install(self) -> None: + self._sheet._clean_merge_range = self.clean_merge_range # type: ignore[attr-defined] + self._sheet.iter_rows = self.iter_rows # type: ignore[method-assign] + self._sheet.iter_cols = self.iter_cols # type: ignore[method-assign] + + def uninstall(self) -> None: + del self._sheet._clean_merge_range # type: ignore[attr-defined] + del self._sheet.iter_rows + del self._sheet.iter_cols + + def _bounds(self) -> DataRegion | None: + # openpyxl binds all cells of a sheet before its ranges, so the first + # call sees all values of the sheet. + if not self._value_bounds: + self._value_bounds.append(_value_bounds(self._sheet)) + return self._value_bounds[0] + + def clean_merge_range(self, merged_range: MergedCellRange) -> None: + region = self._merged_budget.fit( + DataRegion( + merged_range.min_row, + merged_range.max_row, + merged_range.min_col, + merged_range.max_col, + ), + self._bounds(), + ) + merged_range.max_row = region.max_row + merged_range.max_col = region.max_col + self._clean_merge_range(merged_range) + + def _fit( + self, + min_row: int | None, + max_row: int | None, + min_col: int | None, + max_col: int | None, + ) -> DataRegion: + # Same defaults as ``Worksheet.iter_rows`` and ``Worksheet.iter_cols``. + return self._hyperlink_budget.fit( + DataRegion( + min_row or 1, + max_row or self._sheet.max_row, + min_col or 1, + max_col or self._sheet.max_column, + ), + self._bounds(), + ) + + def iter_rows( + self, + min_row: int | None = None, + max_row: int | None = None, + min_col: int | None = None, + max_col: int | None = None, + values_only: bool = False, + ) -> Any: + if not any((min_row, max_row, min_col, max_col)): + return self._iter_rows(values_only=values_only) + region = self._fit(min_row, max_row, min_col, max_col) + return self._iter_rows( + region.min_row, region.max_row, region.min_col, region.max_col, values_only + ) + + def iter_cols( + self, + min_col: int | None = None, + max_col: int | None = None, + min_row: int | None = None, + max_row: int | None = None, + values_only: bool = False, + ) -> Any: + if not any((min_row, max_row, min_col, max_col)): + return self._iter_cols(values_only=values_only) + region = self._fit(min_row, max_row, min_col, max_col) + return self._iter_cols( + region.min_col, region.max_col, region.min_row, region.max_row, values_only + ) + + +def _value_bounds(sheet: Worksheet) -> DataRegion | None: + """Return the smallest region that holds all cells with a value.""" + coords = [coord for coord, cell in sheet._cells.items() if cell.value is not None] + if not coords: + return None + rows = [row for row, _ in coords] + cols = [col for _, col in coords] + return DataRegion(min(rows), max(rows), min(cols), max(cols)) + + +def _read_worksheets_with_clipped_ranges( + workbook: Workbook, read_worksheets: Callable[[], None] +) -> None: + """Run ``read_worksheets`` with ranges clipped on each new worksheet.""" + # The worksheet methods can only be shadowed once a worksheet exists, so + # ``create_sheet`` is shadowed on this workbook instance while its sheets load. + merged_budget = _RangeBudget(_MAX_MERGED_CELLS) + hyperlink_budget = _RangeBudget(_MAX_HYPERLINK_CELLS) + create_sheet = workbook.create_sheet + clippers: list[_SheetRangeClipper] = [] + + def create_bounded_sheet( + title: str | None = None, index: int | None = None + ) -> Worksheet: + sheet = create_sheet(title, index) + clipper = _SheetRangeClipper(sheet, merged_budget, hyperlink_budget) + clipper.install() + clippers.append(clipper) + return sheet + + workbook.create_sheet = create_bounded_sheet # type: ignore[method-assign] + try: + read_worksheets() + finally: + del workbook.create_sheet + for clipper in clippers: + clipper.uninstall() + if merged_budget.clipped: + _log.warning( + "Clipped %d merged cell ranges: together they cover more than %d cells", + merged_budget.clipped, + _MAX_MERGED_CELLS, + ) + if hyperlink_budget.clipped: + _log.warning( + "Clipped %d hyperlink ranges: together they cover more than %d cells", + hyperlink_budget.clipped, + _MAX_HYPERLINK_CELLS, + ) + + +def _load_workbook(filename: str | BytesIO) -> Workbook: + """Load a workbook like ``load_workbook(filename, data_only=True)``. + + Merged cell and hyperlink ranges are clipped so that they cover a bounded + number of cells. + """ + reader = ExcelReader(filename, data_only=True) + read_worksheets = reader.read_worksheets + # ``reader.wb`` exists once ``read`` has read the workbook part, which openpyxl + # does before it reads the worksheets. + reader.read_worksheets = lambda: _read_worksheets_with_clipped_ranges( # type: ignore[method-assign] + reader.wb, read_worksheets + ) + reader.read() + return reader.wb + + class ExcelCell(BaseModel): """Represents an Excel cell. @@ -398,13 +630,9 @@ def __init__( module=r"openpyxl\.reader\.drawings", ) if isinstance(self.path_or_stream, BytesIO): - self.workbook = load_workbook( - filename=self.path_or_stream, data_only=True - ) + self.workbook = _load_workbook(self.path_or_stream) elif isinstance(self.path_or_stream, Path): - self.workbook = load_workbook( - filename=str(self.path_or_stream), data_only=True - ) + self.workbook = _load_workbook(str(self.path_or_stream)) self.valid = self.workbook is not None except Exception as e: diff --git a/tests/test_backend_msexcel.py b/tests/test_backend_msexcel.py index 19d12271be..f35085da07 100644 --- a/tests/test_backend_msexcel.py +++ b/tests/test_backend_msexcel.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: MIT import logging +import tracemalloc +import zipfile from collections.abc import Iterator from datetime import datetime from io import BytesIO @@ -738,6 +740,145 @@ def test_merged_cells_are_indexed_once_and_preserve_semantics(tmp_path: Path) -> assert comment_map[(3, 3)] == [("Codex", "Synthetic note", None)] +def _rewrite_sheet_xml( + path: Path, old: str, new: str, sheet_part: str = "xl/worksheets/sheet1.xml" +) -> None: + """Rewrite a range in a sheet part, as openpyxl cannot save such ranges.""" + with zipfile.ZipFile(path) as source: + parts = {info.filename: source.read(info) for info in source.infolist()} + assert old.encode() in parts[sheet_part] + parts[sheet_part] = parts[sheet_part].replace(old.encode(), new.encode()) + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as target: + for name, data in parts.items(): + target.writestr(name, data) + + +def test_whole_sheet_merged_range_is_clipped_to_data(tmp_path: Path) -> None: + workbook = Workbook() + sheet = workbook.active + sheet["A1"] = "Quarterly report" + sheet["B3"] = "end" + sheet.merge_cells("A1:B2") + file_path = tmp_path / "whole-sheet-merge.xlsx" + workbook.save(file_path) + _rewrite_sheet_xml( + file_path, '', '' + ) + + converter = DocumentConverter(allowed_formats=[InputFormat.XLSX]) + tracemalloc.start() + try: + doc = converter.convert(file_path).document + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert peak < 50 * 1024 * 1024 + assert len(doc.tables) == 1 + table = doc.tables[0].data + assert (table.num_rows, table.num_cols) == (3, 2) + [anchor] = [cell for cell in table.table_cells if cell.text == "Quarterly report"] + assert (anchor.row_span, anchor.col_span) == (3, 2) + + +def test_merged_range_beyond_data_is_clipped(tmp_path: Path) -> None: + workbook = Workbook() + sheet = workbook.active + sheet.append(["Inventory", None, None]) + sheet.append(["Item", "Count", "Price"]) + sheet.append(["Fruit", 3, 1.5]) + sheet.append([None, 4, 2.5]) + sheet.append(["Total", None, None]) + sheet.merge_cells("A1:C1") + sheet.merge_cells("A3:A4") + sheet.merge_cells("A5:B5") + file_path = tmp_path / "partial-merge.xlsx" + workbook.save(file_path) + _rewrite_sheet_xml( + file_path, '', '' + ) + + in_doc = InputDocument( + path_or_stream=file_path, + format=InputFormat.XLSX, + filename=file_path.stem, + backend=MsExcelDocumentBackend, + ) + backend = MsExcelDocumentBackend(in_doc=in_doc, path_or_stream=file_path) + loaded_sheet = backend.workbook.active + assert {str(merged) for merged in loaded_sheet.merged_cells.ranges} == { + "A1:C1", + "A3:A4", + "A5:C5", + } + + tables, _ = backend._find_data_tables(loaded_sheet) + + assert len(tables) == 1 + table = tables[0] + assert (table.anchor, table.num_rows, table.num_cols) == ((0, 0), 5, 3) + spans = { + cell.text: (cell.row_span, cell.col_span) + for cell in table.data + if cell.text in {"Inventory", "Fruit", "Total"} + } + assert spans == {"Inventory": (1, 3), "Fruit": (2, 1), "Total": (1, 3)} + + +def test_hyperlink_range_is_clipped_to_data(tmp_path: Path) -> None: + workbook = Workbook() + sheet = workbook.active + sheet["A1"] = "Docs" + sheet["A1"].hyperlink = "https://example.com/docs" + sheet["B3"] = "end" + small = workbook.create_sheet("Small") + small["A1"] = "Home" + small["A1"].hyperlink = "https://example.com/" + small["C3"] = "end" + file_path = tmp_path / "hyperlink-range.xlsx" + workbook.save(file_path) + _rewrite_sheet_xml(file_path, 'ref="A1" r:id', 'ref="A1:J100000" r:id') + _rewrite_sheet_xml( + file_path, + 'ref="A1" r:id', + 'ref="A1:B2" r:id', + sheet_part="xl/worksheets/sheet2.xml", + ) + + in_doc = InputDocument( + path_or_stream=file_path, + format=InputFormat.XLSX, + filename=file_path.stem, + backend=MsExcelDocumentBackend, + ) + tracemalloc.start() + try: + backend = MsExcelDocumentBackend(in_doc=in_doc, path_or_stream=file_path) + doc = backend.convert() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert peak < 50 * 1024 * 1024 + loaded_sheet = backend.workbook["Sheet"] + assert loaded_sheet["A1"].hyperlink.target == "https://example.com/docs" + assert (loaded_sheet.max_row, loaded_sheet.max_column) == (3, 2) + loaded_small = backend.workbook["Small"] + linked = { + cell.coordinate + for row in loaded_small.iter_rows() + for cell in row + if cell.hyperlink is not None + } + assert linked == {"A1", "A2", "B1", "B2"} + [docs_table] = [ + table.data + for table in doc.tables + if any(cell.text == "Docs" for cell in table.data.table_cells) + ] + assert (docs_table.num_rows, docs_table.num_cols) == (3, 2) + + def test_split_leading_section_label_helper() -> None: backend = object.__new__(MsExcelDocumentBackend)