mirror of
https://github.com/Unstructured-IO/unstructured.git
synced 2025-10-15 01:52:59 +00:00
cf31c9a2c4
Fixes recursion limit error that was being raised when partitioning
Excel documents of a certain size.
Previously we used a recursive method to find subtables within an excel
sheet. However this would run afoul of Python's recursion depth limit
when there was a contiguous block of more than 1000 cells within a
sheet. This function has been updated to use the NetworkX library which
avoids Python recursion issues.
* Updated `_get_connected_components` to use `networkx` graph methods
rather than implementing our own algorithm for finding contiguous groups
of cells within a sheet.
* Added a test and example doc that replicates the `RecursionError`
prior to the change.
* Added `networkx` to `extra_xlsx` dependencies and `pip-compile`d.
#### Testing:
The following run from a Python terminal should raise a `RecursionError`
on `main` and succeed on this branch:
```python
import sys
from unstructured.partition.xlsx import partition_xlsx
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)
```
Note: the recursion limit is different in different contexts. Checking
my own system, the default in a notebook seems to be 3000, but in a
terminal it's 1000. The documented Python default recursion limit is
1000.
247 lines
9.2 KiB
Python
247 lines
9.2 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
from test_unstructured.partition.test_constants import EXPECTED_TABLE, EXPECTED_TEXT, 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
|
|
assert elements[1].metadata.text_as_html == EXPECTED_TABLE
|
|
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
|
|
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
|
|
)
|
|
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
|
|
assert elements[1].metadata.text_as_html == EXPECTED_TABLE
|
|
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
|
|
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
|
|
)
|
|
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
|
|
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
|
|
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)
|