Enable dynamic file type registration (#3946)

The purpose of this PR is to enable registering new file types
dynamically.

The PR enables this through 2 primary functions:

1. `unstructured.file_utils.model.create_file_type` This registers the
new `FileType` enum which enables the rest of unstructured to understand
a new type of file
2. `unstructured.file_utils.model.register_partitioner` Decorator that
enables registering a partitioner function to run for a file type.

---------

Co-authored-by: Roman Isecke <[email protected]>

Author: vangheem

CHANGELOG.md

@@ -1,7 +1,10 @@
-## 0.16.24-dev4
+## 0.16.24-dev5
 
 ### Enhancements
 
+- **Support dynamic partitioner file type registration**. Use `create_file_type` to create new file type that can be handled
+  in unstructured and `register_partitioner` to enable registering your own partitioner for any file type.
+
 - **`extract_image_block_types` now also works for CamelCase elemenet type names**. Previously `NarrativeText` and similar CamelCase element types can't be extracted using the mentioned parameter in `partition`. Now figures for those elements can be extracted like `Image` and `Table` elements
 
 ### Features

test_unstructured/file_utils/test_filetype.py

@@ -25,7 +25,7 @@
     detect_filetype,
     is_json_processable,
 )
-from unstructured.file_utils.model import FileType
+from unstructured.file_utils.model import FileType, create_file_type
 
 is_in_docker = os.path.exists("/.dockerenv")
 
@@ -467,6 +467,13 @@ def test_it_detect_CSV_from_path_and_file_when_content_contains_escaped_commas()
         assert detect_filetype(file=f) == FileType.CSV
 
 
+def test_it_detects_correct_file_type_for_custom_types(tmp_path):
+    file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
+    dumb_file = tmp_path / "dumb.foo"
+    dumb_file.write_bytes(b"38v8df889qw8sdfj")
+    assert detect_filetype(file_path=str(dumb_file), content_type="application/foo") is file_type
+
+
 # ================================================================================================
 # Describe `is_json_processable()`
 # ================================================================================================

test_unstructured/file_utils/test_model.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from unstructured.file_utils.model import FileType
+from unstructured.file_utils.model import FileType, create_file_type, register_partitioner
 
 
 class DescribeFileType:
@@ -225,3 +225,21 @@ def it_provides_access_to_the_partitioner_shortname(
         self, file_type: FileType, expected_value: str
     ):
         assert file_type.partitioner_shortname == expected_value
+
+
+def test_create_file_type():
+    file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
+
+    assert FileType.from_extension(".foo") is file_type
+    assert FileType.from_mime_type("application/foo") is file_type
+
+
+def test_register_partitioner():
+    file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
+
+    @register_partitioner(file_type)
+    def partition_foo():
+        pass
+
+    assert file_type.partitioner_function_name == "partition_foo"
+    assert file_type.partitioner_module_qname == "test_unstructured.file_utils.test_model"

test_unstructured/partition/test_auto.py

@@ -43,7 +43,7 @@
     Title,
 )
 from unstructured.file_utils.filetype import detect_filetype
-from unstructured.file_utils.model import FileType
+from unstructured.file_utils.model import FileType, create_file_type, register_partitioner
 from unstructured.partition.auto import _PartitionerLoader, partition
 from unstructured.partition.common import UnsupportedFileFormatError
 from unstructured.partition.utils.constants import PartitionStrategy
@@ -1331,3 +1331,17 @@ def expected_docx_elements():
         Text("2023"),
         Address("DOYLESTOWN, PA 18901"),
     ]
+
+
+def _test_partition_foo():
+    pass
+
+
+def test_auto_partition_works_with_custom_types(
+    request: FixtureRequest,
+):
+    file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
+
+    register_partitioner(file_type)(_test_partition_foo)
+    loader = _PartitionerLoader()
+    assert loader.get(file_type) is _test_partition_foo

unstructured/__version__.py

@@ -1 +1 @@
-__version__ = "0.16.24-dev4"  # pragma: no cover
+__version__ = "0.16.24-dev5"  # pragma: no cover

unstructured/file_utils/model.py

@@ -3,7 +3,43 @@
 from __future__ import annotations
 
 import enum
-from typing import Iterable, cast
+from typing import TYPE_CHECKING, Callable, Iterable, Type, cast
+
+from typing_extensions import ParamSpec
+
+if TYPE_CHECKING:
+    from unstructured.documents.elements import Element
+else:
+    Element = None
+
+
+def _create_file_type_enum(
+    cls: Type["FileType"],
+    value: str,
+    partitioner_shortname: str | None,
+    importable_package_dependencies: Iterable[str],
+    extra_name: str | None,
+    extensions: Iterable[str],
+    canonical_mime_type: str,
+    alias_mime_types: Iterable[str],
+    partitioner_full_module_path: str | None = None,
+) -> "FileType":
+    """
+    Moving here instead of directly in the FileType.__new__ allows us
+    to dynamically create new enum properties.
+
+    FileType.__new__ does not work with dynamic properties.
+    """
+    val = object.__new__(cls)
+    val._value_ = value
+    val._partitioner_shortname = partitioner_shortname
+    val._importable_package_dependencies = tuple(importable_package_dependencies)
+    val._extra_name = extra_name
+    val._extensions = tuple(extensions)
+    val._canonical_mime_type = canonical_mime_type
+    val._alias_mime_types = tuple(alias_mime_types)
+    val._partitioner_full_module_path = partitioner_full_module_path
+    return val
 
 
 class FileType(enum.Enum):
@@ -30,6 +66,9 @@ class FileType(enum.Enum):
     _alias_mime_types: tuple[str, ...]
     """MIME-types accepted as identifying this file-type."""
 
+    _partitioner_full_module_path: str | None
+    """Fully-qualified name of module providing partitioner for this file-type."""
+
     def __new__(
         cls,
         value: str,
@@ -39,16 +78,19 @@ def __new__(
         extensions: Iterable[str],
         canonical_mime_type: str,
         alias_mime_types: Iterable[str],
+        partitioner_full_module_path: str | None = None,
     ):
-        self = object.__new__(cls)
-        self._value_ = value
-        self._partitioner_shortname = partitioner_shortname
-        self._importable_package_dependencies = tuple(importable_package_dependencies)
-        self._extra_name = extra_name
-        self._extensions = tuple(extensions)
-        self._canonical_mime_type = canonical_mime_type
-        self._alias_mime_types = tuple(alias_mime_types)
-        return self
+        return _create_file_type_enum(
+            cls,
+            value,
+            partitioner_shortname,
+            importable_package_dependencies,
+            extra_name,
+            extensions,
+            canonical_mime_type,
+            alias_mime_types,
+            partitioner_full_module_path,
+        )
 
     def __lt__(self, other: FileType) -> bool:
         """Makes `FileType` members comparable with relational operators, at least with `<`.
@@ -132,7 +174,7 @@ def is_partitionable(self) -> bool:
         distinguishing file-types like WAV, ZIP, EMPTY, and UNK which are legitimate file-types
         but have no associated partitioner.
         """
-        return bool(self._partitioner_shortname)
+        return bool(self._partitioner_shortname) or bool(self._partitioner_full_module_path)
 
     @property
     def mime_type(self) -> str:
@@ -154,6 +196,9 @@ def partitioner_function_name(self) -> str:
         # -- Raise when this property is accessed on a FileType member that has no partitioner
         # -- shortname. This prevents a harder-to-find bug from appearing far away from this call
         # -- when code would try to `getattr(module, None)` or whatever.
+        if full_module_path := self._partitioner_full_module_path:
+            return full_module_path.split(".")[-1]
+
         if (shortname := self._partitioner_shortname) is None:
             raise ValueError(
                 f"`.partitioner_function_name` is undefined because FileType.{self.name} is not"
@@ -171,6 +216,9 @@ def partitioner_module_qname(self) -> str:
         # -- Raise when this property is accessed on a FileType member that has no partitioner
         # -- shortname. This prevents a harder-to-find bug from appearing far away from this call
         # -- when code would try to `importlib.import_module(None)` or whatever.
+        if full_module_path := self._partitioner_full_module_path:
+            return ".".join(full_module_path.split(".")[:-1])
+
         if (shortname := self._partitioner_shortname) is None:
             raise ValueError(
                 f"`.partitioner_module_qname` is undefined because FileType.{self.name} is not"
@@ -446,3 +494,42 @@ def partitioner_shortname(self) -> str | None:
         "inode/x-empty",
         cast(list[str], []),
     )
+
+
+def create_file_type(
+    name: str,
+    *,
+    canonical_mime_type: str,
+    importable_package_dependencies: Iterable[str] | None = None,
+    extra_name: str | None = None,
+    extensions: Iterable[str] | None = None,
+    alias_mime_types: Iterable[str] | None = None,
+) -> FileType:
+    """Register a new FileType member."""
+    type_ = _create_file_type_enum(
+        FileType,
+        name,
+        None,
+        importable_package_dependencies or cast(list[str], []),
+        extra_name,
+        extensions or cast(list[str], []),
+        canonical_mime_type,
+        alias_mime_types or cast(list[str], []),
+        None,
+    )
+    type_._name_ = name
+    FileType._member_map_[name] = type_
+    return type_
+
+
+_P = ParamSpec("_P")
+
+
+def register_partitioner(
+    file_type: FileType,
+) -> Callable[[Callable[_P, list[Element]]], Callable[_P, list[Element]]]:
+    def decorator(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
+        file_type._partitioner_full_module_path = func.__module__ + "." + func.__name__
+        return func
+
+    return decorator