natural-pdf 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

natural_pdf/collections/pdf_collection.py

@@ -1,39 +1,50 @@
-import os
+import copy  # Added for copying options
 import glob as py_glob
 import logging
-from typing import List, Optional, Dict, Any, Union, Iterable, Set, TYPE_CHECKING, Type
+import os
+import re  # Added for safe path generation
 from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Type, Union
+
 from PIL import Image
-import re # Added for safe path generation
-import copy # Added for copying options
 from tqdm import tqdm
 
 # Set up logger early
 logger = logging.getLogger(__name__)
 
 from natural_pdf.core.pdf import PDF
-from natural_pdf.elements.region import Region 
+from natural_pdf.elements.region import Region
 
 # --- Search Imports ---
 try:
     from natural_pdf.search.search_service_protocol import (
-         SearchServiceProtocol, SearchOptions, Indexable
-     )
+        Indexable,
+        SearchOptions,
+        SearchServiceProtocol,
+    )
     from natural_pdf.search.searchable_mixin import SearchableMixin
 except ImportError as e:
     logger_init = logging.getLogger(__name__)
-    logger_init.error(f"Failed to import search components. Search functionality disabled. Error: {e}", exc_info=True)
+    logger_init.warning(
+        f"Failed to import Haystack components. Semantic search functionality disabled.",
+    )
+
     # Dummy definitions
-    class SearchableMixin: pass
+    class SearchableMixin:
+        pass
+
     SearchServiceProtocol, SearchOptions, Indexable = object, object, object
 
-from natural_pdf.search.searchable_mixin import SearchableMixin # Import the new mixin
+from natural_pdf.search.searchable_mixin import SearchableMixin  # Import the new mixin
 
-class PDFCollection(SearchableMixin): # Inherit from the mixin
-    def __init__(self,
-                 source: Union[str, Iterable[Union[str, 'PDF']]],
-                 recursive: bool = True,
-                 **pdf_options: Any):
+
+class PDFCollection(SearchableMixin):  # Inherit from the mixin
+    def __init__(
+        self,
+        source: Union[str, Iterable[Union[str, "PDF"]]],
+        recursive: bool = True,
+        **pdf_options: Any,
+    ):
         """
         Initializes a collection of PDF documents from various sources.
 
@@ -46,27 +57,29 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
                        whether to search recursively (default: True).
             **pdf_options: Keyword arguments passed to the PDF constructor.
         """
-        self._pdfs: List['PDF'] = []
-        self._pdf_options = pdf_options # Store options for potential slicing later
-        self._recursive = recursive # Store setting for potential slicing
+        self._pdfs: List["PDF"] = []
+        self._pdf_options = pdf_options  # Store options for potential slicing later
+        self._recursive = recursive  # Store setting for potential slicing
 
         # Dynamically import PDF class within methods to avoid circular import at module load time
         PDF = self._get_pdf_class()
 
-        if hasattr(source, '__iter__') and not isinstance(source, str):
-             source_list = list(source)
-             if not source_list: return # Empty list source
-             if isinstance(source_list[0], PDF):
-                  if all(isinstance(item, PDF) for item in source_list):
-                       self._pdfs = source_list # Direct assignment
-                       # Don't adopt search context anymore
-                       return
-                  else: raise TypeError("Iterable source has mixed PDF/non-PDF objects.")
-             # If it's an iterable but not PDFs, fall through to resolve sources
+        if hasattr(source, "__iter__") and not isinstance(source, str):
+            source_list = list(source)
+            if not source_list:
+                return  # Empty list source
+            if isinstance(source_list[0], PDF):
+                if all(isinstance(item, PDF) for item in source_list):
+                    self._pdfs = source_list  # Direct assignment
+                    # Don't adopt search context anymore
+                    return
+                else:
+                    raise TypeError("Iterable source has mixed PDF/non-PDF objects.")
+            # If it's an iterable but not PDFs, fall through to resolve sources
 
         # Resolve string, iterable of strings, or single string source to paths/URLs
         resolved_paths_or_urls = self._resolve_sources_to_paths(source)
-        self._initialize_pdfs(resolved_paths_or_urls, PDF) # Pass PDF class
+        self._initialize_pdfs(resolved_paths_or_urls, PDF)  # Pass PDF class
 
         self._iter_index = 0
 
@@ -79,15 +92,21 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
         try:
             # Import needs to resolve path correctly
             from natural_pdf.core.pdf import PDF
+
             return PDF
         except ImportError as e:
-            logger.error("Could not import PDF class from natural_pdf.core.pdf. Ensure it exists and there are no circular imports at runtime.")
+            logger.error(
+                "Could not import PDF class from natural_pdf.core.pdf. Ensure it exists and there are no circular imports at runtime."
+            )
             raise ImportError("PDF class is required but could not be imported.") from e
 
     # --- Internal Helpers ---
 
-    def _is_url(self, s: str) -> bool: return s.startswith(('http://', 'https://'))
-    def _has_glob_magic(self, s: str) -> bool: return py_glob.has_magic(s)
+    def _is_url(self, s: str) -> bool:
+        return s.startswith(("http://", "https://"))
+
+    def _has_glob_magic(self, s: str) -> bool:
+        return py_glob.has_magic(s)
 
     def _execute_glob(self, pattern: str) -> Set[str]:
         """Glob for paths and return a set of valid PDF paths."""
@@ -96,10 +115,10 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
             # Use iglob for potentially large directories/matches
             paths_iter = py_glob.iglob(pattern, recursive=self._recursive)
             for path_str in paths_iter:
-                 # Use Path object for easier checking
-                 p = Path(path_str)
-                 if p.is_file() and p.suffix.lower() == ".pdf":
-                      found_paths.add(str(p.resolve())) # Store resolved absolute path
+                # Use Path object for easier checking
+                p = Path(path_str)
+                if p.is_file() and p.suffix.lower() == ".pdf":
+                    found_paths.add(str(p.resolve()))  # Store resolved absolute path
         except Exception as e:
             logger.error(f"Error processing glob pattern '{pattern}': {e}")
         return found_paths
@@ -111,33 +130,37 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
 
         if isinstance(source, str):
             sources_to_process.append(source)
-        elif hasattr(source, '__iter__'):
+        elif hasattr(source, "__iter__"):
             sources_to_process.extend(list(source))
-        else: # Should not happen based on __init__ checks, but safeguard
-             raise TypeError(f"Unexpected source type in _resolve_sources_to_paths: {type(source)}")
+        else:  # Should not happen based on __init__ checks, but safeguard
+            raise TypeError(f"Unexpected source type in _resolve_sources_to_paths: {type(source)}")
 
         for item in sources_to_process:
-             if not isinstance(item, str):
-                  logger.warning(f"Skipping non-string item in source list: {type(item)}")
-                  continue
-             
-             item_path = Path(item)
-             
-             if self._is_url(item):
-                 final_paths.add(item) # Add URL directly
-             elif self._has_glob_magic(item):
-                 glob_results = self._execute_glob(item)
-                 final_paths.update(glob_results)
-             elif item_path.is_dir():
-                 # Use glob to find PDFs in directory, respecting recursive flag
-                 dir_pattern = str(item_path / "**" / "*.pdf") if self._recursive else str(item_path / "*.pdf")
-                 dir_glob_results = self._execute_glob(dir_pattern)
-                 final_paths.update(dir_glob_results)
-             elif item_path.is_file() and item_path.suffix.lower() == ".pdf":
-                 final_paths.add(str(item_path.resolve())) # Add resolved file path
-             else:
-                 logger.warning(f"Source item ignored (not a valid URL, directory, file, or glob): {item}")
-                 
+            if not isinstance(item, str):
+                logger.warning(f"Skipping non-string item in source list: {type(item)}")
+                continue
+
+            item_path = Path(item)
+
+            if self._is_url(item):
+                final_paths.add(item)  # Add URL directly
+            elif self._has_glob_magic(item):
+                glob_results = self._execute_glob(item)
+                final_paths.update(glob_results)
+            elif item_path.is_dir():
+                # Use glob to find PDFs in directory, respecting recursive flag
+                dir_pattern = (
+                    str(item_path / "**" / "*.pdf") if self._recursive else str(item_path / "*.pdf")
+                )
+                dir_glob_results = self._execute_glob(dir_pattern)
+                final_paths.update(dir_glob_results)
+            elif item_path.is_file() and item_path.suffix.lower() == ".pdf":
+                final_paths.add(str(item_path.resolve()))  # Add resolved file path
+            else:
+                logger.warning(
+                    f"Source item ignored (not a valid URL, directory, file, or glob): {item}"
+                )
+
         return sorted(list(final_paths))
 
     def _initialize_pdfs(self, paths_or_urls: List[str], PDF_cls: Type):
@@ -149,32 +172,38 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
                 pdf_instance = PDF_cls(path_or_url, **self._pdf_options)
                 self._pdfs.append(pdf_instance)
             except Exception as e:
-                 logger.error(f"Failed to load PDF: {path_or_url}. Error: {e}", exc_info=False) # Keep log concise
-                 failed_count += 1
+                logger.error(
+                    f"Failed to load PDF: {path_or_url}. Error: {e}", exc_info=False
+                )  # Keep log concise
+                failed_count += 1
         logger.info(f"Successfully initialized {len(self._pdfs)} PDFs. Failed: {failed_count}")
 
     # --- Public Factory Class Methods (Simplified) ---
 
     @classmethod
-    def from_paths(cls, paths_or_urls: List[str], **pdf_options: Any) -> 'PDFCollection':
+    def from_paths(cls, paths_or_urls: List[str], **pdf_options: Any) -> "PDFCollection":
         """Creates a PDFCollection explicitly from a list of file paths or URLs."""
         # __init__ can handle List[str] directly now
         return cls(paths_or_urls, **pdf_options)
 
     @classmethod
-    def from_glob(cls, pattern: str, recursive: bool = True, **pdf_options: Any) -> 'PDFCollection':
+    def from_glob(cls, pattern: str, recursive: bool = True, **pdf_options: Any) -> "PDFCollection":
         """Creates a PDFCollection explicitly from a single glob pattern."""
         # __init__ can handle single glob string directly
         return cls(pattern, recursive=recursive, **pdf_options)
 
     @classmethod
-    def from_globs(cls, patterns: List[str], recursive: bool = True, **pdf_options: Any) -> 'PDFCollection':
+    def from_globs(
+        cls, patterns: List[str], recursive: bool = True, **pdf_options: Any
+    ) -> "PDFCollection":
         """Creates a PDFCollection explicitly from a list of glob patterns."""
-         # __init__ can handle List[str] containing globs directly
+        # __init__ can handle List[str] containing globs directly
         return cls(patterns, recursive=recursive, **pdf_options)
 
     @classmethod
-    def from_directory(cls, directory_path: str, recursive: bool = True, **pdf_options: Any) -> 'PDFCollection':
+    def from_directory(
+        cls, directory_path: str, recursive: bool = True, **pdf_options: Any
+    ) -> "PDFCollection":
         """Creates a PDFCollection explicitly from PDF files within a directory."""
         # __init__ can handle single directory string directly
         return cls(directory_path, recursive=recursive, **pdf_options)
@@ -183,12 +212,12 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
     def __len__(self) -> int:
         return len(self._pdfs)
 
-    def __getitem__(self, key) -> Union['PDF', 'PDFCollection']:
+    def __getitem__(self, key) -> Union["PDF", "PDFCollection"]:
         # Use dynamic import here as well
         PDF = self._get_pdf_class()
         if isinstance(key, slice):
             # Create a new collection with the sliced PDFs and original options
-            new_collection = PDFCollection.__new__(PDFCollection) # Create blank instance
+            new_collection = PDFCollection.__new__(PDFCollection)  # Create blank instance
             new_collection._pdfs = self._pdfs[key]
             new_collection._pdf_options = self._pdf_options
             new_collection._recursive = self._recursive
@@ -199,9 +228,9 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
             if 0 <= key < len(self._pdfs):
                 return self._pdfs[key]
             else:
-                 raise IndexError(f"PDF index {key} out of range (0-{len(self._pdfs)-1}).")
+                raise IndexError(f"PDF index {key} out of range (0-{len(self._pdfs)-1}).")
         else:
-             raise TypeError(f"PDF indices must be integers or slices, not {type(key)}.")
+            raise TypeError(f"PDF indices must be integers or slices, not {type(key)}.")
 
     def __iter__(self):
         return iter(self._pdfs)
@@ -211,24 +240,23 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
         return f"<PDFCollection(count={len(self)})>"
 
     @property
-    def pdfs(self) -> List['PDF']:
-         """Returns the list of PDF objects held by the collection."""
-         return self._pdfs
+    def pdfs(self) -> List["PDF"]:
+        """Returns the list of PDF objects held by the collection."""
+        return self._pdfs
 
-    # --- Other Methods (e.g., apply_ocr - could leverage service in future?) ---
     def apply_ocr(self, *args, **kwargs):
         PDF = self._get_pdf_class()
         # Delegate to individual PDF objects
         logger.info("Applying OCR to relevant PDFs in collection...")
         results = []
         for pdf in self._pdfs:
-             # We need to figure out which pages belong to which PDF if batching here
-             # For now, simpler to call on each PDF
-             try:
-                 # Assume apply_ocr exists on PDF and accepts similar args
-                 pdf.apply_ocr(*args, **kwargs)
-             except Exception as e:
-                 logger.error(f"Failed applying OCR to {pdf.path}: {e}", exc_info=True)
+            # We need to figure out which pages belong to which PDF if batching here
+            # For now, simpler to call on each PDF
+            try:
+                # Assume apply_ocr exists on PDF and accepts similar args
+                pdf.apply_ocr(*args, **kwargs)
+            except Exception as e:
+                logger.error(f"Failed applying OCR to {pdf.path}: {e}", exc_info=True)
         return self
 
     # --- Advanced Method Placeholders ---
@@ -237,23 +265,44 @@ class PDFCollection(SearchableMixin): # Inherit from the mixin
     def categorize(self, categories: List[str], **kwargs):
         """Categorizes PDFs in the collection based on content or features."""
         # Implementation requires integrating with classification models or logic
-        raise NotImplementedError("categorize requires classification implementation.") 
+        raise NotImplementedError("categorize requires classification implementation.")
+
+    def export_ocr_correction_task(self, output_zip_path: str, **kwargs):
+        """
+        Exports OCR results from all PDFs in this collection into a single
+        correction task package (zip file).
+
+        Args:
+            output_zip_path: The path to save the output zip file.
+            **kwargs: Additional arguments passed to create_correction_task_package
+                      (e.g., image_render_scale, overwrite).
+        """
+        try:
+            from natural_pdf.utils.packaging import create_correction_task_package
+            # Pass the collection itself (self) as the source
+            create_correction_task_package(source=self, output_zip_path=output_zip_path, **kwargs)
+        except ImportError:
+            logger.error("Failed to import 'create_correction_task_package'. Packaging utility might be missing.")
+            # Or raise
+        except Exception as e:
+            logger.error(f"Failed to export correction task for collection: {e}", exc_info=True)
+            raise # Re-raise the exception from the utility function
 
-    # --- Mixin Required Implementation --- 
+    # --- Mixin Required Implementation ---
     def get_indexable_items(self) -> Iterable[Indexable]:
         """Yields Page objects from the collection, conforming to Indexable."""
         if not self._pdfs:
-             return # Return empty iterator if no PDFs
-             
+            return  # Return empty iterator if no PDFs
+
         for pdf in self._pdfs:
-             if not pdf.pages: # Handle case where a PDF might have 0 pages after loading
-                 logger.warning(f"PDF '{pdf.path}' has no pages. Skipping.")
-                 continue
-             for page in pdf.pages:
-                 # Optional: Add filtering here if needed (e.g., skip empty pages)
-                 # Assuming Page object conforms to Indexable
-                 # We might still want the empty page check here for efficiency
-                 # if not page.extract_text(use_exclusions=False).strip():
-                 #     logger.debug(f"Skipping empty page {page.page_number} from PDF '{pdf.path}'.")
-                 #     continue
-                 yield page 
+            if not pdf.pages:  # Handle case where a PDF might have 0 pages after loading
+                logger.warning(f"PDF '{pdf.path}' has no pages. Skipping.")
+                continue
+            for page in pdf.pages:
+                # Optional: Add filtering here if needed (e.g., skip empty pages)
+                # Assuming Page object conforms to Indexable
+                # We might still want the empty page check here for efficiency
+                # if not page.extract_text(use_exclusions=False).strip():
+                #     logger.debug(f"Skipping empty page {page.page_number} from PDF '{pdf.path}'.")
+                #     continue
+                yield page

natural_pdf/core/__init__.py

@@ -1,3 +1,3 @@
 """
 Core classes for Natural PDF.
-"""
+"""