Merge pull request #385 from aiden2480/extra-filters

Add custom filter predicate and header/footer filters

Author: jstockwin

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Added extra filtering methods for ElementList
+
 ## [0.12.0] - 2023-11-10
 
 ### Added

py_pdf_parser/filtering.py

@@ -1,5 +1,6 @@
 from typing import (
     TYPE_CHECKING,
+    Callable,
     FrozenSet,
     Iterable,
     Iterator,
@@ -103,6 +104,21 @@ def add_tag_to_elements(self, tag: str) -> None:
         for element in self:
             element.add_tag(tag)
 
+    def filter(self, predicate: Callable[["PDFElement"], bool]) -> "ElementList":
+        """
+        Filter by elements matching a custom predicate. Anything that passes the
+        predicate is included in the new ElementList.
+
+        Args:
+            predicate (Callable[[PDFElement], bool]): The predicate to filter by.
+
+        Returns:
+            ElementList: The filtered list.
+        """
+
+        new_indexes = set(element._index for element in self if predicate(element))
+        return ElementList(self.document, new_indexes)
+
     def filter_by_tag(self, tag: str) -> "ElementList":
         """
         Filter for elements containing only the given tag.
@@ -113,8 +129,8 @@ def filter_by_tag(self, tag: str) -> "ElementList":
         Returns:
             ElementList: The filtered list.
         """
-        new_indexes = set(element._index for element in self if tag in element.tags)
-        return ElementList(self.document, new_indexes)
+
+        return self.filter(lambda e: tag in e.tags)
 
     def filter_by_tags(self, *tags: str) -> "ElementList":
         """
@@ -126,12 +142,8 @@ def filter_by_tags(self, *tags: str) -> "ElementList":
         Returns:
             ElementList: The filtered list.
         """
-        new_indexes = set(
-            element._index
-            for element in self
-            if any(tag in element.tags for tag in tags)
-        )
-        return ElementList(self.document, new_indexes)
+
+        return self.filter(lambda e: any(tag in e.tags for tag in tags))
 
     def filter_by_text_equal(self, text: str, stripped: bool = True) -> "ElementList":
         """
@@ -145,11 +157,8 @@ def filter_by_text_equal(self, text: str, stripped: bool = True) -> "ElementList
         Returns:
             ElementList: The filtered list.
         """
-        new_indexes = set(
-            element._index for element in self if element.text(stripped) == text
-        )
 
-        return ElementList(self.document, new_indexes)
+        return self.filter(lambda e: e.text(stripped) == text)
 
     def filter_by_text_contains(self, text: str) -> "ElementList":
         """
@@ -161,8 +170,8 @@ def filter_by_text_contains(self, text: str) -> "ElementList":
         Returns:
             ElementList: The filtered list.
         """
-        new_indexes = set(element._index for element in self if text in element.text())
-        return ElementList(self.document, new_indexes)
+
+        return self.filter(lambda e: text in e.text())
 
     def filter_by_regex(
         self,
@@ -216,6 +225,19 @@ def filter_by_fonts(self, *fonts: str) -> "ElementList":
         new_indexes = self.indexes & self.document._element_indexes_with_fonts(*fonts)
         return ElementList(self.document, new_indexes)
 
+    def filter_by_font_size(self, font_size: float) -> "ElementList":
+        """
+        Filter for elements of a particular font size.
+
+        Args:
+            font_size (float): The font size to filter for.
+
+        Returns:
+            ElementList: The filtered list.
+        """
+
+        return self.filter(lambda e: e.font_size == font_size)
+
     def filter_by_page(self, page_number: int) -> "ElementList":
         """
         Filter for elements on the given page.
@@ -786,7 +808,7 @@ def extract_single_element(self) -> "PDFElement":
                 f"There are {len(self.indexes)} elements in the ElementList"
             )
 
-        return self.document._element_list[list(self.indexes)[0]]
+        return self.first()
 
     def add_element(self, element: "PDFElement") -> "ElementList":
         """
@@ -919,6 +941,74 @@ def move_backwards_from(
         """
         return self.move_forwards_from(element, count=-count, capped=capped)
 
+    def filter_out_header(self, bottom_of_header_y: float) -> "ElementList":
+        """
+        Filter out header elements, as specified by a certain y position. Only elements
+        completely within the header are discarded. Partially overlapping elements are
+        kept.
+
+        Args:
+            bottom_of_header_y (float): The Y coordinate of the bottom of the header.
+
+        Note:
+            Y decreases as elements go down the page.
+
+        Returns:
+            ElementList: The filtered list without header elements.
+        """
+
+        return self.filter(lambda e: e.bounding_box.y0 < bottom_of_header_y)
+
+    def filter_out_footer(self, top_of_footer_y: float) -> "ElementList":
+        """
+        Filter out footer elements, as specified by a certain y position. Only elements
+        completely within the footer are discarded. Partially overlapping elements are
+        kept.
+
+        Args:
+            top_of_footer_y (float): The Y coordinate of the top of the footer.
+
+        Note:
+            Y decreases as elements go down the page.
+
+        Returns:
+            ElementList: The filtered list without footer elements.
+        """
+
+        return self.filter(lambda e: e.bounding_box.y1 > top_of_footer_y)
+
+    def first(self) -> "PDFElement":
+        """
+        Returns the first element in the ElementList
+
+        Unlike extract_single_element, an error is not thrown if there is more
+        than one element in the ElementList.
+
+        Raises:
+            NoElementFoundError: If there are no elements in the ElementList
+        """
+
+        if len(self.indexes) == 0:
+            raise NoElementFoundError("There are no elements in the ElementList")
+
+        return self[0]
+
+    def last(self) -> "PDFElement":
+        """
+        Returns the last element in the ElementList
+
+        Unlike extract_single_element, an error is not thrown if there is more
+        than one element in the ElementList.
+
+        Raises:
+            NoElementFoundError: If there are no elements in the ElementList
+        """
+
+        if len(self.indexes) == 0:
+            raise NoElementFoundError("There are no elements in the ElementList")
+
+        return self[-1]
+
     def __intersect_indexes_with_self(self, new_indexes: Set[int]) -> "ElementList":
         return self & ElementList(self.document, new_indexes)
 

tests/test_filtering.py

@@ -1211,6 +1211,63 @@ def test_move_backwards_from(self):
             self.elem_list[-1],
         )
 
+    def test_filter(self):
+        elem1 = FakePDFMinerTextElement(text="even")
+        elem2 = FakePDFMinerTextElement(text="odd")
+        elem3 = FakePDFMinerTextElement(text="even")
+        elem4 = FakePDFMinerTextElement(text="odd")
+        elem5 = FakePDFMinerTextElement(text="even")
+
+        doc = create_pdf_document([elem1, elem2, elem3, elem4, elem5])
+        even_elems = doc.elements.filter(lambda e: len(e.text()) % 2 == 0)
+        odd_elems = doc.elements.filter(lambda e: len(e.text()) % 2 == 1)
+
+        self.assertEqual(ElementList(doc, {0, 2, 4}), even_elems)
+        self.assertEqual(ElementList(doc, {1, 3}), odd_elems)
+
+    def test_filter_by_font_size(self):
+        elem1 = FakePDFMinerTextElement(font_name="foo", font_size=1)
+        elem2 = FakePDFMinerTextElement(font_name="bar", font_size=2)
+        elem3 = FakePDFMinerTextElement(font_name="bat", font_size=2)
+        elem4 = FakePDFMinerTextElement(font_name="baz", font_size=3)
+        doc = create_pdf_document([elem1, elem2, elem3, elem4])
+
+        self.assertEqual(ElementList(doc, {1, 2}), doc.elements.filter_by_font_size(2))
+
+    def test_filter_out_header(self):
+        bbox1 = BoundingBox(20, 30, 75, 80)  # Completely within header - discarded
+        bbox2 = BoundingBox(20, 30, 25, 75)  # Completely outside header - kept
+        bbox3 = BoundingBox(20, 30, 10, 20)  # Partially within header - kept
+
+        elems = [FakePDFMinerTextElement(b) for b in (bbox1, bbox2, bbox3)]
+        doc = create_pdf_document(elems)
+
+        self.assertEqual(ElementList(doc, {1, 2}), doc.elements.filter_out_header(50))
+
+    def test_filter_out_footer(self):
+        bbox1 = BoundingBox(20, 30, 75, 80)  # Completely outside footer - kept
+        bbox2 = BoundingBox(20, 30, 25, 75)  # Partially within footer - kept
+        bbox3 = BoundingBox(20, 30, 10, 20)  # Completely within footer - discarded
+
+        elems = [FakePDFMinerTextElement(b) for b in (bbox1, bbox2, bbox3)]
+        doc = create_pdf_document(elems)
+
+        self.assertEqual(ElementList(doc, {0, 1}), doc.elements.filter_out_footer(50))
+
+    def test_first(self):
+        actual_first_elem = self.extract_element_from_list(self.elem1, self.elem_list)
+        self.assertEqual(self.elem_list.first(), actual_first_elem)
+
+        with self.assertRaises(NoElementFoundError):
+            self.elem_list.filter_by_tag("non_existent_tag").first()
+
+    def test_last(self):
+        actual_last_elem = self.extract_element_from_list(self.elem6, self.elem_list)
+        self.assertEqual(self.elem_list.last(), actual_last_elem)
+
+        with self.assertRaises(NoElementFoundError):
+            self.elem_list.filter_by_tag("non_existent_tag").last()
+
     def test_repr(self):
         self.assertEqual(repr(self.elem_list), "<ElementList of 6 elements>")