luciano/markitdown
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore: make cleaning optional

Browse source
This commit is contained in:
Hew Li Yang 2024-12-17 14:03:39 +08:00
parent 134d35a859
commit c2aae4ddda
2 changed files with 24 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
src/markitdown/_markitdown.py
View file

@ -527,7 +527,16 @@ class XlsxConverter(HtmlConverter):
return "" return ""
return colname return colname
def convert(self, local_path, **kwargs) -> Union[None, DocumentConverterResult]: def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
return (
df.rename(columns=lambda col: self._clean_colname(col))
.dropna(how="all", axis=1)
.dropna(how="all", axis=0)
)
def convert(
self, local_path, beautify: bool = True, **kwargs
) -> Union[None, DocumentConverterResult]:
# Bail if not a XLSX # Bail if not a XLSX
extension = kwargs.get("file_extension", "") extension = kwargs.get("file_extension", "")
if extension.lower() != ".xlsx": if extension.lower() != ".xlsx":
@ -535,14 +544,13 @@ class XlsxConverter(HtmlConverter):
sheets = pd.read_excel(local_path, sheet_name=None) sheets = pd.read_excel(local_path, sheet_name=None)
md_content = "" md_content = ""
for s in sheets: for name, sheet in sheets.items():
md_content += f"## {s}\n" md_content += f"## {name}\n"
sheet = sheets[s] df = self._clean_dataframe(sheet) if beautify else sheet
sheet.columns = list(map(self._clean_colname, sheet.columns))
html_content = ( html_content = (
sheet.dropna(how="all", axis=1) df.to_html(index=False, na_rep="")
.dropna(how="all", axis=0) if beautify
.to_html(index=False, na_rep="") else df.to_html(index=False)
) )
md_content += self._convert(html_content).text_content.strip() + "\n\n" md_content += self._convert(html_content).text_content.strip() + "\n\n"

11
tests/test_markitdown.py
View file

@ -42,6 +42,7 @@ XLSX_TEST_STRINGS = [
"affc7dad-52dc-4b98-9b5d-51e65d8a8ad0", "affc7dad-52dc-4b98-9b5d-51e65d8a8ad0",
] ]
DOCX_TEST_STRINGS = [ DOCX_TEST_STRINGS = [
"314b0a30-5b04-470b-b9f7-eed2c2bec74a", "314b0a30-5b04-470b-b9f7-eed2c2bec74a",
"49e168b7-d2ae-407f-a055-2167576f39a1", "49e168b7-d2ae-407f-a055-2167576f39a1",
@ -139,14 +140,18 @@ def test_markitdown_local() -> None:
markitdown = MarkItDown() markitdown = MarkItDown()
# Test XLSX processing # Test XLSX processing
result = markitdown.convert(os.path.join(TEST_FILES_DIR, "test.xlsx")) # XlsxConverter has an additional kwarg `beautify`, which defaults to True
result = markitdown.convert(
os.path.join(TEST_FILES_DIR, "test.xlsx"), beautify=False
)
result_cleaned = markitdown.convert(os.path.join(TEST_FILES_DIR, "test.xlsx"))
# Check assertions # Check assertions
for test_string in XLSX_TEST_STRINGS: for test_string in XLSX_TEST_STRINGS:
text_content = result.text_content.replace("\\", "") text_content = result.text_content.replace("\\", "")
assert test_string in text_content assert test_string in text_content
# Check negations # Check negations
assert "Unnamed:" not in text_content assert "Unnamed:" not in result_cleaned.text_content
assert "NaN" not in text_content assert "NaN" not in result_cleaned.text_content
# Test DOCX processing # Test DOCX processing
result = markitdown.convert(os.path.join(TEST_FILES_DIR, "test.docx")) result = markitdown.convert(os.path.join(TEST_FILES_DIR, "test.docx"))