99 lines
4.6 KiB
Python
99 lines
4.6 KiB
Python
"""Tests for PDF text merge and heading classification."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
from rag_cut.layout_meta import enrich_layout_metadata
|
|
from rag_cut.models import Block, BlockType
|
|
from rag_cut.parsers.pdf.layout import TextLine, _merge_lines_to_paragraphs
|
|
from rag_cut.parsers.pdf.text_extract import MAX_HEADING_CHARS, _font_heading_level, _looks_like_ui_label
|
|
from rag_cut.parsers.pdf.noise_filter import is_margin_noise_block
|
|
from rag_cut.splitters.heading_splitter import is_heading_block, normalize_heading_block, numbered_heading_level
|
|
|
|
|
|
class ParagraphMergeTest(unittest.TestCase):
|
|
def test_does_not_merge_title_and_body_with_different_font_sizes(self) -> None:
|
|
lines = [
|
|
TextLine(x0=90, y0=90, x1=400, y1=110, text="Master Collection/Pattern Scanning", font_size=18),
|
|
TextLine(x0=90, y0=115, x1=480, y1=132, text="及所有大師/形態/策略指標的位置", font_size=17.5),
|
|
TextLine(
|
|
x0=90,
|
|
y0=145,
|
|
x1=500,
|
|
y1=200,
|
|
text="頁面左側【篩選】欄下方即是【Master Collection】文件夾。點擊此檔夾即可出現下拉頁面。",
|
|
font_size=11,
|
|
),
|
|
]
|
|
merged = _merge_lines_to_paragraphs(lines, page_width=600)
|
|
texts = [ln.text for ln in merged]
|
|
self.assertEqual(len(merged), 2)
|
|
self.assertIn("Master Collection", texts[0])
|
|
self.assertTrue(texts[1].startswith("頁面左側"))
|
|
self.assertNotIn("頁面左側", texts[0])
|
|
|
|
|
|
class FontHeadingTest(unittest.TestCase):
|
|
def test_short_large_text_is_heading(self) -> None:
|
|
level = _font_heading_level(18, 11, "Master Collection 位置說明")
|
|
self.assertIsNotNone(level)
|
|
|
|
def test_title_plus_body_blob_is_not_heading(self) -> None:
|
|
blob = (
|
|
"Master Collection/Pattern Scanning/Strategy Scanning 及所有大師/形態/策略指標的位置 "
|
|
"頁面左側【篩選】欄下方即是【Master Collection】文件夾。點擊此檔夾即可出現下拉頁面，"
|
|
"其中包含了十五個大師的選股策略。如圖所示："
|
|
)
|
|
self.assertGreater(len(blob), MAX_HEADING_CHARS)
|
|
self.assertIsNone(_font_heading_level(18, 11, blob))
|
|
|
|
def test_compact_numbered_chinese_title_is_heading(self) -> None:
|
|
# FAQ style: "10.上升三角形態" with no space after the dot.
|
|
level = _font_heading_level(16, 12, "10.上升三角形態")
|
|
self.assertEqual(level, 1)
|
|
self.assertFalse(_looks_like_ui_label("10.上升三角形態"))
|
|
|
|
def test_cjk_section_banner_is_level1_not_ui_label(self) -> None:
|
|
self.assertFalse(_looks_like_ui_label("形態指標"))
|
|
self.assertEqual(_font_heading_level(16, 12, "形態指標"), 1)
|
|
|
|
|
|
|
|
class HeadingDemoteTest(unittest.TestCase):
|
|
def test_normalize_demotes_overlong_heading(self) -> None:
|
|
blob = "T" * (MAX_HEADING_CHARS + 20)
|
|
block = Block(type=BlockType.HEADING, text=blob, level=1, meta={"page": 1})
|
|
normalized = normalize_heading_block(block)
|
|
self.assertEqual(normalized.type, BlockType.PARAGRAPH)
|
|
self.assertFalse(is_heading_block(normalized))
|
|
|
|
def test_enrich_demotes_overlong_heading_block(self) -> None:
|
|
blob = (
|
|
"Master Collection/Pattern Scanning/Strategy Scanning Master Collection/"
|
|
"Pattern Scanning/Strategy Scanning 及所有大師 /形態/策略指標的位置 "
|
|
"頁面左側【篩選】欄下方即是【Master Collection】文件夾。點擊此檔夾即可出 "
|
|
"現下拉頁面，其中包含了十五個大師的選股策略。"
|
|
)
|
|
blocks = [
|
|
Block(type=BlockType.HEADING, text=blob, level=1, meta={"page": 1, "bbox": [90, 90, 500, 240]}),
|
|
Block(type=BlockType.IMAGE, image_id="a.png", image_path="a.png", meta={"page": 1, "bbox": [90, 250, 500, 480]}),
|
|
]
|
|
enriched = enrich_layout_metadata(blocks)
|
|
self.assertEqual(enriched[0].type, BlockType.PARAGRAPH)
|
|
self.assertNotEqual(enriched[1].meta.get("bound_heading"), blob)
|
|
|
|
def test_top_of_page_numbered_title_not_margin_noise(self) -> None:
|
|
block = Block(
|
|
type=BlockType.PARAGRAPH,
|
|
text="10.上升三角形態",
|
|
meta={"page": 50, "bbox": [90, 74, 209, 90], "page_height": 842},
|
|
)
|
|
self.assertFalse(is_margin_noise_block(block, 842))
|
|
self.assertEqual(numbered_heading_level("10.上升三角形態"), 1)
|
|
self.assertTrue(is_heading_block(normalize_heading_block(block)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|