Skip to content

perf(xlsx): speed up simple table conversion #2546

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Dex Hunter (dexhunter) wants to merge 1 commit into
base: main
Choose a base branch
from
+322 −10
Open

perf(xlsx): speed up simple table conversion #2546

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
199 changes: 189 additions & 10 deletions packages/markitdown/src/markitdown/converters/_xlsx_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,26 +1,32 @@
import inspect
import io
import re
import sys
import zipfile
from collections.abc import Iterator
from contextlib import contextmanager
from typing import BinaryIO, Any, Iterator, Optional
from ._html_converter import HtmlConverter
from html import unescape
from html.parser import HTMLParser
from typing import Any, BinaryIO

from .._base_converter import DocumentConverter, DocumentConverterResult
from .._exceptions import MissingDependencyException, MISSING_DEPENDENCY_MESSAGE
from .._exceptions import MISSING_DEPENDENCY_MESSAGE, MissingDependencyException
from .._stream_info import StreamInfo
from ._html_converter import HtmlConverter
from ._markdownify import _CustomMarkdownify

# Try loading optional (but in this case, required) dependencies
# Save reporting of any exceptions for later
_xlsx_dependency_exc_info = None
try:
import pandas as pd
import openpyxl # noqa: F401
import pandas as pd
except ImportError:
_xlsx_dependency_exc_info = sys.exc_info()

_xls_dependency_exc_info = None
try:
import pandas as pd # noqa: F811
import pandas as pd
import xlrd # noqa: F401
except ImportError:
_xls_dependency_exc_info = sys.exc_info()
Expand All @@ -42,6 +48,88 @@
# workbook is repaired by renaming it. The rename is confined to <sheetView> start tags.
_SHEET_VIEW_START_TAG = re.compile(rb"<sheetView(?=[\s/>])[^>]*>")
_SHOW_ZEROES_ATTRIBUTE = re.compile(rb"(?<=[\s])showZeroes(\s*=)")
_NON_ASCII_SPACE_WHITESPACE = re.compile(r"[^\S ]")

_PANDAS_TABLE_RE = re.compile(
r'\A<table border="1" class="dataframe">\s*<thead>\s*'
r"(?P<header><tr[^>]*>.*?</tr>)\s*</thead>\s*<tbody>\s*"
r"(?P<body>.*?)</tbody>\s*</table>\s*\Z",
re.DOTALL,
)


def _scan_pandas_rows(
section: str,
expected_tag: str,
formatter: _CustomMarkdownify,
expected_width: int | None = None,
) -> tuple[list[str], int] | None:
lines: list[str] = []
position = 0
width = expected_width
escape_has_parent_tags = (
"parent_tags" in inspect.signature(formatter.escape).parameters
)
while position < len(section):
if not section.startswith("<tr", position):
return None
tag_end = section.find(">", position + 3)
row_end = section.find("</tr>", position)
if tag_end < 0 or row_end < 0 or tag_end >= row_end:
return None
row_text = section[tag_end + 1 : row_end]
open_tag = f"<{expected_tag}>"
close_tag = f"</{expected_tag}>"
parts = row_text.split(open_tag)
if len(parts) < 2 or parts[0].strip():
return None
cells: list[str] = []
for part in parts[1:]:
raw, closing, trailing = part.partition(close_tag)
if not closing or "<" in raw or trailing.strip():
return None
if "&" in raw:
raw = unescape(raw)
if _NON_ASCII_SPACE_WHITESPACE.search(raw) is not None:
return None
normalized = " ".join(raw.split()) if " " in raw else raw
if escape_has_parent_tags:
cells.append(formatter.escape(normalized, []))
else:
cells.append(formatter.escape(normalized))
if not cells or (width is not None and len(cells) != width):
return None
width = len(cells)
lines.append("| " + " | ".join(cells) + " |")
position = row_end + len("</tr>")
if position < len(section):
next_tag = section.find("<", position)
if next_tag < 0:
if section[position:].strip():
return None
position = len(section)
elif section[position:next_tag].strip():
return None
else:
position = next_tag
return lines, width if width is not None else 0


def _pandas_table_to_markdown_scanner(html_content: str) -> str | None:
table_match = _PANDAS_TABLE_RE.fullmatch(html_content.lstrip())
if table_match is None:
return None
formatter = _CustomMarkdownify()
header = _scan_pandas_rows(table_match.group("header"), "th", formatter)
if header is None or len(header[0]) != 1 or header[1] == 0:
return None
body = _scan_pandas_rows(table_match.group("body"), "td", formatter, header[1])
if body is None:
return None
lines = [header[0][0]]
lines.append("| " + " | ".join("---" for _ in range(header[1])) + " |")
lines.extend(body[0])
return "\n".join(lines)


@contextmanager
Expand Down Expand Up @@ -93,6 +181,93 @@ def _repair_sheetview_show_zeroes(
return repaired_stream


class _PandasTableParser(HTMLParser):
"""Parse the restricted table HTML emitted by pandas.DataFrame.to_html."""

def __init__(self) -> None:
super().__init__(convert_charrefs=True)
self.rows: list[list[str]] = []
self._formatter = _CustomMarkdownify()
self._escape_has_parent_tags = (
"parent_tags" in inspect.signature(self._formatter.escape).parameters
)
self._row: list[str] | None = None
self._cell_text: list[str] | None = None
self.valid = True

def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
del attrs
if tag == "tr":
if self._row is not None or self._cell_text is not None:
self.valid = False
self._row = []
elif tag in {"th", "td"}:
if self._row is None or self._cell_text is not None:
self.valid = False
self._cell_text = []
elif tag not in {"table", "thead", "tbody"}:
self.valid = False

def handle_data(self, data: str) -> None:
if self._cell_text is not None:
self._cell_text.append(data)
elif data.strip():
self.valid = False

def handle_endtag(self, tag: str) -> None:
if tag in {"th", "td"}:
if self._row is None or self._cell_text is None:
self.valid = False
return
raw = "".join(self._cell_text)
if _NON_ASCII_SPACE_WHITESPACE.search(raw) is not None:
self.valid = False
normalized = " ".join(raw.split())
if self._escape_has_parent_tags:
self._row.append(self._formatter.escape(normalized, []))
else:
self._row.append(self._formatter.escape(normalized))
self._cell_text = None
elif tag == "tr":
if self._row is None or self._cell_text is not None:
self.valid = False
return
self.rows.append(self._row)
self._row = None
elif tag not in {"table", "thead", "tbody"}:
self.valid = False


def _pandas_table_to_markdown_fastpath(html_content: str) -> str | None:
"""Convert simple pandas tables without traversing a general HTML DOM."""
if not html_content.lstrip().startswith('<table border="1" class="dataframe">'):
return None

scanned = _pandas_table_to_markdown_scanner(html_content)
if scanned is not None:
return scanned

parser = _PandasTableParser()
parser.feed(html_content)
parser.close()
if (
not parser.valid
or parser._row is not None
or parser._cell_text is not None
or not parser.rows
):
return None

width = len(parser.rows[0])
if width == 0 or any(len(row) != width for row in parser.rows):
return None

lines = ["| " + " | ".join(parser.rows[0]) + " |"]
lines.append("| " + " | ".join("---" for _ in range(width)) + " |")
lines.extend("| " + " | ".join(row) + " |" for row in parser.rows[1:])
return "\n".join(lines)


class XlsxConverter(DocumentConverter):
"""
Converts XLSX files to Markdown, with each sheet presented as a separate Markdown table.
Expand Down Expand Up @@ -151,12 +326,16 @@ def convert(
for s in sheets:
md_content += f"## {s}\n"
html_content = sheets[s].to_html(index=False)
md_content += (
self._html_converter.convert_string(
framework_keys = {"_parent_converters", "file_extension", "url"}
formatting_kwargs = {key for key in kwargs if key not in framework_keys}
table_markdown = None
if not formatting_kwargs and images is None:
table_markdown = _pandas_table_to_markdown_fastpath(html_content)
if table_markdown is None:
table_markdown = self._html_converter.convert_string(
html_content, **kwargs
).markdown.strip()
+ "\n\n"
)
md_content += table_markdown + "\n\n"
if images is not None:
image_content = images.to_html(s, self._image_to_html, kwargs)
if image_content:
Expand All @@ -174,7 +353,7 @@ def _image_to_html(
image_stream: BinaryIO,
stream_info: StreamInfo,
**kwargs: Any,
) -> Optional[str]:
) -> str | None:
"""Override to render an embedded image as an HTML fragment.

The stream is borrowed, seekable, and positioned at zero; do not close
Expand Down
133 changes: 133 additions & 0 deletions packages/markitdown/tests/test_xlsx_table_conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
"""Keep the XLSX table shortcut aligned with the public conversion output."""

import inspect
import io
from unittest.mock import Mock

import openpyxl
import pandas as pd
import pytest
from markitdown import MarkItDown, StreamInfo
from markitdown.converters import XlsxConverter, _xlsx_converter
from markitdown.converters._html_converter import HtmlConverter
from markitdown.converters._markdownify import _CustomMarkdownify


def _workbook(*sheets: tuple[str, list[list[object]]]) -> bytes:
workbook = openpyxl.Workbook()
workbook.remove(workbook.active)
for name, rows in sheets:
sheet = workbook.create_sheet(name)
for row in rows:
sheet.append(row)
stream = io.BytesIO()
workbook.save(stream)
workbook.close()
return stream.getvalue()


def _legacy_output(data: bytes, **kwargs: object) -> str:
sheets = pd.read_excel(io.BytesIO(data), sheet_name=None, engine="openpyxl")
converter = HtmlConverter()
return "\n\n".join(
f"## {name}\n"
+ converter.convert_string(
sheet.to_html(index=False), **kwargs
).markdown.strip()
for name, sheet in sheets.items()
)


def test_public_xlsx_conversion_keeps_table_output_without_html_fallback(
monkeypatch: pytest.MonkeyPatch,
) -> None:
data = _workbook(
("People", [["Name", "Count"], ["Ada", 2], ["Björk & Co", 3]]),
("Notes", [["Text"], ["<tag> *bold* | pipe"]]),
)
expected = _legacy_output(data)
assert "## People\n| Name | Count |\n| --- | --- |\n| Ada | 2 |" in expected
assert "Björk & Co" in expected
assert "\\*bold\\*" in expected

def unexpected_fallback(*args: object, **kwargs: object) -> None:
raise AssertionError("ordinary pandas tables should use the table shortcut")

monkeypatch.setattr(HtmlConverter, "convert_string", unexpected_fallback)
actual = (
MarkItDown().convert_stream(io.BytesIO(data), file_extension=".xlsx").markdown
)
assert actual == expected


def test_unicode_whitespace_uses_legacy_html_conversion(
monkeypatch: pytest.MonkeyPatch,
) -> None:
data = _workbook(("Unicode", [["Text"], ["one\u00a0two"]]))
expected = _legacy_output(data)
converter = XlsxConverter()
original = converter._html_converter.convert_string
fallback = Mock(wraps=original)
monkeypatch.setattr(converter._html_converter, "convert_string", fallback)

actual = converter.convert(io.BytesIO(data), StreamInfo(extension=".xlsx")).markdown

assert actual == expected
assert "one" in actual and "two" in actual
fallback.assert_called_once()


def test_formatting_options_still_reach_legacy_html_converter(
monkeypatch: pytest.MonkeyPatch,
) -> None:
data = _workbook(("Options", [["first_name"], ["a_b"]]))
options = {"escape_underscores": False}
expected = _legacy_output(data, **options)
converter = XlsxConverter()
original = converter._html_converter.convert_string
fallback = Mock(wraps=original)
monkeypatch.setattr(converter._html_converter, "convert_string", fallback)

actual = converter.convert(
io.BytesIO(data), StreamInfo(extension=".xlsx"), **options
).markdown

assert actual == expected
assert "a_b" in actual
fallback.assert_called_once()
assert fallback.call_args.kwargs == options


def test_empty_and_header_only_sheets_keep_native_output() -> None:
data = _workbook(
("Empty", []),
("Headers", [["First", "Second"]]),
)
actual = (
MarkItDown().convert_stream(io.BytesIO(data), file_extension=".xlsx").markdown
)
assert actual == _legacy_output(data)
assert actual.startswith("## Empty\n")
assert "## Headers\n| First | Second |" in actual


@pytest.mark.parametrize("value", ["a_b *c*", "one\u00a0two"])
def test_older_markdownify_escape_signature_keeps_public_xlsx_output(
monkeypatch: pytest.MonkeyPatch, value: str
) -> None:
data = _workbook(("Legacy", [["Text"], [value]]))
expected = _legacy_output(data)
original_escape = _CustomMarkdownify.escape
has_parent_tags = "parent_tags" in inspect.signature(original_escape).parameters

class LegacyEscapeFormatter(_CustomMarkdownify):
def escape(self, text: str) -> str:
if has_parent_tags:
return original_escape(self, text, [])
return original_escape(self, text)

monkeypatch.setattr(_xlsx_converter, "_CustomMarkdownify", LegacyEscapeFormatter)
actual = (
MarkItDown().convert_stream(io.BytesIO(data), file_extension=".xlsx").markdown
)
assert actual == expected