natural-pdf 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

natural_pdf/core/page.py

@@ -1138,31 +1138,145 @@ class Page(ClassificationMixin, ExtractionMixin, ShapeDetectionMixin):
         logger.debug(f"Page {self.number}: extract_text finished, result length: {len(result)}.")
         return result
 
-    def extract_table(self, table_settings={}) -> List[Any]:
+    def extract_table(
+        self,
+        method: Optional[str] = None,
+        table_settings: Optional[dict] = None,
+        use_ocr: bool = False,
+        ocr_config: Optional[dict] = None,
+        text_options: Optional[Dict] = None,
+        cell_extraction_func: Optional[Callable[["Region"], Optional[str]]] = None,
+        show_progress: bool = False,
+    ) -> List[List[Optional[str]]]:
         """
-        Extract the largest table from this page.
+        Extract the largest table from this page using enhanced region-based extraction.
 
         Args:
-            table_settings: Additional extraction parameters
+            method: Method to use: 'tatr', 'pdfplumber', 'text', 'stream', 'lattice', or None (auto-detect).
+            table_settings: Settings for pdfplumber table extraction.
+            use_ocr: Whether to use OCR for text extraction (currently only applicable with 'tatr' method).
+            ocr_config: OCR configuration parameters.
+            text_options: Dictionary of options for the 'text' method.
+            cell_extraction_func: Optional callable function that takes a cell Region object
+                                  and returns its string content. For 'text' method only.
+            show_progress: If True, display a progress bar during cell text extraction for the 'text' method.
 
         Returns:
-            List of extracted tables (or None if no table found)
+            Table data as a list of rows, where each row is a list of cell values (str or None).
         """
-        # pdfplumber returns None if no table found
-        return self._page.extract_table(table_settings)
+        # Create a full-page region and delegate to its enhanced extract_table method
+        page_region = self.create_region(0, 0, self.width, self.height)
+        return page_region.extract_table(
+            method=method,
+            table_settings=table_settings,
+            use_ocr=use_ocr,
+            ocr_config=ocr_config,
+            text_options=text_options,
+            cell_extraction_func=cell_extraction_func,
+            show_progress=show_progress,
+        )
 
-    def extract_tables(self, table_settings={}) -> List[Any]:
+    def extract_tables(
+        self,
+        method: Optional[str] = None,
+        table_settings: Optional[dict] = None,
+        check_tatr: bool = True,
+    ) -> List[List[List[str]]]:
         """
-        Extract tables from this page.
+        Extract all tables from this page with enhanced method support.
 
         Args:
-            table_settings: Additional extraction parameters
+            method: Method to use: 'pdfplumber', 'stream', 'lattice', or None (auto-detect).
+                    'stream' uses text-based strategies, 'lattice' uses line-based strategies.
+                    Note: 'tatr' and 'text' methods are not supported for extract_tables.
+            table_settings: Settings for pdfplumber table extraction.
+            check_tatr: If True (default), first check for TATR-detected table regions
+                        and extract from those before falling back to pdfplumber methods.
 
         Returns:
-            List of extracted tables
+            List of tables, where each table is a list of rows, and each row is a list of cell values.
         """
-        # pdfplumber returns list of tables
-        return self._page.extract_tables(table_settings)
+        if table_settings is None:
+            table_settings = {}
+
+        # Check for TATR-detected table regions first if enabled
+        if check_tatr:
+            try:
+                tatr_tables = self.find_all("region[type=table][model=tatr]")
+                if tatr_tables:
+                    logger.debug(f"Page {self.number}: Found {len(tatr_tables)} TATR table regions, extracting from those...")
+                    extracted_tables = []
+                    for table_region in tatr_tables:
+                        try:
+                            table_data = table_region.extract_table(method="tatr")
+                            if table_data:  # Only add non-empty tables
+                                extracted_tables.append(table_data)
+                        except Exception as e:
+                            logger.warning(f"Failed to extract table from TATR region {table_region.bbox}: {e}")
+                    
+                    if extracted_tables:
+                        logger.debug(f"Page {self.number}: Successfully extracted {len(extracted_tables)} tables from TATR regions")
+                        return extracted_tables
+                    else:
+                        logger.debug(f"Page {self.number}: TATR regions found but no tables extracted, falling back to pdfplumber")
+                else:
+                    logger.debug(f"Page {self.number}: No TATR table regions found, using pdfplumber methods")
+            except Exception as e:
+                logger.debug(f"Page {self.number}: Error checking TATR regions: {e}, falling back to pdfplumber")
+
+        # Auto-detect method if not specified (try lattice first, then stream)
+        if method is None:
+            logger.debug(f"Page {self.number}: Auto-detecting tables extraction method...")
+            
+            # Try lattice first
+            try:
+                lattice_settings = table_settings.copy()
+                lattice_settings.setdefault("vertical_strategy", "lines")
+                lattice_settings.setdefault("horizontal_strategy", "lines")
+                
+                logger.debug(f"Page {self.number}: Trying 'lattice' method first for tables...")
+                lattice_result = self._page.extract_tables(lattice_settings)
+                
+                # Check if lattice found meaningful tables
+                if (lattice_result and len(lattice_result) > 0 and 
+                    any(any(any(cell and cell.strip() for cell in row if cell) for row in table if table) for table in lattice_result)):
+                    logger.debug(f"Page {self.number}: 'lattice' method found {len(lattice_result)} tables")
+                    return lattice_result
+                else:
+                    logger.debug(f"Page {self.number}: 'lattice' method found no meaningful tables")
+                    
+            except Exception as e:
+                logger.debug(f"Page {self.number}: 'lattice' method failed: {e}")
+            
+            # Fall back to stream
+            logger.debug(f"Page {self.number}: Falling back to 'stream' method for tables...")
+            stream_settings = table_settings.copy()
+            stream_settings.setdefault("vertical_strategy", "text")
+            stream_settings.setdefault("horizontal_strategy", "text")
+            
+            return self._page.extract_tables(stream_settings)
+
+        effective_method = method
+
+        # Handle method aliases
+        if effective_method == "stream":
+            logger.debug("Using 'stream' method alias for 'pdfplumber' with text-based strategies.")
+            effective_method = "pdfplumber"
+            table_settings.setdefault("vertical_strategy", "text")
+            table_settings.setdefault("horizontal_strategy", "text")
+        elif effective_method == "lattice":
+            logger.debug("Using 'lattice' method alias for 'pdfplumber' with line-based strategies.")
+            effective_method = "pdfplumber"
+            table_settings.setdefault("vertical_strategy", "lines")
+            table_settings.setdefault("horizontal_strategy", "lines")
+
+        # Use the selected method
+        if effective_method == "pdfplumber":
+            return self._page.extract_tables(table_settings)
+        else:
+            raise ValueError(
+                f"Unknown tables extraction method: '{method}'. Choose from 'pdfplumber', 'stream', 'lattice'."
+            )
 
     def _load_elements(self):
         """Load all elements from the page via ElementManager."""
@@ -2198,7 +2312,7 @@ class Page(ClassificationMixin, ExtractionMixin, ShapeDetectionMixin):
     def viewer(
         self,
         # elements_to_render: Optional[List['Element']] = None, # No longer needed, from_page handles it
-        # include_element_types: List[str] = ['word', 'line', 'rect', 'region'] # No longer needed
+        # include_source_types: List[str] = ['word', 'line', 'rect', 'region'] # No longer needed
     ) -> Optional["SimpleInteractiveViewerWidget"]:  # Return type hint updated
         """
         Creates and returns an interactive ipywidget for exploring elements on this page.

natural_pdf/elements/base.py

@@ -59,7 +59,7 @@ class DirectionalMixin:
         direction: str,
         size: Optional[float] = None,
         cross_size: str = "full",
-        include_element: bool = False,
+        include_source: bool = False,
         until: Optional[str] = None,
         include_endpoint: bool = True,
         **kwargs,
@@ -71,7 +71,7 @@ class DirectionalMixin:
             direction: 'left', 'right', 'above', or 'below'
             size: Size in the primary direction (width for horizontal, height for vertical)
             cross_size: Size in the cross direction ('full' or 'element')
-            include_element: Whether to include this element/region's area in the result
+            include_source: Whether to include this element/region's area in the result
             until: Optional selector string to specify a boundary element
             include_endpoint: Whether to include the boundary element found by 'until'
             **kwargs: Additional parameters for the 'until' selector search
@@ -85,7 +85,7 @@ class DirectionalMixin:
         is_positive = direction in ("right", "below")  # right/below are positive directions
         pixel_offset = 1  # Offset for excluding elements/endpoints
 
-        # 1. Determine initial boundaries based on direction and include_element
+        # 1. Determine initial boundaries based on direction and include_source
         if is_horizontal:
             # Initial cross-boundaries (vertical)
             y0 = 0 if cross_size == "full" else self.top
@@ -93,11 +93,11 @@ class DirectionalMixin:
 
             # Initial primary boundaries (horizontal)
             if is_positive:  # right
-                x0_initial = self.x0 if include_element else self.x1 + pixel_offset
+                x0_initial = self.x0 if include_source else self.x1 + pixel_offset
                 x1_initial = self.x1  # This edge moves
             else:  # left
                 x0_initial = self.x0  # This edge moves
-                x1_initial = self.x1 if include_element else self.x0 - pixel_offset
+                x1_initial = self.x1 if include_source else self.x0 - pixel_offset
         else:  # Vertical
             # Initial cross-boundaries (horizontal)
             x0 = 0 if cross_size == "full" else self.x0
@@ -105,11 +105,11 @@ class DirectionalMixin:
 
             # Initial primary boundaries (vertical)
             if is_positive:  # below
-                y0_initial = self.top if include_element else self.bottom + pixel_offset
+                y0_initial = self.top if include_source else self.bottom + pixel_offset
                 y1_initial = self.bottom  # This edge moves
             else:  # above
                 y0_initial = self.top  # This edge moves
-                y1_initial = self.bottom if include_element else self.top - pixel_offset
+                y1_initial = self.bottom if include_source else self.top - pixel_offset
 
         # 2. Calculate the final primary boundary, considering 'size' or page limits
         if is_horizontal:
@@ -195,7 +195,7 @@ class DirectionalMixin:
 
         result = Region(self.page, final_bbox)
         result.source_element = self
-        result.includes_source = include_element
+        result.includes_source = include_source
         # Optionally store the boundary element if found
         if target:
             result.boundary_element = target
@@ -206,7 +206,7 @@ class DirectionalMixin:
         self,
         height: Optional[float] = None,
         width: str = "full",
-        include_element: bool = False,
+        include_source: bool = False,
         until: Optional[str] = None,
         include_endpoint: bool = True,
         **kwargs,
@@ -217,7 +217,7 @@ class DirectionalMixin:
         Args:
             height: Height of the region above, in points
             width: Width mode - "full" for full page width or "element" for element width
-            include_element: Whether to include this element/region in the result (default: False)
+            include_source: Whether to include this element/region in the result (default: False)
             until: Optional selector string to specify an upper boundary element
             include_endpoint: Whether to include the boundary element in the region (default: True)
             **kwargs: Additional parameters
@@ -229,7 +229,7 @@ class DirectionalMixin:
             direction="above",
             size=height,
             cross_size=width,
-            include_element=include_element,
+            include_source=include_source,
             until=until,
             include_endpoint=include_endpoint,
             **kwargs,
@@ -239,7 +239,7 @@ class DirectionalMixin:
         self,
         height: Optional[float] = None,
         width: str = "full",
-        include_element: bool = False,
+        include_source: bool = False,
         until: Optional[str] = None,
         include_endpoint: bool = True,
         **kwargs,
@@ -250,7 +250,7 @@ class DirectionalMixin:
         Args:
             height: Height of the region below, in points
             width: Width mode - "full" for full page width or "element" for element width
-            include_element: Whether to include this element/region in the result (default: False)
+            include_source: Whether to include this element/region in the result (default: False)
             until: Optional selector string to specify a lower boundary element
             include_endpoint: Whether to include the boundary element in the region (default: True)
             **kwargs: Additional parameters
@@ -262,7 +262,7 @@ class DirectionalMixin:
             direction="below",
             size=height,
             cross_size=width,
-            include_element=include_element,
+            include_source=include_source,
             until=until,
             include_endpoint=include_endpoint,
             **kwargs,
@@ -272,7 +272,7 @@ class DirectionalMixin:
         self,
         width: Optional[float] = None,
         height: str = "full",
-        include_element: bool = False,
+        include_source: bool = False,
         until: Optional[str] = None,
         include_endpoint: bool = True,
         **kwargs,
@@ -283,7 +283,7 @@ class DirectionalMixin:
         Args:
             width: Width of the region to the left, in points
             height: Height mode - "full" for full page height or "element" for element height
-            include_element: Whether to include this element/region in the result (default: False)
+            include_source: Whether to include this element/region in the result (default: False)
             until: Optional selector string to specify a left boundary element
             include_endpoint: Whether to include the boundary element in the region (default: True)
             **kwargs: Additional parameters
@@ -295,7 +295,7 @@ class DirectionalMixin:
             direction="left",
             size=width,
             cross_size=height,
-            include_element=include_element,
+            include_source=include_source,
             until=until,
             include_endpoint=include_endpoint,
             **kwargs,
@@ -305,7 +305,7 @@ class DirectionalMixin:
         self,
         width: Optional[float] = None,
         height: str = "full",
-        include_element: bool = False,
+        include_source: bool = False,
         until: Optional[str] = None,
         include_endpoint: bool = True,
         **kwargs,
@@ -316,7 +316,7 @@ class DirectionalMixin:
         Args:
             width: Width of the region to the right, in points
             height: Height mode - "full" for full page height or "element" for element height
-            include_element: Whether to include this element/region in the result (default: False)
+            include_source: Whether to include this element/region in the result (default: False)
             until: Optional selector string to specify a right boundary element
             include_endpoint: Whether to include the boundary element in the region (default: True)
             **kwargs: Additional parameters
@@ -328,7 +328,7 @@ class DirectionalMixin:
             direction="right",
             size=width,
             cross_size=height,
-            include_element=include_element,
+            include_source=include_source,
             until=until,
             include_endpoint=include_endpoint,
             **kwargs,