feat: add --metadata-include and --metadata-exclude parameters to unstructured-ingest (#368)

* added metadata in/exclude params

* updated process_file

* existing tests

* remove default behavior

* changelog and ci

* line length

* import

* import

* import sorted

* import

* type

* line length

* main

* ci

* json

* dict

* type ignore

* lint

* unit tests for process_file

* lint

* type changed to Optional(str)

* ci

* line length

* added mutex check

* nit

Author: natygyoon

Diff for: CHANGELOG.md

@@ -6,6 +6,7 @@
 
 ### Features
 
+* Add `--metadata-include` and `--metadata-exclude` parameters to `unstructured-ingest`
 * Add `clean_non_ascii_chars` to remove non-ascii characters from unicode string
 
 ### Fixes

Diff for: test_unstructured_ingest/test-ingest-azure.sh

@@ -4,6 +4,7 @@ SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
 cd "$SCRIPT_DIR"/.. || exit 1
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --remote-url abfs://container1/ \
     --azure-account-name azureunstructured1 \
     --structured-output-dir azure-ingest-output \

Diff for: test_unstructured_ingest/test-ingest-biomed-api.sh

@@ -10,6 +10,7 @@ if [[ "$(find test_unstructured_ingest/expected-structured-output/biomed-ingest-
 fi
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+   --metadata-exclude filename \
    --biomed-api-from "2019-01-02" \
    --biomed-api-until "2019-01-02+00:03:10" \
    --structured-output-dir biomed-ingest-output-api  \

Diff for: test_unstructured_ingest/test-ingest-biomed-path.sh

@@ -10,6 +10,7 @@ if [[ "$(find test_unstructured_ingest/expected-structured-output/biomed-ingest-
 fi
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --biomed-path "oa_pdf/07/07/sbaa031.073.PMC7234218.pdf" \
     --structured-output-dir biomed-ingest-output-path \
     --num-processes 2 \

Diff for: test_unstructured_ingest/test-ingest-github.sh

@@ -12,7 +12,12 @@ if [[ "$CI" == "true" ]]; then
 fi
 
 
-PYTHONPATH=. ./unstructured/ingest/main.py --github-url dcneiner/Downloadify --git-file-glob '*.html,*.txt' --structured-output-dir github-downloadify-output --verbose
+PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
+    --github-url dcneiner/Downloadify \
+    --git-file-glob '*.html,*.txt' \
+    --structured-output-dir github-downloadify-output \
+    --verbose
 
 if ! diff -ru test_unstructured_ingest/expected-structured-output/github-downloadify github-downloadify-output ; then
    echo

Diff for: test_unstructured_ingest/test-ingest-gitlab.sh

@@ -4,6 +4,7 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 cd "$SCRIPT_DIR"/.. || exit 1
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --gitlab-url https://gitlab.com/gitlab-com/content-sites/docsy-gitlab \
     --git-file-glob '*.md,*.txt' \
     --structured-output-dir gitlab-ingest-output \

Diff for: test_unstructured_ingest/test-ingest-s3.sh

@@ -9,7 +9,11 @@ if [[ "$(find test_unstructured_ingest/expected-structured-output/s3-small-batch
     exit 1
 fi
 
-PYTHONPATH=. ./unstructured/ingest/main.py --s3-url s3://utic-dev-tech-fixtures/small-pdf-set/ --s3-anonymous --structured-output-dir s3-small-batch-output
+PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
+    --s3-url s3://utic-dev-tech-fixtures/small-pdf-set/ \
+    --s3-anonymous \
+    --structured-output-dir s3-small-batch-output \
 
 if ! diff -ru test_unstructured_ingest/expected-structured-output/s3-small-batch s3-small-batch-output ; then
     echo

Diff for: test_unstructured_ingest/test-ingest-wikipedia.sh

@@ -4,6 +4,7 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 cd "$SCRIPT_DIR"/.. || exit 1
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --wikipedia-page-title "Open Source Software" \
     --structured-output-dir wikipedia-ingest-output \
     --num-processes 2 \

Diff for: test_unstructured_ingest/test_interfaces.py

@@ -0,0 +1,78 @@
+import os
+import pathlib
+
+import pytest
+
+from unstructured.ingest.connector.git import GitIngestDoc, SimpleGitConfig
+
+DIRECTORY = pathlib.Path(__file__).parent.resolve()
+EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "..", "example-docs")
+
+test_files = [
+    "layout-parser-paper-fast.jpg",
+    "layout-parser-paper-fast.pdf",
+]
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_include_filename(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_include="filename",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k == "filename"
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_include_filename_pagenum(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_include="filename,page_number",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k in ["filename", "page_number"]
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_exclude_filename(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_exclude="filename",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k != "filename"
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_exclude_filename_pagenum(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_exclude="filename,page_number",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k not in ["filename", "page_number"]

Diff for: unstructured/ingest/connector/biomed.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 from ftplib import FTP, error_perm
 from pathlib import Path
-from typing import List, Union
+from typing import List, Optional, Union
 
 import requests
 from bs4 import BeautifulSoup
@@ -48,6 +48,8 @@ class SimpleBiomedConfig(BaseConnectorConfig):
     output_dir: str
     re_download: bool = False
     preserve_downloads: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
     def _validate_date_args(self, date):
         date_formats = ["%Y-%m-%d", "%Y-%m-%d+%H:%M:%S"]

Diff for: unstructured/ingest/connector/fsspec.py

@@ -3,7 +3,7 @@
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Type
+from typing import Optional, Type
 
 from unstructured.ingest.interfaces import (
     BaseConnector,
@@ -29,6 +29,8 @@ class SimpleFsspecConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
     # fsspec specific options
     access_kwargs: dict = field(default_factory=dict)

Diff for: unstructured/ingest/connector/git.py

@@ -26,6 +26,8 @@ class SimpleGitConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
     repo_path: str = field(init=False, repr=False)
 

Diff for: unstructured/ingest/connector/google_drive.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 from mimetypes import guess_extension
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Optional
 
 from unstructured.file_utils.filetype import EXT_TO_FILETYPE
 from unstructured.file_utils.google_filetype import GOOGLE_DRIVE_EXPORT_TYPES
@@ -77,6 +77,8 @@ class SimpleGoogleDriveConfig(BaseConnectorConfig):
     output_dir: str
     re_download: bool = False
     preserve_downloads: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
     recursive: bool = False
 

Diff for: unstructured/ingest/connector/reddit.py

@@ -2,7 +2,7 @@
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from unstructured.ingest.interfaces import (
     BaseConnector,
@@ -31,6 +31,8 @@ class SimpleRedditConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
     def __post_init__(self):
         if self.num_posts <= 0:

Diff for: unstructured/ingest/connector/wikipedia.py

@@ -2,7 +2,7 @@
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from unstructured.ingest.interfaces import (
     BaseConnector,
@@ -26,6 +26,8 @@ class SimpleWikipediaConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
 
 @dataclass

Diff for: unstructured/ingest/interfaces.py

@@ -2,6 +2,7 @@
 through Unstructured."""
 
 from abc import ABC, abstractmethod
+from typing import Optional
 
 from unstructured.ingest.logger import logger
 from unstructured.partition.auto import partition
@@ -47,6 +48,8 @@ class BaseConnectorConfig(ABC):
     # where to write structured data outputs
     output_dir: str
     re_download: bool = False
+    metadata_include: Optional[str] = None
+    metadata_exclude: Optional[str] = None
 
 
 class BaseIngestDoc(ABC):
@@ -58,6 +61,8 @@ class BaseIngestDoc(ABC):
     Crucially, it is not responsible for the actual processing of the raw document.
     """
 
+    config: BaseConnectorConfig
+
     @property
     @abstractmethod
     def filename(self):
@@ -94,7 +99,24 @@ def process_file(self):
         self.isd_elems_no_filename = []
         for elem in isd_elems:
             # type: ignore
-            elem["metadata"].pop("filename", None)  # type: ignore[attr-defined]
+            if (
+                self.config.metadata_exclude is not None
+                and self.config.metadata_include is not None
+            ):
+                raise ValueError(
+                    "Arguments `--metadata-include` and `--metadata-exclude` are "
+                    "mutually exclusive with each other.",
+                )
+            elif self.config.metadata_exclude is not None:
+                ex_list = self.config.metadata_exclude.split(",")
+                for ex in ex_list:
+                    elem["metadata"].pop(ex, None)  # type: ignore[attr-defined]
+            elif self.config.metadata_include is not None:
+                in_list = self.config.metadata_include.split(",")
+                for k in elem["metadata"]:
+                    if k not in in_list:
+                        elem["metadata"].pop(k, None)  # type: ignore[attr-defined]
+
             elem.pop("coordinates")  # type: ignore[attr-defined]
             self.isd_elems_no_filename.append(elem)