mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 13:40:27 +02:00
Merge pull request #126 from mohangk/webui
Further refactoring of book_handler logic into PDFParser
This commit is contained in:
Deniz Şafak
File diff suppressed because it is too large
Load Diff
+132
-15
@@ -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):
|
||||
# PDF metadata extraction can be added here if needed
|
||||
# For now, base class metadata is empty dict
|
||||
pass
|
||||
def _get_page_title(self, page_num, text):
|
||||
title = f"Page {page_num + 1}"
|
||||
if text:
|
||||
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):
|
||||
# PDF specific implementation because it doesn't use nav structure
|
||||
chapters = []
|
||||
if self.pdf_doc:
|
||||
for i in range(len(self.pdf_doc)):
|
||||
chapters.append((f"page_{i+1}", f"Page {i+1}"))
|
||||
return chapters
|
||||
def _build_structure_from_toc(self, toc):
|
||||
# 1. Flatten TOC to easier list (page_num, title, level)
|
||||
# fitz TOC is [[lvl, title, page, dest], ...]
|
||||
|
||||
bookmarks = []
|
||||
for entry in toc:
|
||||
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):
|
||||
|
||||
+3
-212
@@ -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.parser.file_type == "pdf":
|
||||
self._build_pdf_tree()
|
||||
elif self.processed_nav_structure:
|
||||
if 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,15 +800,12 @@ class HandlerDialog(QDialog):
|
||||
continue
|
||||
|
||||
identifier = item.data(0, Qt.ItemDataRole.UserRole)
|
||||
|
||||
if not identifier:
|
||||
iterator += 1
|
||||
continue
|
||||
|
||||
if (
|
||||
not identifier.startswith("page_")
|
||||
or self.content_lengths.get(identifier, 0) > 0
|
||||
):
|
||||
# Logic: Check item if it has content (already handled by ItemIsUserCheckable flag really)
|
||||
# But duplicate logic from previous implementation:
|
||||
item.setCheckState(0, Qt.CheckState.Checked)
|
||||
|
||||
iterator += 1
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -175,9 +175,10 @@ 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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user