mirror of
https://github.com/Unstructured-IO/unstructured.git
synced 2025-10-07 22:31:40 +00:00
89fa88f076
The current code assumes the first line of csv and tsv files are a
header line. Most csv and tsv files don't have a header line, and even
for those that do, dropping this line may not be the desired behavior.
Here is a snippet of code that demonstrates the current behavior and the
proposed fix
```
import pandas as pd
from lxml.html.soupparser import fromstring as soupparser_fromstring
c1 = """
Stanley Cups,,
Team,Location,Stanley Cups
Blues,STL,1
Flyers,PHI,2
Maple Leafs,TOR,13
"""
f = "./test.csv"
with open(f, 'w') as ff:
ff.write(c1)
print("Suggested Improvement Keep First Line")
table = pd.read_csv(f, header=None)
html_text = table.to_html(index=False, header=False, na_rep="")
text = soupparser_fromstring(html_text).text_content()
print(text)
print("\n\nOriginal Looses First Line")
table = pd.read_csv(f)
html_text = table.to_html(index=False, header=False, na_rep="")
text = soupparser_fromstring(html_text).text_content()
print(text)
```
---------
Co-authored-by: cragwolfe <crag@unstructured.io>
Co-authored-by: Yao You <theyaoyou@gmail.com>
Co-authored-by: Yao You <yao@unstructured.io>
251 lines
9.2 KiB
Python
251 lines
9.2 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
from test_unstructured.partition.test_constants import (
|
|
EXPECTED_TABLE_XLSX,
|
|
EXPECTED_TEXT_XLSX,
|
|
EXPECTED_TITLE,
|
|
)
|
|
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
|
from unstructured.cleaners.core import clean_extra_whitespace
|
|
from unstructured.documents.elements import Table, Text, Title
|
|
from unstructured.partition.xlsx import partition_xlsx
|
|
|
|
EXPECTED_FILETYPE = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
|
|
|
EXCEPTED_PAGE_NAME = "Stanley Cups"
|
|
|
|
|
|
def test_partition_xlsx_from_filename(filename="example-docs/stanley-cups.xlsx"):
|
|
elements = partition_xlsx(filename=filename, include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert len(elements) == 4
|
|
|
|
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
|
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
|
assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
|
|
assert elements[1].metadata.page_number == 1
|
|
assert elements[1].metadata.filetype == EXPECTED_FILETYPE
|
|
assert elements[1].metadata.page_name == EXCEPTED_PAGE_NAME
|
|
assert elements[1].metadata.filename == "stanley-cups.xlsx"
|
|
|
|
|
|
def test_partition_xlsx_from_filename_with_emoji(filename="example-docs/emoji.xlsx"):
|
|
elements = partition_xlsx(filename=filename, include_header=False)
|
|
assert sum(isinstance(element, Text) for element in elements) == 1
|
|
assert len(elements) == 1
|
|
assert clean_extra_whitespace(elements[0].text) == "🤠😅"
|
|
|
|
|
|
def test_partition_xlsx_from_filename_with_metadata_filename(
|
|
filename="example-docs/stanley-cups.xlsx",
|
|
):
|
|
elements = partition_xlsx(filename=filename, metadata_filename="test", include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert sum(isinstance(element, Title) for element in elements) == 2
|
|
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
|
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
|
assert elements[0].metadata.filename == "test"
|
|
|
|
|
|
def test_partition_xlsx_from_filename_with_header(filename="example-docs/stanley-cups.xlsx"):
|
|
elements = partition_xlsx(filename=filename, include_header=True)
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert len(elements) == 2
|
|
assert (
|
|
clean_extra_whitespace(elements[0].text)
|
|
== "Stanley Cups Unnamed: 1 Unnamed: 2 " + EXPECTED_TEXT_XLSX
|
|
)
|
|
assert "<thead>" in elements[0].metadata.text_as_html
|
|
|
|
|
|
def test_partition_xlsx_from_file(filename="example-docs/stanley-cups.xlsx"):
|
|
with open(filename, "rb") as f:
|
|
elements = partition_xlsx(file=f, include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert len(elements) == 4
|
|
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
|
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
|
assert elements[1].metadata.text_as_html == EXPECTED_TABLE_XLSX
|
|
assert elements[1].metadata.page_number == 1
|
|
assert elements[1].metadata.filetype == EXPECTED_FILETYPE
|
|
assert elements[1].metadata.page_name == EXCEPTED_PAGE_NAME
|
|
assert elements[1].metadata.filename is None
|
|
|
|
|
|
def test_partition_xlsx_from_file_with_metadata_filename(filename="example-docs/stanley-cups.xlsx"):
|
|
with open(filename, "rb") as f:
|
|
elements = partition_xlsx(file=f, metadata_filename="test", include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
|
assert elements[1].metadata.filename == "test"
|
|
|
|
|
|
def test_partition_xlsx_from_file_with_header(filename="example-docs/stanley-cups.xlsx"):
|
|
with open(filename, "rb") as f:
|
|
elements = partition_xlsx(file=f, include_header=True)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert len(elements) == 2
|
|
assert (
|
|
clean_extra_whitespace(elements[0].text)
|
|
== "Stanley Cups Unnamed: 1 Unnamed: 2 " + EXPECTED_TEXT_XLSX
|
|
)
|
|
assert "<thead>" in elements[0].metadata.text_as_html
|
|
|
|
|
|
def test_partition_xlsx_filename_exclude_metadata(filename="example-docs/stanley-cups.xlsx"):
|
|
elements = partition_xlsx(filename=filename, include_metadata=False, include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert len(elements) == 4
|
|
|
|
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
|
assert elements[1].metadata.text_as_html is None
|
|
assert elements[1].metadata.page_number is None
|
|
assert elements[1].metadata.filetype is None
|
|
assert elements[1].metadata.page_name is None
|
|
assert elements[1].metadata.filename is None
|
|
|
|
|
|
def test_partition_xlsx_from_file_exclude_metadata(filename="example-docs/stanley-cups.xlsx"):
|
|
with open(filename, "rb") as f:
|
|
elements = partition_xlsx(file=f, include_metadata=False, include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert sum(isinstance(element, Title) for element in elements) == 2
|
|
assert len(elements) == 4
|
|
|
|
assert clean_extra_whitespace(elements[0].text) == EXPECTED_TITLE
|
|
assert clean_extra_whitespace(elements[1].text) == EXPECTED_TEXT_XLSX
|
|
assert elements[0].metadata.text_as_html is None
|
|
assert elements[0].metadata.page_number is None
|
|
assert elements[0].metadata.filetype is None
|
|
assert elements[0].metadata.page_name is None
|
|
assert elements[0].metadata.filename is None
|
|
|
|
|
|
def test_partition_xlsx_metadata_date(
|
|
mocker,
|
|
filename="example-docs/stanley-cups.xlsx",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.xlsx.get_last_modified_date",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
elements = partition_xlsx(
|
|
filename=filename,
|
|
)
|
|
|
|
assert elements[0].metadata.last_modified == mocked_last_modification_date
|
|
|
|
|
|
def test_partition_xlsx_with_custom_metadata_date(
|
|
mocker,
|
|
filename="example-docs/stanley-cups.xlsx",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
expected_last_modification_date = "2020-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.xlsx.get_last_modified_date",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
elements = partition_xlsx(
|
|
filename=filename,
|
|
metadata_last_modified=expected_last_modification_date,
|
|
)
|
|
|
|
assert elements[0].metadata.last_modified == expected_last_modification_date
|
|
|
|
|
|
def test_partition_xlsx_from_file_metadata_date(
|
|
mocker,
|
|
filename="example-docs/stanley-cups.xlsx",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.xlsx.get_last_modified_date_from_file",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
with open(filename, "rb") as f:
|
|
elements = partition_xlsx(
|
|
file=f,
|
|
)
|
|
|
|
assert elements[0].metadata.last_modified == mocked_last_modification_date
|
|
|
|
|
|
def test_partition_xlsx_from_file_with_custom_metadata_date(
|
|
mocker,
|
|
filename="example-docs/stanley-cups.xlsx",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
expected_last_modification_date = "2020-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.xlsx.get_last_modified_date_from_file",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
with open(filename, "rb") as f:
|
|
elements = partition_xlsx(file=f, metadata_last_modified=expected_last_modification_date)
|
|
|
|
assert elements[0].metadata.last_modified == expected_last_modification_date
|
|
|
|
|
|
def test_partition_xlsx_with_json():
|
|
elements = partition_xlsx(example_doc_path("stanley-cups.xlsx"), include_header=False)
|
|
assert_round_trips_through_JSON(elements)
|
|
|
|
|
|
@pytest.mark.skip("Needs to fix language detection for table. Currently detected as 'tur'")
|
|
def test_partition_xlsx_metadata_language_from_filename(filename="example-docs/stanley-cups.xlsx"):
|
|
elements = partition_xlsx(filename=filename, include_header=False)
|
|
|
|
assert sum(isinstance(element, Table) for element in elements) == 2
|
|
assert len(elements) == 4
|
|
|
|
assert elements[0].metadata.languages == ["eng"]
|
|
|
|
|
|
def test_partition_xlsx_subtables(filename="example-docs/vodafone.xlsx"):
|
|
elements = partition_xlsx(filename)
|
|
assert sum(isinstance(element, Table) for element in elements) == 3
|
|
assert len(elements) == 6
|
|
|
|
|
|
def test_partition_xlsx_element_metadata_has_languages():
|
|
filename = "example-docs/stanley-cups.xlsx"
|
|
elements = partition_xlsx(filename=filename)
|
|
assert elements[0].metadata.languages == ["eng"]
|
|
|
|
|
|
def test_partition_eml_respects_detect_language_per_element():
|
|
filename = "example-docs/language-docs/eng_spa.xlsx"
|
|
elements = partition_xlsx(filename=filename, detect_language_per_element=True)
|
|
langs = {element.metadata.languages[0] for element in elements}
|
|
assert "eng" in langs
|
|
assert "spa" in langs
|
|
|
|
|
|
def test_partition_xlsx_with_more_than_1k_cells():
|
|
old_recursion_limit = sys.getrecursionlimit()
|
|
try:
|
|
sys.setrecursionlimit(1000)
|
|
filename = "example-docs/more-than-1k-cells.xlsx"
|
|
partition_xlsx(filename=filename)
|
|
finally:
|
|
sys.setrecursionlimit(old_recursion_limit)
|