refactor: address several input formats with same mime type

Signed-off-by: Cesar Berrospi Ramis <75900930+ceberam@users.noreply.github.com>

docling/backend/xml_uspto_backend.py

@@ -89,7 +89,10 @@ def _set_parser(self, doctype: str) -> None:
             self.parser = PatentUsptoIce()
         elif "us-grant-025" in doctype_line:
             self.parser = PatentUsptoGrantV2()
-        elif "pap-v1" in doctype_line:
+        elif all(
+            item in doctype_line
+            for item in ("patent-application-publication", "pap-v1")
+        ):
             self.parser = PatentUsptoAppV1()
         else:
             self.parser = None

docling/datamodel/base_models.py

@@ -86,8 +86,10 @@ class OutputFormat(str, Enum):
     InputFormat.XML_USPTO: ["application/xml", "text/plain"],
 }
 
-MimeTypeToFormat = {
-    mime: fmt for fmt, mimes in FormatToMimeType.items() for mime in mimes
+MimeTypeToFormat: dict[str, list[InputFormat]] = {
+    mime: [fmt for fmt in FormatToMimeType if mime in FormatToMimeType[fmt]]
+    for value in FormatToMimeType.values()
+    for mime in value
 }
 
 

docling/datamodel/document.py

@@ -3,7 +3,17 @@
 from enum import Enum
 from io import BytesIO
 from pathlib import Path, PurePath
-from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Set, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Set,
+    Type,
+    Union,
+)
 
 import filetype
 from docling_core.types.doc import (
@@ -235,31 +245,31 @@ def docs(
             if isinstance(obj, Path):
                 yield InputDocument(
                     path_or_stream=obj,
-                    format=format,
+                    format=format,  # type: ignore[arg-type]
                     filename=obj.name,
                     limits=self.limits,
                     backend=backend,
                 )
             elif isinstance(obj, DocumentStream):
                 yield InputDocument(
                     path_or_stream=obj.stream,
-                    format=format,
+                    format=format,  # type: ignore[arg-type]
                     filename=obj.name,
                     limits=self.limits,
                     backend=backend,
                 )
             else:
                 raise RuntimeError(f"Unexpected obj type in iterator: {type(obj)}")
 
-    def _guess_format(self, obj: Union[Path, DocumentStream]):
+    def _guess_format(self, obj: Union[Path, DocumentStream]) -> Optional[InputFormat]:
         content = b""  # empty binary blob
-        format = None
+        formats: list[InputFormat] = []
 
         if isinstance(obj, Path):
             mime = filetype.guess_mime(str(obj))
             if mime is None:
                 ext = obj.suffix[1:]
-                mime = self._mime_from_extension(ext)
+                mime = _DocumentConversionInput._mime_from_extension(ext)
             if mime is None:  # must guess from
                 with obj.open("rb") as f:
                     content = f.read(1024)  # Read first 1KB
@@ -274,15 +284,52 @@ def _guess_format(self, obj: Union[Path, DocumentStream]):
                     if ("." in obj.name and not obj.name.startswith("."))
                     else ""
                 )
-                mime = self._mime_from_extension(ext)
+                mime = _DocumentConversionInput._mime_from_extension(ext)
 
-        mime = mime or self._detect_html_xhtml(content)
+        mime = mime or _DocumentConversionInput._detect_html_xhtml(content)
         mime = mime or "text/plain"
+        formats = MimeTypeToFormat.get(mime, [])
+        if formats:
+            if len(formats) == 1 and mime != "text/plain":
+                return formats[0]
+            else:  # ambiguity cases
+                return _DocumentConversionInput._guess_from_content(
+                    content, mime, formats
+                )
+        else:
+            return None
+
+    @staticmethod
+    def _guess_from_content(
+        content: bytes, mime: str, formats: list[InputFormat]
+    ) -> Optional[InputFormat]:
+        """Guess the input format of a document by checking part of its content."""
+        input_format: Optional[InputFormat] = None
+        content_str = content.decode("utf-8")
+
+        if mime == "application/xml":
+            match_doctype = re.search(r"<!DOCTYPE [^>]+>", content_str)
+            if match_doctype:
+                xml_doctype = match_doctype.group()
+                if InputFormat.XML_USPTO in formats and any(
+                    item in xml_doctype
+                    for item in (
+                        "us-patent-application-v4",
+                        "us-patent-grant-v4",
+                        "us-grant-025",
+                        "patent-application-publication",
+                    )
+                ):
+                    input_format = InputFormat.XML_USPTO
+
+        elif mime == "text/plain":
+            if InputFormat.XML_USPTO in formats and content_str.startswith("PATN\r\n"):
+                input_format = InputFormat.XML_USPTO
 
-        format = MimeTypeToFormat.get(mime)
-        return format
+        return input_format
 
-    def _mime_from_extension(self, ext):
+    @staticmethod
+    def _mime_from_extension(ext):
         mime = None
         if ext in FormatToExtensions[InputFormat.ASCIIDOC]:
             mime = FormatToMimeType[InputFormat.ASCIIDOC][0]
@@ -293,7 +340,19 @@ def _mime_from_extension(self, ext):
 
         return mime
 
-    def _detect_html_xhtml(self, content):
+    @staticmethod
+    def _detect_html_xhtml(
+        content: bytes,
+    ) -> Optional[Literal["application/xhtml+xml", "application/xml", "text/html"]]:
+        """Guess the mime type of an XHTML, HTML, or XML file from its content.
+
+        Args:
+            content: A short piece of a document from its beginning.
+
+        Returns:
+            The mime type of an XHTML, HTML, or XML file, or None if the content does
+              not match any of these formats.
+        """
         content_str = content.decode("ascii", errors="ignore").lower()
         # Remove XML comments
         content_str = re.sub(r"<!--(.*?)-->", "", content_str, flags=re.DOTALL)
@@ -302,6 +361,8 @@ def _detect_html_xhtml(self, content):
         if re.match(r"<\?xml", content_str):
             if "xhtml" in content_str[:1000]:
                 return "application/xhtml+xml"
+            else:
+                return "application/xml"
 
         if re.match(r"<!doctype\s+html|<html|<head|<body", content_str):
             return "text/html"

tests/test_input_doc.py

@@ -3,7 +3,7 @@
 
 from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
 from docling.datamodel.base_models import DocumentStream, InputFormat
-from docling.datamodel.document import InputDocument
+from docling.datamodel.document import InputDocument, _DocumentConversionInput
 
 
 def test_in_doc_from_valid_path():
@@ -39,6 +39,42 @@ def test_in_doc_from_invalid_buf():
     assert doc.valid == False
 
 
+def test_guess_format():
+    """Test docling.datamodel.document._DocumentConversionInput.__guess_format"""
+    dci = _DocumentConversionInput(path_or_stream_iterator=[])
+
+    # Valid PDF
+    buf = BytesIO(Path("./tests/data/2206.01062.pdf").open("rb").read())
+    stream = DocumentStream(name="my_doc.pdf", stream=buf)
+    assert dci._guess_format(stream) == InputFormat.PDF
+    doc_path = Path("./tests/data/2206.01062.pdf")
+    assert dci._guess_format(doc_path) == InputFormat.PDF
+
+    # Valid MS Doc
+    buf = BytesIO(Path("./tests/data/docx/lorem_ipsum.docx").open("rb").read())
+    stream = DocumentStream(name="lorem_ipsum.docx", stream=buf)
+    assert dci._guess_format(stream) == InputFormat.DOCX
+    doc_path = Path("./tests/data/docx/lorem_ipsum.docx")
+    assert dci._guess_format(doc_path) == InputFormat.DOCX
+
+    # Valid XML USPTO patent
+    buf = BytesIO(Path("./tests/data/uspto/ipa20110039701.xml").open("rb").read())
+    stream = DocumentStream(name="ipa20110039701.xml", stream=buf)
+    assert dci._guess_format(stream) == InputFormat.XML_USPTO
+    doc_path = Path("./tests/data/uspto/ipa20110039701.xml")
+    assert dci._guess_format(doc_path) == InputFormat.XML_USPTO
+
+    buf = BytesIO(Path("./tests/data/uspto/pftaps057006474.txt").open("rb").read())
+    stream = DocumentStream(name="pftaps057006474.txt", stream=buf)
+    assert dci._guess_format(stream) == InputFormat.XML_USPTO
+    doc_path = Path("./tests/data/uspto/pftaps057006474.txt")
+    assert dci._guess_format(doc_path) == InputFormat.XML_USPTO
+
+    # Invalid XML USPTO patent
+    stream = DocumentStream(name="pftaps057006474.txt", stream=BytesIO(b"xyz"))
+    assert dci._guess_format(stream) == None
+
+
 def _make_input_doc(path):
     in_doc = InputDocument(
         path_or_stream=path,