feat: add metadata_filename parameter across all partition functions (#811)

* fix conflicts

* add tests and clean metadata_filename in partitions

* fix test_email and remove comments

* make tidy/check

* update changelog and version

* fix tests

* make tidy again

Author: Coniferish

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features
 
+* Add metadata_filename parameter across all partition functions
+
 ### Fixes
 * Fix KeyError when `isd_to_elements` doesn't find a type
 

test_unstructured/partition/test_csv.py

@@ -12,6 +12,16 @@ def test_partition_csv_from_filename(filename="example-docs/stanley-cups.csv"):
     assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
     assert elements[0].metadata.text_as_html == EXPECTED_TABLE
     assert elements[0].metadata.filetype == EXPECTED_FILETYPE
+    assert elements[0].metadata.filename == "stanley-cups.csv"
+
+
+def test_partition_csv_from_filename_with_metadata_filename(
+    filename="example-docs/stanley-cups.csv",
+):
+    elements = partition_csv(filename=filename, metadata_filename="test")
+
+    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
+    assert elements[0].metadata.filename == "test"
 
 
 def test_partition_csv_from_file(filename="example-docs/stanley-cups.csv"):
@@ -22,6 +32,15 @@ def test_partition_csv_from_file(filename="example-docs/stanley-cups.csv"):
     assert isinstance(elements[0], Table)
     assert elements[0].metadata.text_as_html == EXPECTED_TABLE
     assert elements[0].metadata.filetype == EXPECTED_FILETYPE
+    assert elements[0].metadata.filename is None
+
+
+def test_partition_csv_from_file_with_metadata_filename(filename="example-docs/stanley-cups.csv"):
+    with open(filename, "rb") as f:
+        elements = partition_csv(file=f, metadata_filename="test")
+
+    assert clean_extra_whitespace(elements[0].text) == EXPECTED_TEXT
+    assert elements[0].metadata.filename == "test"
 
 
 def test_partition_csv_can_exclude_metadata(filename="example-docs/stanley-cups.csv"):
@@ -31,3 +50,4 @@ def test_partition_csv_can_exclude_metadata(filename="example-docs/stanley-cups.
     assert isinstance(elements[0], Table)
     assert elements[0].metadata.text_as_html is None
     assert elements[0].metadata.filetype is None
+    assert elements[0].metadata.filename is None

test_unstructured/partition/test_doc.py

@@ -56,28 +56,40 @@ def expected_elements():
     ]
 
 
-def test_partition_doc_with_filename(mock_document, expected_elements, tmpdir, capsys):
+def test_partition_doc_from_filename(mock_document, expected_elements, tmpdir, capsys):
     docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
     mock_document.save(docx_filename)
     convert_office_doc(docx_filename, tmpdir.dirname, "doc")
-
     elements = partition_doc(filename=doc_filename)
     assert elements == expected_elements
     assert elements[0].metadata.filename == "mock_document.doc"
     assert elements[0].metadata.file_directory == tmpdir.dirname
-
     assert capsys.readouterr().out == ""
     assert capsys.readouterr().err == ""
 
 
-def test_partition_doc_matches_partition_docx(mock_document, expected_elements, tmpdir):
+def test_partition_doc_from_filename_with_metadata_filename(
+    mock_document,
+    expected_elements,
+    tmpdir,
+):
     docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
     mock_document.save(docx_filename)
     convert_office_doc(docx_filename, tmpdir.dirname, "doc")
 
-    partition_doc(filename=doc_filename) == partition_docx(filename=docx_filename)
+    elements = partition_doc(filename=doc_filename, metadata_filename="test")
+    assert elements == expected_elements
+    assert all(element.metadata.filename == "test" for element in elements)
+
+
+def test_partition_doc_matches_partition_docx(mock_document, expected_elements, tmpdir):
+    docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
+    doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
+    mock_document.save(docx_filename)
+    convert_office_doc(docx_filename, tmpdir.dirname, "doc")
+    assert partition_doc(filename=doc_filename) == partition_docx(filename=docx_filename)
 
 
 def test_partition_raises_with_missing_doc(mock_document, expected_elements, tmpdir):
@@ -87,7 +99,7 @@ def test_partition_raises_with_missing_doc(mock_document, expected_elements, tmp
         partition_doc(filename=doc_filename)
 
 
-def test_partition_doc_with_file(mock_document, expected_elements, tmpdir, capsys):
+def test_partition_doc_from_file(mock_document, expected_elements, tmpdir, capsys):
     docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
     mock_document.save(docx_filename)
@@ -96,9 +108,22 @@ def test_partition_doc_with_file(mock_document, expected_elements, tmpdir, capsy
     with open(doc_filename, "rb") as f:
         elements = partition_doc(file=f)
     assert elements == expected_elements
-
     assert capsys.readouterr().out == ""
     assert capsys.readouterr().err == ""
+    for element in elements:
+        assert element.metadata.filename is None
+
+
+def test_partition_doc_from_file_with_metadata_filename(mock_document, tmpdir):
+    docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
+    doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
+    mock_document.save(docx_filename)
+    convert_office_doc(docx_filename, tmpdir.dirname, "doc")
+
+    with open(doc_filename, "rb") as f:
+        elements = partition_doc(file=f, metadata_filename="test")
+    for element in elements:
+        assert element.metadata.filename == "test"
 
 
 def test_partition_doc_raises_with_both_specified(mock_document, tmpdir):
@@ -116,7 +141,7 @@ def test_partition_doc_raises_with_neither():
         partition_doc()
 
 
-def test_partition_doc_with_file_exclude_metadata(mock_document, tmpdir):
+def test_partition_doc_from_file_exclude_metadata(mock_document, tmpdir):
     docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
     mock_document.save(docx_filename)
@@ -130,7 +155,7 @@ def test_partition_doc_with_file_exclude_metadata(mock_document, tmpdir):
     assert elements[0].metadata.filename is None
 
 
-def test_partition_doc_with_filename_exclude_metadata(mock_document, tmpdir):
+def test_partition_doc_from_filename_exclude_metadata(mock_document, tmpdir):
     docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
     mock_document.save(docx_filename)

test_unstructured/partition/test_docx.py

@@ -57,13 +57,22 @@ def expected_elements():
     ]
 
 
-def test_partition_docx_with_filename(mock_document, expected_elements, tmpdir):
+def test_partition_docx_from_filename(mock_document, expected_elements, tmpdir):
     filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     mock_document.save(filename)
 
     elements = partition_docx(filename=filename)
     assert elements == expected_elements
     assert elements[0].metadata.page_number is None
+    for element in elements:
+        assert element.metadata.filename == "mock_document.docx"
+
+
+def test_partition_docx_from_filename_with_metadata_filename(mock_document, tmpdir):
+    filename = os.path.join(tmpdir.dirname, "mock_document.docx")
+    mock_document.save(filename)
+    elements = partition_docx(filename=filename, metadata_filename="test")
+    assert all(element.metadata.filename == "test" for element in elements)
 
 
 def test_partition_docx_with_spooled_file(mock_document, expected_elements, tmpdir):
@@ -79,15 +88,30 @@ def test_partition_docx_with_spooled_file(mock_document, expected_elements, tmpd
         spooled_temp_file.seek(0)
         elements = partition_docx(file=spooled_temp_file)
         assert elements == expected_elements
+        for element in elements:
+            assert element.metadata.filename is None
 
 
-def test_partition_docx_with_file(mock_document, expected_elements, tmpdir):
+def test_partition_docx_from_file(mock_document, expected_elements, tmpdir):
     filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     mock_document.save(filename)
 
     with open(filename, "rb") as f:
         elements = partition_docx(file=f)
     assert elements == expected_elements
+    for element in elements:
+        assert element.metadata.filename is None
+
+
+def test_partition_docx_from_file_with_metadata_filename(mock_document, expected_elements, tmpdir):
+    filename = os.path.join(tmpdir.dirname, "mock_document.docx")
+    mock_document.save(filename)
+
+    with open(filename, "rb") as f:
+        elements = partition_docx(file=f, metadata_filename="test")
+    assert elements == expected_elements
+    for element in elements:
+        assert element.metadata.filename == "test"
 
 
 def test_partition_docx_raises_with_both_specified(mock_document, tmpdir):
@@ -125,30 +149,36 @@ def test_partition_docx_grabs_header_and_footer(filename="example-docs/handbook-
     elements = partition_docx(filename=filename)
     assert elements[0] == Header("US Trustee Handbook")
     assert elements[-1] == Footer("Copyright")
+    for element in elements:
+        assert element.metadata.filename == "handbook-1p.docx"
 
 
 def test_partition_docx_includes_pages_if_present(filename="example-docs/handbook-1p.docx"):
     elements = partition_docx(filename=filename, include_page_breaks=False)
     assert "PageBreak" not in [elem.category for elem in elements]
     assert elements[1].metadata.page_number == 1
     assert elements[-2].metadata.page_number == 2
+    for element in elements:
+        assert element.metadata.filename == "handbook-1p.docx"
 
 
 def test_partition_docx_includes_page_breaks(filename="example-docs/handbook-1p.docx"):
     elements = partition_docx(filename=filename, include_page_breaks=True)
     assert "PageBreak" in [elem.category for elem in elements]
     assert elements[1].metadata.page_number == 1
     assert elements[-2].metadata.page_number == 2
+    for element in elements:
+        assert element.metadata.filename == "handbook-1p.docx"
 
 
-def test_partition_docx_with_filename_exclude_metadata(filename="example-docs/handbook-1p.docx"):
+def test_partition_docx_from_filename_exclude_metadata(filename="example-docs/handbook-1p.docx"):
     elements = partition_docx(filename=filename, include_metadata=False)
     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_docx_with_file_exclude_metadata(mock_document, tmpdir):
+def test_partition_docx_from_file_exclude_metadata(mock_document, tmpdir):
     filename = os.path.join(tmpdir.dirname, "mock_document.docx")
     mock_document.save(filename)