Skip to content

Commit ccf78f0

Browse files
finish parse merged cells
1 parent e57e332 commit ccf78f0

4 files changed

Lines changed: 140 additions & 2 deletions

File tree

packages/markitdown/src/markitdown/converters/_xlsx_converter.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,10 @@ def convert(
8181
)
8282

8383
sheets = pd.read_excel(file_stream, sheet_name=None, engine="openpyxl")
84+
if kwargs.get("fill_merged_cells", False):
85+
md_content = self._parse_merged_cells(file_stream, sheets, **kwargs)
86+
return DocumentConverterResult(markdown=md_content.strip())
87+
8488
md_content = ""
8589
for s in sheets:
8690
md_content += f"## {s}\n"
@@ -94,6 +98,63 @@ def convert(
9498

9599
return DocumentConverterResult(markdown=md_content.strip())
96100

101+
def _parse_merged_cells(
102+
self, file_stream: BinaryIO, sheets: dict[str, pd.DataFrame], **kwargs: Any
103+
) -> str:
104+
"""Use openpyxl to parse merged cells
105+
106+
Args:
107+
file_stream: BinaryIO
108+
Returns:
109+
str
110+
"""
111+
wb = openpyxl.load_workbook(file_stream)
112+
113+
merged_cells_info = {}
114+
for sheet in wb.worksheets:
115+
merged_cells = {}
116+
for row in sheet.merged_cells.ranges:
117+
min_col, min_row, max_col, max_row = row.bounds
118+
common_value = sheet.cell(row=min_row, column=min_col).value
119+
for row in range(min_row, max_row + 1):
120+
for col in range(min_col, max_col + 1):
121+
merged_cells[(row, col)] = common_value
122+
# Merged header at first (sort by row)
123+
merged_cells = dict(
124+
sorted(merged_cells.items(), key=lambda x: (x[0][0], x[0][1]))
125+
)
126+
merged_cells_info[sheet.title] = merged_cells
127+
128+
wb.close()
129+
130+
md_content = ""
131+
for s in sheets:
132+
md_content += f"## {s}\n"
133+
df = sheets[s]
134+
for (row, col), value in merged_cells_info[s].items():
135+
if row == 1:
136+
# Header row merged.
137+
if col > len(df.columns):
138+
# Insert new column
139+
df.insert(col - 1, f"{value} {col-1}", "NaN")
140+
elif str(df.columns[col - 1]).startswith("Unnamed"):
141+
# Rename unnamed column
142+
df.rename(
143+
columns={df.columns[col - 1]: f"{value} {col-1}"},
144+
inplace=True,
145+
)
146+
else:
147+
df.at[row - 2, df.columns[col - 1]] = value
148+
html_content = df.to_html(index=False)
149+
md_content += (
150+
self._html_converter.convert_string(
151+
html_content, **kwargs
152+
).markdown.strip()
153+
+ "\n\n"
154+
)
155+
156+
return md_content
157+
97158

98159
class XlsConverter(DocumentConverter):
99160
"""

packages/markitdown/tests/_test_vectors.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,3 +277,22 @@ class FileTestVector(object):
277277
],
278278
),
279279
]
280+
281+
282+
MERGED_CELLS_TEST_VECTORS = [
283+
FileTestVector(
284+
filename="test.xlsx",
285+
mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
286+
charset=None,
287+
url=None,
288+
must_include=[
289+
"722.0 | NaN | NaN",
290+
"NaN | 42.000000 | NaN",
291+
"Merged Column | Merged Column 6 | Merged Column 2 | Merged Column 2 8",
292+
"## 09060124-b5e7-4717-9d07-3c046eb",
293+
"6ff4173b-42a5-4784-9b19-f49caff4d93d",
294+
"affc7dad-52dc-4b98-9b5d-51e65d8a8ad0",
295+
],
296+
must_not_include=["Unnamed"],
297+
),
298+
]
855 Bytes
Binary file not shown.

packages/markitdown/tests/test_module_vectors.py

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,17 @@
99
from pathlib import Path
1010

1111
if __name__ == "__main__":
12-
from _test_vectors import GENERAL_TEST_VECTORS, DATA_URI_TEST_VECTORS
12+
from _test_vectors import (
13+
GENERAL_TEST_VECTORS,
14+
DATA_URI_TEST_VECTORS,
15+
MERGED_CELLS_TEST_VECTORS,
16+
)
1317
else:
14-
from ._test_vectors import GENERAL_TEST_VECTORS, DATA_URI_TEST_VECTORS
18+
from ._test_vectors import (
19+
GENERAL_TEST_VECTORS,
20+
DATA_URI_TEST_VECTORS,
21+
MERGED_CELLS_TEST_VECTORS,
22+
)
1523

1624
from markitdown import (
1725
MarkItDown,
@@ -200,6 +208,44 @@ def test_convert_stream_keep_data_uris(test_vector):
200208
for string in test_vector.must_not_include:
201209
assert string not in result.markdown
202210

211+
@pytest.mark.parametrize("test_vector", MERGED_CELLS_TEST_VECTORS)
212+
def test_convert_xlsx(test_vector):
213+
"""Test the conversion of an XLSX file."""
214+
markitdown = MarkItDown()
215+
216+
result = markitdown.convert(
217+
os.path.join(TEST_FILES_DIR, test_vector.filename),
218+
fill_merged_cells=True,
219+
url=test_vector.url,
220+
)
221+
222+
for string in test_vector.must_include:
223+
assert string in result.markdown
224+
for string in test_vector.must_not_include:
225+
assert string not in result.markdown
226+
227+
228+
@pytest.mark.parametrize("test_vector", MERGED_CELLS_TEST_VECTORS)
229+
def test_convert_stream_xlsx(test_vector):
230+
"""Test the conversion of an XLSX file."""
231+
markitdown = MarkItDown()
232+
233+
stream_info = StreamInfo(
234+
extension=os.path.splitext(test_vector.filename)[1],
235+
mimetype=test_vector.mimetype,
236+
charset=test_vector.charset,
237+
)
238+
239+
with open(os.path.join(TEST_FILES_DIR, test_vector.filename), "rb") as stream:
240+
result = markitdown.convert(
241+
stream, stream_info=stream_info, fill_merged_cells=True, url=test_vector.url
242+
)
243+
244+
for string in test_vector.must_include:
245+
assert string in result.markdown
246+
for string in test_vector.must_not_include:
247+
assert string not in result.markdown
248+
203249

204250
def test_convert_docx_with_style_missing_type(tmp_path):
205251
"""DOCX conversion should not fail when a style entry is missing w:type."""
@@ -264,4 +310,16 @@ def test_convert_docx_with_style_missing_type(tmp_path):
264310
test_function(test_vector)
265311
print("OK")
266312

313+
# XLSX parse merged cells tests
314+
for test_function in [
315+
test_convert_xlsx,
316+
test_convert_stream_xlsx,
317+
]:
318+
for test_vector in MERGED_CELLS_TEST_VECTORS:
319+
print(
320+
f"Running {test_function.__name__} on {test_vector.filename}...", end=""
321+
)
322+
test_function(test_vector)
323+
print("OK")
324+
267325
print("All tests passed!")

0 commit comments

Comments
 (0)