the MIME type is not extension

Author: ianardee

mindee/documents/base.py

@@ -38,7 +38,7 @@ def __init__(
         if input_file:
             self.filepath = input_file.filepath
             self.filename = input_file.filename
-            self.file_extension = input_file.file_extension
+            self.file_extension = input_file.file_mimetype
 
         self.type = document_type
 

mindee/documents/financial_document.py

@@ -165,7 +165,7 @@ def request(
         :param include_words: Include Mindee vision words in http_response
         :param close_file: Whether to `close()` the file after parsing it.
         """
-        if "pdf" in input_file.file_extension:
+        if "pdf" in input_file.file_mimetype:
             # invoices is index 0, receipts 1 (this should be cleaned up)
             index = 0
         else:

mindee/inputs.py

@@ -8,7 +8,7 @@
 
 from mindee.logger import logger
 
-ALLOWED_EXTENSIONS = [
+ALLOWED_MIME_TYPES = [
     "image/png",
     "image/jpg",
     "image/jpeg",
@@ -20,12 +20,12 @@
 INPUT_TYPE_BASE64 = "base64"
 INPUT_TYPE_BYTES = "bytes"
 INPUT_TYPE_PATH = "path"
-INPUT_TYPE_DUMMY = "dummy"
 
 
 class InputDocument:
     file_object: BinaryIO
     filename: str
+    file_mimetype: str
     input_type: str
     filepath: Optional[str] = None
     cut_pdf: bool
@@ -39,17 +39,9 @@ def __init__(
     ):
         assert 0 < n_pdf_pages <= 3
         self.input_type = input_type
-        self.file_extension = guess_type(self.filename)[0]
+        self._check_mimetype()
 
-        if (
-            self.file_extension not in ALLOWED_EXTENSIONS
-            and self.input_type != INPUT_TYPE_DUMMY
-        ):
-            raise AssertionError(
-                "File type not allowed, must be in {%s}" % ", ".join(ALLOWED_EXTENSIONS)
-            )
-
-        if self.file_extension == "application/pdf":
+        if self.file_mimetype == "application/pdf":
             self.check_pdf_open()
             count_pages = self.count_pdf_pages()
             if cut_pdf is True:
@@ -61,6 +53,19 @@ def __init__(
                 raise AssertionError(f"PDF pages are empty in: {self.filename}")
         logger.debug("Loaded new document '%s' from %s", self.filename, self.input_type)
 
+    def _check_mimetype(self) -> None:
+        file_mimetype = guess_type(self.filename)[0]
+        if file_mimetype:
+            self.file_mimetype = file_mimetype
+        else:
+            raise AssertionError(f"Could not determine MIME type of '{self.filename}'")
+
+        if self.file_mimetype not in ALLOWED_MIME_TYPES:
+            raise AssertionError(
+                "File type not allowed, must be one of {%s}"
+                % ", ".join(ALLOWED_MIME_TYPES)
+            )
+
     def count_pdf_pages(self) -> int:
         """
         Count the pages in the PDF.
@@ -73,10 +78,10 @@ def count_pdf_pages(self) -> int:
 
     def merge_pdf_pages(self, page_numbers: list) -> None:
         """
-        Create a new PDF from pages.
+        Create a new PDF from pages and set it to ``file_object``.
 
         :param page_numbers: List of pages number to use for merging in the original PDF.
-        :return: (void) Set the Input.file with the reconstructed pdf stream
+        :return: None
         """
         self.file_object.seek(0)
         new_pdf = pikepdf.Pdf.new()

mindee/response.py

@@ -5,9 +5,37 @@
 from mindee.logger import logger
 
 
+class DocumentResponse:
+    http_response: Dict[str, Any]
+    """Raw HTTP response JSON"""
+    document_type: str
+    """Document type"""
+
+    def __init__(
+        self,
+        doc_config: DocumentConfig,
+        http_response: dict,
+        pages: List[Document],
+        document_type: str,
+        document=TypeDocument,
+    ):
+        """
+        Container for the raw API response and the parsed document.
+
+        :param http_response: Raw HTTP response object
+        :param pages: List of document objects, page level
+        :param document: reconstructed object from all pages
+        :param document_type: Document class
+        """
+        self.http_response = http_response
+        self.document_type = document_type
+        setattr(self, doc_config.singular_name, document)
+        setattr(self, doc_config.plural_name, pages)
+
+
 def format_response(
     doc_config: DocumentConfig, http_response: dict, document_type: str, input_file
-):
+) -> DocumentResponse:
     """
     Create a `DocumentResponse`.
 
@@ -21,7 +49,7 @@ def format_response(
     http_response["input_type"] = input_file.input_type
     http_response["filename"] = input_file.filename
     http_response["filepath"] = input_file.filepath
-    http_response["file_extension"] = input_file.file_extension
+    http_response["file_extension"] = input_file.file_mimetype
     pages = []
 
     logger.debug("Handling API response")
@@ -50,31 +78,3 @@ def format_response(
         document_type=document_type,
         document=document_level,
     )
-
-
-class DocumentResponse:
-    http_response: Dict[str, Any]
-    """Raw HTTP response JSON"""
-    document_type: str
-    """Document type"""
-
-    def __init__(
-        self,
-        doc_config: DocumentConfig,
-        http_response: dict,
-        pages: List[Document],
-        document_type: str,
-        document=TypeDocument,
-    ):
-        """
-        Container for the raw API response and the parsed document.
-
-        :param http_response: Raw HTTP response object
-        :param pages: List of document objects, page level
-        :param document: reconstructed object from all pages
-        :param document_type: Document class
-        """
-        self.http_response = http_response
-        self.document_type = document_type
-        setattr(self, doc_config.singular_name, document)
-        setattr(self, doc_config.plural_name, pages)