natural-pdf 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

natural_pdf/collections/mixins.py

@@ -0,0 +1,111 @@
+import logging
+from typing import Any, Callable, Iterable, TypeVar
+
+from tqdm.auto import tqdm
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")  # Generic type for items in the collection
+
+
+class DirectionalCollectionMixin:
+    """
+    Mixin providing directional methods for collections of elements/regions.
+    """
+
+    def below(self, **kwargs) -> "ElementCollection":
+        """Find regions below all elements in this collection."""
+        return self.apply(lambda element: element.below(**kwargs))
+
+    def above(self, **kwargs) -> "ElementCollection":
+        """Find regions above all elements in this collection."""
+        return self.apply(lambda element: element.above(**kwargs))
+
+    def left(self, **kwargs) -> "ElementCollection":
+        """Find regions to the left of all elements in this collection."""
+        return self.apply(lambda element: element.left(**kwargs))
+
+    def right(self, **kwargs) -> "ElementCollection":
+        """Find regions to the right of all elements in this collection."""
+        return self.apply(lambda element: element.right(**kwargs))
+
+    def expand(self, **kwargs) -> "ElementCollection":
+        """Expand all elements in this collection."""
+        return self.apply(lambda element: element.expand(**kwargs))
+
+
+class ApplyMixin:
+    """
+    Mixin class providing an `.apply()` method for collections.
+
+    Assumes the inheriting class implements `__iter__` and `__len__` appropriately
+    for the items to be processed by `apply`.
+    """
+
+    def _get_items_for_apply(self) -> Iterable[Any]:
+        """
+        Returns the iterable of items to apply the function to.
+        Defaults to iterating over `self`. Subclasses should override this
+        if the default iteration is not suitable for the apply operation.
+        """
+        # Default to standard iteration over the collection itself
+        return iter(self)
+
+    def apply(self: Any, func: Callable[[Any, ...], Any], *args, **kwargs) -> Iterable[Any]:
+        """
+        Applies a function to each item in the collection.
+
+        Args:
+            func: The function to apply to each item. The item itself
+                  will be passed as the first argument to the function.
+            *args: Additional positional arguments to pass to func.
+            **kwargs: Additional keyword arguments to pass to func.
+                      A special keyword argument 'show_progress' (bool, default=False)
+                      can be used to display a progress bar.
+        """
+        show_progress = kwargs.pop("show_progress", False)
+        # Derive unit name from class name
+        unit_name = self.__class__.__name__.lower()
+        items_iterable = self._get_items_for_apply()
+
+        # Need total count for tqdm, assumes __len__ is implemented by the inheriting class
+        total_items = 0
+        try:
+            total_items = len(self)
+        except TypeError:  # Handle cases where __len__ might not be defined on self
+            logger.warning(f"Could not determine collection length for progress bar.")
+
+        if show_progress and total_items > 0:
+            items_iterable = tqdm(
+                items_iterable, total=total_items, desc=f"Applying {func.__name__}", unit=unit_name
+            )
+        elif show_progress:
+            logger.info(
+                f"Applying {func.__name__} (progress bar disabled for zero/unknown length)."
+            )
+
+        results = [func(item, *args, **kwargs) for item in items_iterable]
+
+        # If results is empty, return an empty list
+        if not results:
+            return []
+
+        # Import here to avoid circular imports
+        from natural_pdf import PDF, Page
+        from natural_pdf.collections.pdf_collection import PDFCollection
+        from natural_pdf.elements.base import Element
+        from natural_pdf.elements.collections import ElementCollection, PageCollection
+        from natural_pdf.elements.region import Region
+
+        first_non_none = next((r for r in results if r is not None), None)
+        first_type = type(first_non_none) if first_non_none is not None else None
+
+        # Return the appropriate collection based on result type (...generally)
+        if issubclass(first_type, Element) or issubclass(first_type, Region):
+            return ElementCollection(results)
+        elif first_type == PDF:
+            return PDFCollection(results)
+        elif first_type == Page:
+            return PageCollection(results)
+
+        return results

natural_pdf/collections/pdf_collection.py

@@ -1,19 +1,47 @@
+import concurrent.futures  # Import concurrent.futures
 import copy  # Added for copying options
 import glob as py_glob
 import logging
 import os
 import re  # Added for safe path generation
+import threading  # Import threading for logging thread information
+import time  # Import time for logging timestamps
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
 
 from PIL import Image
 from tqdm import tqdm
+from tqdm.auto import tqdm as auto_tqdm
+from tqdm.notebook import tqdm as notebook_tqdm
+
+from natural_pdf.utils.tqdm_utils import get_tqdm
+
+# Get the appropriate tqdm class once
+tqdm = get_tqdm()
 
 # Set up logger early
+# Configure logging to include thread information
+# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(threadName)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 from natural_pdf.core.pdf import PDF
 from natural_pdf.elements.region import Region
+from natural_pdf.export.mixin import ExportMixin
 
 # --- Search Imports ---
 try:
@@ -35,10 +63,12 @@ except ImportError as e:
 
     SearchServiceProtocol, SearchOptions, Indexable = object, object, object
 
+# Import the ApplyMixin
+from natural_pdf.collections.mixins import ApplyMixin
 from natural_pdf.search.searchable_mixin import SearchableMixin  # Import the new mixin
 
 
-class PDFCollection(SearchableMixin):  # Inherit from the mixin
+class PDFCollection(SearchableMixin, ApplyMixin, ExportMixin):  # Add ExportMixin
     def __init__(
         self,
         source: Union[str, Iterable[Union[str, "PDF"]]],
@@ -237,30 +267,257 @@ class PDFCollection(SearchableMixin):  # Inherit from the mixin
 
     def __repr__(self) -> str:
         # Removed search status
-        return f"<PDFCollection(count={len(self)})>"
+        return f"<PDFCollection(count={len(self._pdfs)})>"
+        return f"<PDFCollection(count={len(self._pdfs)})>"
 
     @property
     def pdfs(self) -> List["PDF"]:
         """Returns the list of PDF objects held by the collection."""
         return self._pdfs
 
-    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 = []
+    @overload
+    def find_all(
+        self,
+        *,
+        text: str,
+        apply_exclusions: bool = True,
+        regex: bool = False,
+        case: bool = True,
+        **kwargs,
+    ) -> "ElementCollection": ...
+
+    @overload
+    def find_all(
+        self,
+        selector: str,
+        *,
+        apply_exclusions: bool = True,
+        regex: bool = False,
+        case: bool = True,
+        **kwargs,
+    ) -> "ElementCollection": ...
+
+    def find_all(
+        self,
+        selector: Optional[str] = None,  # Now optional
+        *,
+        text: Optional[str] = None,  # New text parameter
+        apply_exclusions: bool = True,
+        regex: bool = False,
+        case: bool = True,
+        **kwargs,
+    ) -> "ElementCollection":
+        """
+        Find all elements matching the selector OR text across all PDFs in the collection.
+
+        Provide EITHER `selector` OR `text`, but not both.
+
+        This creates an ElementCollection that can span multiple PDFs. Note that
+        some ElementCollection methods have limitations when spanning PDFs.
+
+        Args:
+            selector: CSS-like selector string to query elements.
+            text: Text content to search for (equivalent to 'text:contains(...)').
+            apply_exclusions: Whether to exclude elements in exclusion regions (default: True).
+            regex: Whether to use regex for text search (`selector` or `text`) (default: False).
+            case: Whether to do case-sensitive text search (`selector` or `text`) (default: True).
+            **kwargs: Additional keyword arguments passed to the find_all method of each PDF.
+
+        Returns:
+            ElementCollection containing all matching elements across all PDFs.
+        """
+        # Validation happens within pdf.find_all
+
+        # Collect elements from all PDFs
+        all_elements = []
         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)
+                # Pass the relevant arguments down to each PDF's find_all
+                elements = pdf.find_all(
+                    selector=selector,
+                    text=text,
+                    apply_exclusions=apply_exclusions,
+                    regex=regex,
+                    case=case,
+                    **kwargs,
+                )
+                all_elements.extend(elements.elements)
+            except Exception as e:
+                logger.error(f"Error finding elements in {pdf.path}: {e}", exc_info=True)
+
+        return ElementCollection(all_elements)
+
+    def apply_ocr(
+        self,
+        engine: Optional[str] = None,
+        languages: Optional[List[str]] = None,
+        min_confidence: Optional[float] = None,
+        device: Optional[str] = None,
+        resolution: Optional[int] = None,
+        apply_exclusions: bool = True,
+        detect_only: bool = False,
+        replace: bool = True,
+        options: Optional[Any] = None,
+        pages: Optional[Union[slice, List[int]]] = None,
+        max_workers: Optional[int] = None,
+    ) -> "PDFCollection":
+        """
+        Apply OCR to all PDFs in the collection, potentially in parallel.
+
+        Args:
+            engine: OCR engine to use (e.g., 'easyocr', 'paddleocr', 'surya')
+            languages: List of language codes for OCR
+            min_confidence: Minimum confidence threshold for text detection
+            device: Device to use for OCR (e.g., 'cpu', 'cuda')
+            resolution: DPI resolution for page rendering
+            apply_exclusions: Whether to apply exclusion regions
+            detect_only: If True, only detect text regions without extracting text
+            replace: If True, replace existing OCR elements
+            options: Engine-specific options
+            pages: Specific pages to process (None for all pages)
+            max_workers: Maximum number of threads to process PDFs concurrently.
+                         If None or 1, processing is sequential. (default: None)
+
+        Returns:
+            Self for method chaining
+        """
+        PDF = self._get_pdf_class()
+        logger.info(
+            f"Applying OCR to {len(self._pdfs)} PDFs in collection (max_workers={max_workers})..."
+        )
+
+        # Worker function takes PDF object again
+        def _process_pdf(pdf: PDF):
+            """Helper function to apply OCR to a single PDF, handling errors."""
+            thread_id = threading.current_thread().name  # Get thread name for logging
+            pdf_path = pdf.path  # Get path for logging
+            logger.debug(f"[{thread_id}] Starting OCR process for: {pdf_path}")
+            start_time = time.monotonic()
+            try:
+                pdf.apply_ocr(  # Call apply_ocr on the original PDF object
+                    pages=pages,
+                    engine=engine,
+                    languages=languages,
+                    min_confidence=min_confidence,
+                    device=device,
+                    resolution=resolution,
+                    apply_exclusions=apply_exclusions,
+                    detect_only=detect_only,
+                    replace=replace,
+                    options=options,
+                    # Note: We might want a max_workers here too for page rendering?
+                    # For now, PDF.apply_ocr doesn't have it.
+                )
+                end_time = time.monotonic()
+                logger.debug(
+                    f"[{thread_id}] Finished OCR process for: {pdf_path} (Duration: {end_time - start_time:.2f}s)"
+                )
+                return pdf_path, None
             except Exception as e:
-                logger.error(f"Failed applying OCR to {pdf.path}: {e}", exc_info=True)
+                end_time = time.monotonic()
+                logger.error(
+                    f"[{thread_id}] Failed OCR process for {pdf_path} after {end_time - start_time:.2f}s: {e}",
+                    exc_info=False,
+                )
+                return pdf_path, e  # Return path and error
+
+        # Use ThreadPoolExecutor for parallel processing if max_workers > 1
+        if max_workers is not None and max_workers > 1:
+            futures = []
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=max_workers, thread_name_prefix="OCRWorker"
+            ) as executor:
+                for pdf in self._pdfs:
+                    # Submit the PDF object to the worker function
+                    futures.append(executor.submit(_process_pdf, pdf))
+
+            # Use the selected tqdm class with as_completed for progress tracking
+            progress_bar = tqdm(
+                concurrent.futures.as_completed(futures),
+                total=len(self._pdfs),
+                desc="Applying OCR (Parallel)",
+                unit="pdf",
+            )
+
+            for future in progress_bar:
+                pdf_path, error = future.result()  # Get result (or exception)
+                if error:
+                    progress_bar.set_postfix_str(f"Error: {pdf_path}", refresh=True)
+                # Progress is updated automatically by tqdm
+
+        else:  # Sequential processing (max_workers is None or 1)
+            logger.info("Applying OCR sequentially...")
+            # Use the selected tqdm class for sequential too for consistency
+            # Iterate over PDF objects directly for sequential
+            for pdf in tqdm(self._pdfs, desc="Applying OCR (Sequential)", unit="pdf"):
+                _process_pdf(pdf)  # Call helper directly with PDF object
+
+        logger.info("Finished applying OCR across the collection.")
         return self
 
-    # --- Advanced Method Placeholders ---
-    # Placeholder for categorize removed as find_relevant is now implemented
+    def correct_ocr(
+        self,
+        correction_callback: Callable[[Any], Optional[str]],
+        max_workers: Optional[int] = None,
+        progress_callback: Optional[Callable[[], None]] = None,
+    ) -> "PDFCollection":
+        """
+        Apply OCR correction to all relevant elements across all pages and PDFs
+        in the collection using a single progress bar.
+
+        Args:
+            correction_callback: Function to apply to each OCR element.
+                                 It receives the element and should return
+                                 the corrected text (str) or None.
+            max_workers: Max threads to use for parallel execution within each page.
+            progress_callback: Optional callback function to call after processing each element.
+
+        Returns:
+            Self for method chaining.
+        """
+        PDF = self._get_pdf_class()  # Ensure PDF class is available
+        if not callable(correction_callback):
+            raise TypeError("`correction_callback` must be a callable function.")
+
+        logger.info(f"Gathering OCR elements from {len(self._pdfs)} PDFs for correction...")
+
+        # 1. Gather all target elements using the collection's find_all
+        #    Crucially, set apply_exclusions=False to include elements in headers/footers etc.
+        all_ocr_elements = self.find_all("text[source=ocr]", apply_exclusions=False).elements
+
+        if not all_ocr_elements:
+            logger.info("No OCR elements found in the collection to correct.")
+            return self
+
+        total_elements = len(all_ocr_elements)
+        logger.info(
+            f"Found {total_elements} OCR elements across the collection. Starting correction process..."
+        )
+
+        # 2. Initialize the progress bar
+        progress_bar = tqdm(total=total_elements, desc="Correcting OCR Elements", unit="element")
+
+        # 3. Iterate through PDFs and delegate to PDF.correct_ocr
+        #    PDF.correct_ocr handles page iteration and passing the progress callback down.
+        for pdf in self._pdfs:
+            if not pdf.pages:
+                continue
+            try:
+                pdf.correct_ocr(
+                    correction_callback=correction_callback,
+                    max_workers=max_workers,
+                    progress_callback=progress_bar.update,  # Pass the bar's update method
+                )
+            except Exception as e:
+                logger.error(
+                    f"Error occurred during correction process for PDF {pdf.path}: {e}",
+                    exc_info=True,
+                )
+                # Decide if we should stop or continue? For now, continue.
+
+        progress_bar.close()
+
+        return self
 
     def categorize(self, categories: List[str], **kwargs):
         """Categorizes PDFs in the collection based on content or features."""
@@ -309,3 +566,165 @@ class PDFCollection(SearchableMixin):  # Inherit from the mixin
                 #     logger.debug(f"Skipping empty page {page.page_number} from PDF '{pdf.path}'.")
                 #     continue
                 yield page
+
+    # --- Classification Method --- #
+    def classify_all(
+        self,
+        categories: List[str],
+        model: str = "text",
+        max_workers: Optional[int] = None,
+        **kwargs,
+    ) -> "PDFCollection":
+        """
+        Classify all pages across all PDFs in the collection, potentially in parallel.
+
+        This method uses the unified `classify_all` approach, delegating page
+        classification to each PDF's `classify_pages` method.
+        It displays a progress bar tracking individual pages.
+
+        Args:
+            categories: A list of string category names.
+            model: Model identifier ('text', 'vision', or specific HF ID).
+            max_workers: Maximum number of threads to process PDFs concurrently.
+                         If None or 1, processing is sequential.
+            **kwargs: Additional arguments passed down to `pdf.classify_pages` and
+                      subsequently to `page.classify` (e.g., device,
+                      confidence_threshold, resolution).
+
+        Returns:
+            Self for method chaining.
+
+        Raises:
+            ValueError: If categories list is empty.
+            ClassificationError: If classification fails for any page (will stop processing).
+            ImportError: If classification dependencies are missing.
+        """
+        PDF = self._get_pdf_class()
+        if not categories:
+            raise ValueError("Categories list cannot be empty.")
+
+        logger.info(
+            f"Starting classification for {len(self._pdfs)} PDFs in collection (model: '{model}')..."
+        )
+
+        # Calculate total pages for the progress bar
+        total_pages = sum(len(pdf.pages) for pdf in self._pdfs if pdf.pages)
+        if total_pages == 0:
+            logger.warning("No pages found in the PDF collection to classify.")
+            return self
+
+        progress_bar = tqdm(
+            total=total_pages, desc=f"Classifying Pages (model: {model})", unit="page"
+        )
+
+        # Worker function
+        def _process_pdf_classification(pdf: PDF):
+            thread_id = threading.current_thread().name
+            pdf_path = pdf.path
+            logger.debug(f"[{thread_id}] Starting classification process for: {pdf_path}")
+            start_time = time.monotonic()
+            try:
+                # Call classify_pages on the PDF, passing the progress callback
+                pdf.classify_pages(
+                    categories=categories,
+                    model=model,
+                    progress_callback=progress_bar.update,
+                    **kwargs,
+                )
+                end_time = time.monotonic()
+                logger.debug(
+                    f"[{thread_id}] Finished classification for: {pdf_path} (Duration: {end_time - start_time:.2f}s)"
+                )
+                return pdf_path, None  # Return path and no error
+            except Exception as e:
+                end_time = time.monotonic()
+                # Error is logged within classify_pages, but log summary here
+                logger.error(
+                    f"[{thread_id}] Failed classification process for {pdf_path} after {end_time - start_time:.2f}s: {e}",
+                    exc_info=False,
+                )
+                # Close progress bar immediately on error to avoid hanging
+                progress_bar.close()
+                # Re-raise the exception to stop the entire collection processing
+                raise
+
+        # Use ThreadPoolExecutor for parallel processing if max_workers > 1
+        try:
+            if max_workers is not None and max_workers > 1:
+                logger.info(f"Classifying PDFs in parallel with {max_workers} workers.")
+                futures = []
+                with concurrent.futures.ThreadPoolExecutor(
+                    max_workers=max_workers, thread_name_prefix="ClassifyWorker"
+                ) as executor:
+                    for pdf in self._pdfs:
+                        futures.append(executor.submit(_process_pdf_classification, pdf))
+
+                    # Wait for all futures to complete (progress updated by callback)
+                    # Exceptions are raised by future.result() if worker failed
+                    for future in concurrent.futures.as_completed(futures):
+                        future.result()  # Raise exception if worker failed
+
+            else:  # Sequential processing
+                logger.info("Classifying PDFs sequentially.")
+                for pdf in self._pdfs:
+                    _process_pdf_classification(pdf)
+
+            logger.info("Finished classification across the collection.")
+
+        finally:
+            # Ensure progress bar is closed even if errors occurred elsewhere
+            if not progress_bar.disable and progress_bar.n < progress_bar.total:
+                progress_bar.close()
+            elif progress_bar.disable is False:
+                progress_bar.close()
+
+        return self
+
+    # --- End Classification Method --- #
+
+    def _gather_analysis_data(
+        self,
+        analysis_keys: List[str],
+        include_content: bool,
+        include_images: bool,
+        image_dir: Optional[Path],
+        image_format: str,
+        image_resolution: int,
+    ) -> List[Dict[str, Any]]:
+        """
+        Gather analysis data from all PDFs in the collection.
+
+        Args:
+            analysis_keys: Keys in the analyses dictionary to export
+            include_content: Whether to include extracted text
+            include_images: Whether to export images
+            image_dir: Directory to save images
+            image_format: Format to save images
+            image_resolution: Resolution for exported images
+
+        Returns:
+            List of dictionaries containing analysis data
+        """
+        if not self._pdfs:
+            logger.warning("No PDFs found in collection")
+            return []
+
+        all_data = []
+
+        for pdf in tqdm(self._pdfs, desc="Gathering PDF data", leave=False):
+            # PDF level data
+            pdf_data = {
+                "pdf_path": pdf.path,
+                "pdf_filename": Path(pdf.path).name,
+                "total_pages": len(pdf.pages) if hasattr(pdf, "pages") else 0,
+            }
+
+            # Add metadata if available
+            if hasattr(pdf, "metadata") and pdf.metadata:
+                for k, v in pdf.metadata.items():
+                    if v:  # Only add non-empty metadata
+                        pdf_data[f"metadata.{k}"] = str(v)
+
+            all_data.append(pdf_data)
+
+        return all_data

natural_pdf/core/element_manager.py

@@ -539,3 +539,86 @@ class ElementManager:
         """Get all region elements."""
         self.load_elements()
         return self._elements.get("regions", [])
+
+    def remove_ocr_elements(self):
+        """
+        Remove all elements with source="ocr" from the elements dictionary.
+        This should be called before adding new OCR elements if replacement is desired.
+
+        Returns:
+            int: Number of OCR elements removed
+        """
+        # Load elements if not already loaded
+        self.load_elements()
+
+        removed_count = 0
+
+        # Filter out OCR elements from words
+        if "words" in self._elements:
+            original_len = len(self._elements["words"])
+            self._elements["words"] = [
+                word for word in self._elements["words"] if getattr(word, "source", None) != "ocr"
+            ]
+            removed_count += original_len - len(self._elements["words"])
+
+        # Filter out OCR elements from chars
+        if "chars" in self._elements:
+            original_len = len(self._elements["chars"])
+            self._elements["chars"] = [
+                char
+                for char in self._elements["chars"]
+                if (isinstance(char, dict) and char.get("source") != "ocr")
+                or (not isinstance(char, dict) and getattr(char, "source", None) != "ocr")
+            ]
+            removed_count += original_len - len(self._elements["chars"])
+
+        logger.info(f"Page {self._page.number}: Removed {removed_count} OCR elements.")
+        return removed_count
+
+    def remove_element(self, element, element_type="words"):
+        """
+        Remove a specific element from the managed elements.
+
+        Args:
+            element: The element to remove
+            element_type: The type of element ('words', 'chars', etc.)
+
+        Returns:
+            bool: True if removed successfully, False otherwise
+        """
+        # Load elements if not already loaded
+        self.load_elements()
+
+        # Check if the collection exists
+        if element_type not in self._elements:
+            logger.warning(f"Cannot remove element: collection '{element_type}' does not exist")
+            return False
+
+        # Try to remove the element
+        try:
+            if element in self._elements[element_type]:
+                self._elements[element_type].remove(element)
+                logger.debug(f"Removed element from {element_type}: {element}")
+                return True
+            else:
+                logger.debug(f"Element not found in {element_type}: {element}")
+                return False
+        except Exception as e:
+            logger.error(f"Error removing element from {element_type}: {e}", exc_info=True)
+            return False
+
+    def has_elements(self) -> bool:
+        """
+        Check if any significant elements (words, rects, lines, regions)
+        have been loaded or added.
+
+        Returns:
+            True if any elements exist, False otherwise.
+        """
+        self.load_elements()
+
+        for key in ["words", "rects", "lines", "regions"]:
+            if self._elements.get(key):
+                return True
+
+        return False