fix: type of path_or_stream in PdfDocumentBackend

Signed-off-by: Michele Dolfi <[email protected]>
DS4SD · Aug 7, 2024 · d125b25 · d125b25
1 parent 9550db8
commit d125b25
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/docling/backend/abstract_backend.py b/docling/backend/abstract_backend.py
@@ -35,7 +35,7 @@ def unload(self):
 
 class PdfDocumentBackend(ABC):
     @abstractmethod
-    def __init__(self, path_or_stream: Iterable[Union[BytesIO, Path]]):
+    def __init__(self, path_or_stream: Union[BytesIO, Path]):
         pass
 
     @abstractmethod

diff --git a/docling/backend/docling_parse_backend.py b/docling/backend/docling_parse_backend.py
@@ -146,11 +146,12 @@ def unload(self):
 
 
 class DoclingParseDocumentBackend(PdfDocumentBackend):
-    def __init__(self, path_or_stream: Iterable[Union[BytesIO, Path]]):
+    def __init__(self, path_or_stream: Union[BytesIO, Path]):
         super().__init__(path_or_stream)
         self._pdoc = pdfium.PdfDocument(path_or_stream)
         # Parsing cells with docling_parser call
-        print("PARSING WITH DOCLING PARSE")
+        if isinstance(path_or_stream, BytesIO):
+            raise NotImplemented("This backend does not support byte streams yet.")
         parser = pdf_parser()
         self._parser_doc = parser.find_cells(str(path_or_stream))
 

diff --git a/docling/backend/pypdfium2_backend.py b/docling/backend/pypdfium2_backend.py
@@ -199,7 +199,7 @@ def unload(self):
 
 
 class PyPdfiumDocumentBackend(PdfDocumentBackend):
-    def __init__(self, path_or_stream: Iterable[Union[BytesIO, Path]]):
+    def __init__(self, path_or_stream: Union[BytesIO, Path]):
         super().__init__(path_or_stream)
         self._pdoc = pdfium.PdfDocument(path_or_stream)