natural-pdf 0.1.38__py3-none-any.whl → 0.2.0__py3-none-any.whl

natural_pdf/elements/base.py

@@ -2,11 +2,12 @@
 Base Element class for natural-pdf.
 """
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, overload
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, overload
 
 from PIL import Image
 
 from natural_pdf.classification.mixin import ClassificationMixin
+from natural_pdf.core.render_spec import RenderSpec, Visualizable
 from natural_pdf.describe.mixin import DescribeMixin
 
 # Import selector parsing functions
@@ -15,7 +16,7 @@ from natural_pdf.selectors.parser import parse_selector, selector_to_filter_func
 if TYPE_CHECKING:
     from natural_pdf.classification.manager import ClassificationManager  # noqa: F401
     from natural_pdf.core.page import Page
-    from natural_pdf.elements.collections import ElementCollection
+    from natural_pdf.elements.element_collection import ElementCollection
     from natural_pdf.elements.region import Region
 
 
@@ -563,7 +564,56 @@ class DirectionalMixin:
         return matches[0]
 
 
-class Element(DirectionalMixin, ClassificationMixin, DescribeMixin):
+class HighlightableMixin:
+    """
+    Mixin that provides the highlighting protocol for elements.
+
+    This protocol enables ElementCollection.show() to work with mixed content
+    including FlowRegions and elements from multiple pages by providing a
+    standard way to get highlight specifications.
+    """
+
+    def get_highlight_specs(self) -> List[Dict[str, Any]]:
+        """
+        Get highlight specifications for this element.
+
+        Returns a list of dictionaries, each containing:
+        - page: The Page object to highlight on
+        - page_index: The 0-based index of the page
+        - bbox: The bounding box (x0, y0, x1, y1) to highlight
+        - polygon: Optional polygon coordinates for non-rectangular highlights
+        - element: Reference to the element being highlighted
+
+        For regular elements, this returns a single spec.
+        For FlowRegions, this returns specs for all constituent regions.
+
+        Returns:
+            List of highlight specification dictionaries
+        """
+        # Default implementation for regular elements
+        if not hasattr(self, "page") or self.page is None:
+            return []
+
+        if not hasattr(self, "bbox") or self.bbox is None:
+            return []
+
+        spec = {
+            "page": self.page,
+            "page_index": self.page.index if hasattr(self.page, "index") else 0,
+            "bbox": self.bbox,
+            "element": self,
+        }
+
+        # Add polygon if available
+        if hasattr(self, "polygon") and hasattr(self, "has_polygon") and self.has_polygon:
+            spec["polygon"] = self.polygon
+
+        return [spec]
+
+
+class Element(
+    DirectionalMixin, ClassificationMixin, DescribeMixin, HighlightableMixin, Visualizable
+):
     """Base class for all PDF elements.
 
     This class provides common properties and methods for all PDF elements,
@@ -1024,7 +1074,7 @@ class Element(DirectionalMixin, ClassificationMixin, DescribeMixin):
         label: str = "",
         color: Optional[Tuple[float, float, float]] = None,
         use_color_cycling: bool = True,
-        include_attrs: Optional[List[str]] = None,
+        annotate: Optional[List[str]] = None,
         existing: str = "append",
     ) -> "Element":
         """Highlight the element with the specified colour.
@@ -1042,7 +1092,7 @@ class Element(DirectionalMixin, ClassificationMixin, DescribeMixin):
             "label": label,
             "use_color_cycling": use_color_cycling,
             "element": self,  # Pass the element itself so attributes can be accessed
-            "include_attrs": include_attrs,
+            "annotate": annotate,
             "existing": existing,
         }
 
@@ -1056,84 +1106,67 @@ class Element(DirectionalMixin, ClassificationMixin, DescribeMixin):
 
         return self
 
-    def show(
+    def _get_render_specs(
         self,
-        resolution: Optional[float] = None,
-        labels: bool = True,
-        legend_position: str = "right",
-        color: Optional[Union[Tuple, str]] = "red",  # Default color for single element
+        mode: Literal["show", "render"] = "show",
+        color: Optional[Union[str, Tuple[int, int, int]]] = None,
+        highlights: Optional[List[Dict[str, Any]]] = None,
+        crop: Union[bool, Literal["content"]] = False,
+        crop_bbox: Optional[Tuple[float, float, float, float]] = None,
         label: Optional[str] = None,
-        width: Optional[int] = None,  # Add width parameter
-        crop: bool = False,  # NEW: Crop to element bounds before legend
-    ) -> Optional["Image.Image"]:
-        """
-        Show the page with only this element highlighted temporarily.
+        **kwargs,
+    ) -> List[RenderSpec]:
+        """Get render specifications for this element.
 
         Args:
-            resolution: Resolution in DPI for rendering (default: uses global options, fallback to 144 DPI)
-            labels: Whether to include a legend for the highlight
-            legend_position: Position of the legend
-            color: Color to highlight this element (default: red)
-            label: Optional label for this element in the legend
-            width: Optional width for the output image in pixels
-            crop: If True, crop the rendered image to this element's
-                        bounding box before legends/overlays are added.
+            mode: Rendering mode - 'show' includes highlights, 'render' is clean
+            color: Color for highlighting this element in show mode
+            highlights: Additional highlight groups to show
+            crop: Whether to crop to element bounds
+            crop_bbox: Explicit crop bounds
+            label: Optional label for this element
+            **kwargs: Additional parameters
 
         Returns:
-            PIL Image of the page with only this element highlighted, or None if error.
+            List with single RenderSpec for this element's page
         """
-        # Apply global options as defaults
-        import natural_pdf
-
-        if resolution is None:
-            if natural_pdf.options.image.resolution is not None:
-                resolution = natural_pdf.options.image.resolution
-            else:
-                resolution = 144  # Default resolution when none specified
-        if not hasattr(self, "page") or not self.page:
-            logger.warning(f"Cannot show element, missing 'page' attribute: {self}")
-            return None
-        if not hasattr(self.page, "_highlighter") or not self.page._highlighter:
-            logger.warning(f"Cannot show element, page lacks highlighter service: {self}")
-            return None
-
-        service = self.page._highlighter
-
-        # Determine the label if not provided
-        display_label = label if label is not None else f"{self.__class__.__name__}"
-
-        # Prepare temporary highlight data for just this element
-        temp_highlight_data = {
-            "page_index": self.page.index,
-            "bbox": self.bbox if not self.has_polygon else None,
-            "polygon": self.polygon if self.has_polygon else None,
-            "color": color,  # Use provided or default color
-            "label": display_label,
-            "use_color_cycling": False,  # Explicitly false for single preview
-        }
+        if not hasattr(self, "page") or self.page is None:
+            return []
+
+        spec = RenderSpec(page=self.page)
+
+        # Handle cropping
+        if crop_bbox:
+            spec.crop_bbox = crop_bbox
+        elif crop == "content" or crop is True:
+            # Crop to element bounds
+            if hasattr(self, "bbox") and self.bbox:
+                spec.crop_bbox = self.bbox
+
+        # Add highlight in show mode
+        if mode == "show":
+            # Use provided label or generate one
+            element_label = label if label is not None else self.__class__.__name__
+
+            spec.add_highlight(
+                element=self,
+                color=color or "red",  # Default red for single element
+                label=element_label,
+            )
 
-        # Determine crop bbox
-        crop_bbox = self.bbox if crop else None
+            # Add additional highlight groups if provided
+            if highlights:
+                for group in highlights:
+                    group_elements = group.get("elements", [])
+                    group_color = group.get("color", color)
+                    group_label = group.get("label")
 
-        # Check if we actually got geometry data
-        if temp_highlight_data["bbox"] is None and temp_highlight_data["polygon"] is None:
-            logger.warning(f"Cannot show element, failed to get bbox or polygon: {self}")
-            return None
+                    for elem in group_elements:
+                        # Only add if element is on same page
+                        if hasattr(elem, "page") and elem.page == self.page:
+                            spec.add_highlight(element=elem, color=group_color, label=group_label)
 
-        # Use render_preview to show only this highlight
-        try:
-            return service.render_preview(
-                page_index=self.page.index,
-                temporary_highlights=[temp_highlight_data],
-                resolution=resolution,
-                width=width,  # Pass the width parameter
-                labels=labels,
-                legend_position=legend_position,
-                crop_bbox=crop_bbox,
-            )
-        except Exception as e:
-            logger.error(f"Error calling render_preview for element {self}: {e}", exc_info=True)
-            return None
+        return [spec]
 
     def save(
         self,
@@ -1346,22 +1379,14 @@ class Element(DirectionalMixin, ClassificationMixin, DescribeMixin):
             resolution = kwargs.get("resolution", 150)
             from natural_pdf.elements.region import Region  # Local import to avoid cycles
 
-            return self.expand().to_image(
+            # Use render() for clean image without highlights
+            return self.expand().render(
                 resolution=resolution,
-                include_highlights=False,
                 crop=True,
             )
         else:
             raise ValueError(f"Unsupported model_type for classification: {model_type}")
 
-    # ------------------------------------------------------------------
-    # Lightweight to_image proxy (vision models, previews, etc.)
-    # ------------------------------------------------------------------
-
-    def to_image(self, *args, **kwargs):  # type: ignore[override]
-        """Generate an image of this element by delegating to a temporary Region."""
-        return self.expand().to_image(*args, **kwargs)
-
     # ------------------------------------------------------------------
     # Unified analysis storage (maps to metadata["analysis"])
     # ------------------------------------------------------------------