Merge pull request #126 from mohangk/webui

Further refactoring of book_handler logic into PDFParser

abogen/book_parser.py

@@ -110,41 +110,158 @@ class PdfParser(BaseBookParser):
             logging.error(f"Error loading PDF {self.book_path}: {e}")
             raise
 
+    def _extract_book_metadata(self):
+        # PDF metadata extraction can be added here if needed
+        # For now, base class metadata is empty dict
+        pass
+
     def process_content(self, replace_single_newlines=True):
         if not self.pdf_doc:
             self.load()
 
+        # 1. Extract text from all pages first
         for page_num in range(len(self.pdf_doc)):
             text = clean_text(self.pdf_doc[page_num].get_text())
-
+            
             # Clean up common PDF artifacts:
-            # - Remove bracketed numbers often used for citations [1]
             text = _BRACKETED_NUMBERS_PATTERN.sub("", text)
-            # - Remove standalone page numbers often found in headers/footers
             text = _STANDALONE_PAGE_NUMBERS_PATTERN.sub("", text)
-            # - Remove page numbers at end of lines
             text = _PAGE_NUMBERS_AT_END_PATTERN.sub("", text)
-            # - Remove page numbers with dashes - 4 -
             text = _PAGE_NUMBERS_WITH_DASH_PATTERN.sub("", text)
 
             page_id = f"page_{page_num + 1}"
             self.content_texts[page_id] = text
             self.content_lengths[page_id] = calculate_text_length(text)
 
+        # 2. Build Navigation Structure
+        toc = self.pdf_doc.get_toc()
+        
+        if not toc:
+            # Fallback: Flat list of pages if no TOC
+            self.processed_nav_structure = []
+            pages_node = {
+                "title": "Pages",
+                "src": None,
+                "children": [],
+                "has_content": False
+            }
+            # Add all pages as children
+            for page_num in range(len(self.pdf_doc)):
+               page_id = f"page_{page_num + 1}"
+               title = self._get_page_title(page_num, self.content_texts.get(page_id, ""))
+               pages_node["children"].append({
+                   "title": title,
+                   "src": page_id,
+                   "children": [],
+                   "has_content": True
+               })
+            self.processed_nav_structure.append(pages_node)
+        else:
+            self.processed_nav_structure = self._build_structure_from_toc(toc)
+
         return self.content_texts, self.content_lengths
 
-    def _extract_book_metadata(self):
+    def _get_page_title(self, page_num, text):
-        # PDF metadata extraction can be added here if needed
+        title = f"Page {page_num + 1}"
-        # For now, base class metadata is empty dict
+        if text:
-        pass
+            first_line = text.split("\n", 1)[0].strip()
+            if first_line and len(first_line) < 100:
+                title += f" - {first_line}"
+        return title
 
-    def get_chapters(self):
+    def _build_structure_from_toc(self, toc):
-        # PDF specific implementation because it doesn't use nav structure
+        # 1. Flatten TOC to easier list (page_num, title, level)
-        chapters = []
+        # fitz TOC is [[lvl, title, page, dest], ...]
-        if self.pdf_doc:
+        
-            for i in range(len(self.pdf_doc)):
+        bookmarks = []
-                chapters.append((f"page_{i+1}", f"Page {i+1}"))
+        for entry in toc:
-        return chapters
+            lvl, title, page = entry[:3]
+            if isinstance(page, int):
+                page_idx = page - 1
+            else:
+                 # Handle potential complex destinations if necessary, but usually simple int
+                 # PyMuPDF docs say int.
+                 page_idx = -1 
+            
+            if page_idx >= 0:
+                bookmarks.append({"level": lvl, "title": title, "page": page_idx})
+
+        
+        root_children = []
+        stack = [] # Stack of (level, list_to_append_to)
+        stack.append((0, root_children)) 
+
+        # Step 1: Build the Skeleton Tree from TOC
+        # And keep a flat list of these nodes to associate with pages.
+        
+        processed_nodes = [] # List of (page_idx, node_dict)
+        
+        for entry in bookmarks:
+            node = {
+                "title": entry["title"],
+                "src": f"page_{entry['page'] + 1}",
+                "children": [],
+                "has_content": True
+            }
+            
+            # Find parent
+            level = entry["level"]
+            
+            # Adjust stack
+            while stack and stack[-1][0] >= level:
+                stack.pop()
+            
+            parent_list = stack[-1][1]
+            parent_list.append(node)
+            
+            stack.append((level, node["children"]))
+            processed_nodes.append((entry["page"], node))
+            
+        # Step 3: Add gap pages.
+        # Sort processed_nodes by page index to find ranges.
+        sorted_bookmarks = sorted(processed_nodes, key=lambda x: x[0])
+        
+        # Set of pages that are "bookmarks"
+        bookmarked_pages = set(p for p, n in sorted_bookmarks)
+        
+        current_node = None
+        # We need a way to look up bookmarks starting at p
+        bookmarks_by_page = {}
+        for p, node in processed_nodes:
+            if p not in bookmarks_by_page:
+                bookmarks_by_page[p] = []
+            bookmarks_by_page[p].append(node)
+
+        
+        # Let's iterate.
+        for page_num in range(len(self.pdf_doc)):
+            page_id = f"page_{page_num + 1}"
+            
+            # Check if this page STARTS bookmarks
+            if page_num in bookmarks_by_page:
+                
+                starts = bookmarks_by_page[page_num]
+                current_node = starts[-1] 
+                
+                continue
+
+            # If page is NOT a bookmark, it's a "gap page".
+            # Add as child to current_node
+            title = self._get_page_title(page_num, self.content_texts.get(page_id, ""))
+            page_node = {
+                "title": title,
+                "src": page_id,
+                "children": [],
+                "has_content": True
+            }
+            
+            if current_node:
+                current_node["children"].append(page_node)
+            else:
+                # No preceding bookmark. Add to root.
+                root_children.append(page_node)
+                
+        return root_children
 
 
 class MarkdownParser(BaseBookParser):

abogen/pyqt/book_handler.py

@@ -394,10 +394,7 @@ class HandlerDialog(QDialog):
         font.setBold(True)
         info_item.setFont(0, font)
 
-        # TODO look into why we need this here and not just rely on logic in PDFParser
+        if self.processed_nav_structure:
-        if self.parser.file_type == "pdf":
-            self._build_pdf_tree()
-        elif self.processed_nav_structure:
             self._build_tree_from_nav(
                 self.processed_nav_structure, self.treeWidget
             )
@@ -471,200 +468,6 @@ class HandlerDialog(QDialog):
                 self._build_tree_from_nav(children, item, seen_content_hashes)
 
 
-    # TODO look into why we need this here and not just rely on logic in PDFParser
-    def _build_pdf_tree(self):
-        pdf_doc = getattr(self.parser, "pdf_doc", None)
-        if not pdf_doc:
-            return
-
-        outline = pdf_doc.get_toc()
-        self.has_pdf_bookmarks = bool(outline)
-        self.bookmark_items_map = {}
-
-        if not outline:
-            self._build_pdf_pages_tree()
-            return
-
-        bookmark_pages = []
-        page_to_bookmark = {}
-        next_page_boundaries = {}
-        added_pages = set()
-
-        def extract_page_numbers(entries):
-            for entry in entries:
-                if len(entry) >= 3:
-                    _, title, page = entry[:3]
-                    page_num = (
-                        page - 1
-                        if isinstance(page, int)
-                        else pdf_doc.resolve_link(page)[0]
-                    )
-                    bookmark_pages.append((page_num, title))
-
-                    if len(entry) > 3 and isinstance(entry[3], list):
-                        extract_page_numbers(entry[3])
-
-        extract_page_numbers(outline)
-        bookmark_pages.sort()
-
-        for i, (page_num, title) in enumerate(bookmark_pages):
-            if i < len(bookmark_pages) - 1:
-                next_page_boundaries[page_num] = bookmark_pages[i + 1][0]
-            page_to_bookmark[page_num] = title
-
-        def build_outline_tree(entries, parent_item):
-            for entry in entries:
-                if len(entry) >= 3:
-                    entry_level, title, page = entry[:3]
-                    page_num = (
-                        page - 1
-                        if isinstance(page, int)
-                        else pdf_doc.resolve_link(page)[0]
-                    )
-                    page_id = f"page_{page_num + 1}"
-                    # attach chapters on same page under original
-                    if page_num in self.bookmark_items_map:
-                        orig = self.bookmark_items_map[page_num]
-                        child = QTreeWidgetItem(orig, [f"{title} (Same page)"])
-                        child.setData(0, Qt.ItemDataRole.UserRole, page_id)
-                        child.setFlags(child.flags() & ~Qt.ItemFlag.ItemIsUserCheckable)
-                        continue
-                    bookmark_item = QTreeWidgetItem(parent_item, [title])
-                    bookmark_item.setData(0, Qt.ItemDataRole.UserRole, page_id)
-                    # only allow checking if this chapter has content
-                    if self.content_lengths.get(page_id, 0) > 0:
-                        bookmark_item.setFlags(
-                            bookmark_item.flags() | Qt.ItemFlag.ItemIsUserCheckable
-                        )
-                        bookmark_item.setCheckState(
-                            0,
-                            (
-                                Qt.CheckState.Checked
-                                if page_id in self.checked_chapters
-                                else Qt.CheckState.Unchecked
-                            ),
-                        )
-                    else:
-                        bookmark_item.setFlags(
-                            bookmark_item.flags() & ~Qt.ItemFlag.ItemIsUserCheckable
-                        )
-                    # map for uncategorized pages
-                    self.bookmark_items_map[page_num] = bookmark_item
-
-                    added_pages.add(page_num)
-
-                    next_page = next_page_boundaries.get(page_num, len(pdf_doc))
-                    for sub_page_num in range(page_num + 1, next_page):
-                        if (
-                            sub_page_num in page_to_bookmark
-                            or sub_page_num in added_pages
-                        ):
-                            continue
-
-                        page_id = f"page_{sub_page_num + 1}"
-                        page_title = f"Page {sub_page_num + 1}"
-
-                        page_text = self.content_texts.get(page_id, "").strip()
-                        if page_text:
-                            first_line = page_text.split("\n", 1)[0].strip()
-                            if first_line and len(first_line) < 100:
-                                page_title += f" - {first_line}"
-
-                        page_item = QTreeWidgetItem(bookmark_item, [page_title])
-                        page_item.setData(0, Qt.ItemDataRole.UserRole, page_id)
-                        # only allow checking if this sub-page has content
-                        if self.content_lengths.get(page_id, 0) > 0:
-                            page_item.setFlags(
-                                page_item.flags() | Qt.ItemFlag.ItemIsUserCheckable
-                            )
-                            page_item.setCheckState(
-                                0,
-                                (
-                                    Qt.CheckState.Checked
-                                    if page_id in self.checked_chapters
-                                    else Qt.CheckState.Unchecked
-                                ),
-                            )
-                        else:
-                            page_item.setFlags(
-                                page_item.flags() & ~Qt.ItemFlag.ItemIsUserCheckable
-                            )
-
-                        added_pages.add(sub_page_num)
-
-        build_outline_tree(outline, self.treeWidget)
-
-        covered_pages = set(added_pages)
-        # attach any pages without direct bookmarks under nearest preceding chapter
-        uncategorized_pages = [
-            i for i in range(len(pdf_doc)) if i not in covered_pages
-        ]
-        for page_num in uncategorized_pages:
-            # find nearest previous bookmark
-            prev_nums = [n for n in sorted(self.bookmark_items_map) if n < page_num]
-            parent_item = (
-                self.bookmark_items_map[prev_nums[-1]] if prev_nums else self.treeWidget
-            )
-            page_id = f"page_{page_num + 1}"
-            title = f"Page {page_num + 1}"
-            text = self.content_texts.get(page_id, "").strip()
-            if text:
-                first = text.split("\n", 1)[0].strip()
-                if first and len(first) < 100:
-                    title += f" - {first}"
-            page_item = QTreeWidgetItem(parent_item, [title])
-            page_item.setData(0, Qt.ItemDataRole.UserRole, page_id)
-            # only allow checking if uncategorized page has content
-            if self.content_lengths.get(page_id, 0) > 0:
-                page_item.setFlags(page_item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
-                page_item.setCheckState(
-                    0,
-                    (
-                        Qt.CheckState.Checked
-                        if page_id in self.checked_chapters
-                        else Qt.CheckState.Unchecked
-                    ),
-                )
-            else:
-                page_item.setFlags(page_item.flags() & ~Qt.ItemFlag.ItemIsUserCheckable)
-
-    def _build_pdf_pages_tree(self):
-        pdf_doc = getattr(self.parser, "pdf_doc", None)
-        if not pdf_doc:
-            return
-
-        pages_item = QTreeWidgetItem(self.treeWidget, ["Pages"])
-        pages_item.setFlags(pages_item.flags() & ~Qt.ItemFlag.ItemIsUserCheckable)
-        font = pages_item.font(0)
-        font.setBold(True)
-        pages_item.setFont(0, font)
-
-        for page_num in range(len(pdf_doc)):
-            page_id = f"page_{page_num + 1}"
-            page_title = f"Page {page_num + 1}"
-
-            page_text = self.content_texts.get(page_id, "").strip()
-            if page_text:
-                first_line = page_text.split("\n", 1)[0].strip()
-                if first_line and len(first_line) < 100:
-                    page_title += f" - {first_line}"
-
-            page_item = QTreeWidgetItem(pages_item, [page_title])
-            page_item.setData(0, Qt.ItemDataRole.UserRole, page_id)
-            # only allow checking if standalone page has content
-            if self.content_lengths.get(page_id, 0) > 0:
-                page_item.setFlags(page_item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
-                page_item.setCheckState(
-                    0,
-                    (
-                        Qt.CheckState.Checked
-                        if page_id in self.checked_chapters
-                        else Qt.CheckState.Unchecked
-                    ),
-                )
-            else:
-                page_item.setFlags(page_item.flags() & ~Qt.ItemFlag.ItemIsUserCheckable)
-
     def _are_provided_checks_relevant(self):
         if not self.checked_chapters:
             return False
@@ -989,15 +792,6 @@ class HandlerDialog(QDialog):
             iterator += 1
 
     def _run_pdf_auto_check(self):
-        if hasattr(self, "has_pdf_bookmarks") and not self.has_pdf_bookmarks:
-            iterator = QTreeWidgetItemIterator(self.treeWidget)
-            while iterator.value():
-                item = iterator.value()
-                if item.flags() & Qt.ItemFlag.ItemIsUserCheckable:
-                    item.setCheckState(0, Qt.CheckState.Checked)
-                iterator += 1
-            return
-
         iterator = QTreeWidgetItemIterator(self.treeWidget)
         while iterator.value():
             item = iterator.value()
@@ -1006,16 +800,13 @@ class HandlerDialog(QDialog):
                 continue
 
             identifier = item.data(0, Qt.ItemDataRole.UserRole)
-
             if not identifier:
-                iterator += 1
+                 iterator += 1
-                continue
+                 continue
-
+                 
-            if (
+            # Logic: Check item if it has content (already handled by ItemIsUserCheckable flag really)
-                not identifier.startswith("page_")
+            # But duplicate logic from previous implementation:
-                or self.content_lengths.get(identifier, 0) > 0
+            item.setCheckState(0, Qt.CheckState.Checked)
-            ):
-                item.setCheckState(0, Qt.CheckState.Checked)
 
             iterator += 1
 

tests/test_book_handler_regression.py

@@ -8,7 +8,7 @@ from PyQt6.QtWidgets import QApplication
 # Ensure we can import the module
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
-from abogen.book_handler import HandlerDialog
+from abogen.pyqt.book_handler import HandlerDialog
 from ebooklib import epub
 
 # We need a QApplication instance for QWriter/QDialog

tests/test_book_parser.py

@@ -175,10 +175,11 @@ class TestBookParser(unittest.TestCase):
         """Test get_chapters returns correct list for different parsers."""
         # PDF
         parser_pdf = get_book_parser(self.sample_pdf_path)
+        parser_pdf.process_content()
         chapters = parser_pdf.get_chapters()
         self.assertEqual(len(chapters), 2)
-        self.assertEqual(chapters[0], ("page_1", "Page 1"))
+        self.assertEqual(chapters[0], ("page_1", "Page 1 - Page 1 content"))
-
+        
         # MD
         parser_md = get_book_parser(self.sample_md_path)
         parser_md.process_content()  # Must process to get structure

tests/test_pdf_structure.py

@@ -0,0 +1,117 @@
+import unittest
+import os
+import shutil
+import fitz  # PyMuPDF
+from abogen.book_parser import PdfParser
+
+class TestPdfStructure(unittest.TestCase):
+
+    def setUp(self):
+        self.test_dir = "tests/test_data_pdf"
+        if os.path.exists(self.test_dir):
+            shutil.rmtree(self.test_dir)
+        os.makedirs(self.test_dir)
+        self.pdf_path = os.path.join(self.test_dir, "structure_test.pdf")
+
+    def tearDown(self):
+        if os.path.exists(self.test_dir):
+            shutil.rmtree(self.test_dir)
+
+    def test_pdf_structure_with_toc(self):
+        # Create PDF
+        doc = fitz.open()
+        p1 = doc.new_page()
+        p1.insert_text((50,50), "Page 1 Content")
+        p2 = doc.new_page()
+        p2.insert_text((50,50), "Page 2 Content")
+        p3 = doc.new_page() # Chapter 2 start
+        p3.insert_text((50,50), "Page 3 Content")
+        p4 = doc.new_page()
+        p4.insert_text((50,50), "Page 4 Content")
+        
+        # Add TOC:
+        # 1. "Chap 1" -> Page 1
+        # 2. "Chap 2" -> Page 3
+        doc.set_toc([[1, "Chap 1", 1], [1, "Chap 2", 3]])
+        doc.save(self.pdf_path)
+        doc.close()
+
+        parser = PdfParser(self.pdf_path)
+        parser.process_content()
+        
+        nav = parser.processed_nav_structure
+        
+        # Expect 2 top level items
+        self.assertEqual(len(nav), 2)
+        self.assertEqual(nav[0]['title'], "Chap 1")
+        self.assertEqual(nav[0]['src'], "page_1")
+        self.assertEqual(nav[1]['title'], "Chap 2")
+        self.assertEqual(nav[1]['src'], "page_3")
+        
+        # Check children of Chap 1 (Page 2 should be there)
+        children_c1 = nav[0]['children']
+        self.assertEqual(len(children_c1), 1)
+        # The child should likely be titled "Page 2 - Page 2 Content" or similar
+        self.assertIn("Page 2", children_c1[0]['title'])
+        self.assertEqual(children_c1[0]['src'], "page_2")
+
+        # Check children of Chap 2 (Page 4 should be there)
+        children_c2 = nav[1]['children']
+        self.assertEqual(len(children_c2), 1)
+        self.assertIn("Page 4", children_c2[0]['title'])
+        self.assertEqual(children_c2[0]['src'], "page_4")
+
+    def test_pdf_structure_without_toc(self):
+        # Create PDF without TOC
+        doc = fitz.open()
+        p1 = doc.new_page()
+        p1.insert_text((50,50), "Start")
+        p2 = doc.new_page()
+        p2.insert_text((50,50), "End")
+        doc.save(self.pdf_path)
+        doc.close()
+
+        parser = PdfParser(self.pdf_path)
+        parser.process_content()
+        
+        nav = parser.processed_nav_structure
+        
+        # Expect 1 top level item (Pages)
+        self.assertEqual(len(nav), 1)
+        self.assertEqual(nav[0]['title'], "Pages")
+        
+        # Check children (all pages)
+        children = nav[0]['children']
+        self.assertEqual(len(children), 2)
+        self.assertIn("Page 1", children[0]['title'])
+        self.assertIn("Page 2", children[1]['title'])
+
+    def test_pdf_structure_nested_toc(self):
+        # Create PDF
+        doc = fitz.open()
+        p1 = doc.new_page() # Chap 1
+        p2 = doc.new_page() # Sec 1.1
+        p3 = doc.new_page() # Chap 2
+        
+        doc.set_toc([
+            [1, "Chap 1", 1],
+            [2, "Sec 1.1", 2],
+            [1, "Chap 2", 3]
+        ])
+        doc.save(self.pdf_path)
+        doc.close()
+
+        parser = PdfParser(self.pdf_path)
+        parser.process_content()
+        nav = parser.processed_nav_structure
+
+        self.assertEqual(len(nav), 2) # Chap 1, Chap 2
+        self.assertEqual(nav[0]['title'], "Chap 1")
+        
+        # Chap 1 should have child Sec 1.1
+        self.assertEqual(len(nav[0]['children']), 1)
+        self.assertEqual(nav[0]['children'][0]['title'], "Sec 1.1")
+        self.assertEqual(nav[0]['children'][0]['src'], "page_2")
+
+if __name__ == "__main__":
+    unittest.main()