chenxin/hunjing
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

初始化提交

Browse Source
This commit is contained in:
chenxin committed 2026-08-21 16:00:13 +08:00
commit 887c1f449e
963 files changed
+166869
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tianmu/tests/test_layout.py
+531
View File
@@ -0,0 +1,531 @@
"""
布局后处理单元测试（纯逻辑，不依赖引擎和服务）。
测试 `src.engine.layout` 中的纯函数：
- sort_and_group：阅读顺序排序 + 行归组
- group_paragraphs：段落分组
- build_markdown_text：Markdown 格式化
- 辅助函数
"""
from __future__ import annotations
import pytest
from src.engine.layout import (
_block_center_x,
_block_center_y,
_block_height,
_detect_columns,
_detect_heading_level,
_detect_line_type,
_sort_and_group_single,
build_markdown_text,
group_paragraphs,
sort_and_group,
)
from src.schemas.ocr import BBox, TextBlock
# ── 工厂函数 ──────────────────────────────────────────────
def _make_block(
text: str,
x1: float = 0,
y1: float = 0,
x2: float = 100,
y2: float = 30,
confidence: float = 0.95,
) -> TextBlock:
"""创建一个测试用的 TextBlock。"""
return TextBlock(
text=text,
confidence=confidence,
bbox=BBox(points=[[x1, y1], [x2, y1], [x2, y2], [x1, y2]]),
)
# ═══════════════════════════════════════════════════════════
# 辅助函数
# ═══════════════════════════════════════════════════════════
class TestHelpers:
"""辅助函数测试。"""
def test_block_height(self) -> None:
block = _make_block("测试", y1=10, y2=50)
assert _block_height(block) == 40.0
def test_block_height_empty(self) -> None:
block = _make_block("测试", x1=0, y1=0, x2=0, y2=0)
assert _block_height(block) == 0.0
def test_block_center_y(self) -> None:
block = _make_block("测试", y1=10, y2=50)
assert _block_center_y(block) == 30.0
def test_block_center_x(self) -> None:
block = _make_block("测试", x1=20, x2=100)
assert _block_center_x(block) == 60.0
def test_detect_heading_chinese(self) -> None:
"""中文短行 → heading。"""
assert _detect_line_type("第一章 绪论") == "heading"
def test_detect_heading_short_no_punct(self) -> None:
"""短行不以句号结尾 → heading。"""
assert _detect_line_type("摘要") == "heading"
def test_detect_not_heading_with_punct(self) -> None:
"""以句号结尾的短行 → text。"""
assert _detect_line_type("这是一个句子。") == "text"
def test_detect_long_text(self) -> None:
"""长行（>= 40 字符）→ text。"""
long = "一二三四五六七八九十" * 4
assert _detect_line_type(long) == "text"
def test_detect_list_item_dash(self) -> None:
assert _detect_line_type("- 列表项") == "list_item"
def test_detect_list_item_bullet(self) -> None:
assert _detect_line_type("• 项目") == "list_item"
def test_detect_numbered(self) -> None:
assert _detect_line_type("1. 首先") == "numbered"
assert _detect_line_type("2) 其次") == "numbered"
def test_detect_separator(self) -> None:
assert _detect_line_type("---") == "separator"
assert _detect_line_type("===") == "separator"
assert _detect_line_type("***") == "separator"
def test_detect_code(self) -> None:
assert _detect_line_type(" def foo():") == "code"
assert _detect_line_type("\t\tindented") == "code"
def test_detect_blockquote_as_text(self) -> None:
"""> 前缀在 OCR 语境中通常是排版伪影，不作为代码块。"""
assert _detect_line_type("> 引用") == "text"
def test_detect_empty_text(self) -> None:
assert _detect_line_type("") == "text"
assert _detect_line_type(" ") == "text"
def test_detect_centered_heading(self) -> None:
"""居中短行 → 即使非中文也应视为 heading。"""
assert _detect_line_type("Hello World", center_x_ratio=0.5) == "heading"
def test_detect_not_centered(self) -> None:
"""不居中短行 + 非中文锁 → 不回退。"""
# "hello world" 非中文、非大写、不居中 → text
assert _detect_line_type("hello world", center_x_ratio=0.1) == "text"
def test_detect_uppercase_heading(self) -> None:
"""全大写短行 → heading。"""
assert _detect_line_type("INTRODUCTION") == "heading"
# ── 新增：编号变体 ────────────────────────────────
def test_detect_numbered_without_space(self) -> None:
"""编号后无空格（"1.内容"）→ numbered。"""
assert _detect_line_type("1.内容") == "numbered"
assert _detect_line_type("1)内容") == "numbered"
def test_detect_numbered_parenthesized(self) -> None:
"""括号编号 → numbered。"""
assert _detect_line_type("(1) 内容") == "numbered"
assert _detect_line_type("(2)内容") == "numbered"
def test_detect_numbered_chinese_numeral(self) -> None:
"""中文数字序号 → numbered。"""
assert _detect_line_type("① 第一项") == "numbered"
assert _detect_line_type("②第二项") == "numbered"
# ── 新增：TOC 目录行 ──────────────────────────────
def test_detect_toc_line(self) -> None:
"""TOC 目录行（含连续点号）→ numbered。"""
assert _detect_line_type("第一章 .............. 1") == "numbered"
assert _detect_line_type("1.1 概述 ............ 10") == "numbered"
# ── 新增：更多 bullet 变体 ────────────────────────
def test_detect_bullet_variants(self) -> None:
"""更多 bullet 变体 → list_item。"""
assert _detect_line_type("▪ 项目") == "list_item"
assert _detect_line_type("▸ 子项") == "list_item"
assert _detect_line_type("→ 链接") == "list_item"
assert _detect_line_type("· 条目") == "list_item"
def test_detect_standalone_bullet(self) -> None:
"""独立 bullet 字符 → list_item。"""
assert _detect_line_type("•") == "list_item"
assert _detect_line_type("▪") == "list_item"
assert _detect_line_type("→") == "list_item"
# ── 新增：字体大小辅助标题检测 ────────────────────
def test_detect_larger_font_heading(self) -> None:
"""大字（rel_height > 1.3）→ heading。"""
assert _detect_line_type("这是一段长文本但字号大", rel_height=1.5) == "heading"
# 中号字（rel_height < 1.3）且不居中 → text
assert _detect_line_type("这是一段长文本字号普通", rel_height=1.1) == "text"
# ═══════════════════════════════════════════════════════════
# 标题级别检测
# ═══════════════════════════════════════════════════════════
class TestHeadingLevel:
"""标题级别检测。"""
def test_h1_centered_large(self) -> None:
assert _detect_heading_level(rel_height=1.5, is_centered=True) == 1
def test_h1_large_alone(self) -> None:
assert _detect_heading_level(rel_height=1.8) == 1
def test_h2_large(self) -> None:
assert _detect_heading_level(rel_height=1.4) == 2
def test_h2_centered(self) -> None:
assert _detect_heading_level(is_centered=True) == 2
def test_h3_default(self) -> None:
assert _detect_heading_level() == 3
assert _detect_heading_level(rel_height=1.1) == 3
# ═══════════════════════════════════════════════════════════
# sort_and_group
# ═══════════════════════════════════════════════════════════
class TestSortAndGroup:
"""阅读顺序排序 + 行归组。"""
def test_empty(self) -> None:
assert sort_and_group([]) == []
def test_single_block(self) -> None:
blocks = [_make_block("你好")]
lines = sort_and_group(blocks)
assert len(lines) == 1
assert len(lines[0]) == 1
assert lines[0][0].text == "你好"
def test_top_to_bottom_sort(self) -> None:
"""按 Y 坐标从上到下排序。"""
blocks = [
_make_block("第二行", y1=100, y2=130),
_make_block("第一行", y1=10, y2=40),
]
lines = sort_and_group(blocks)
assert len(lines) == 2
assert lines[0][0].text == "第一行"
assert lines[1][0].text == "第二行"
def test_same_line_grouping(self) -> None:
"""Y 坐标接近的 blocks 归为同一行，按 X 排序。"""
blocks = [
_make_block("右", x1=100, x2=130, y1=10, y2=30),
_make_block("左", x1=0, x2=30, y1=10, y2=30),
]
lines = sort_and_group(blocks)
assert len(lines) == 1
assert len(lines[0]) == 2
assert lines[0][0].text == "左"
assert lines[0][1].text == "右"
def test_multi_line_grouping(self) -> None:
"""多行正确归组。"""
blocks = [
_make_block("B1", x1=0, x2=30, y1=100, y2=130),
_make_block("A1", x1=0, x2=30, y1=10, y2=40),
_make_block("A2", x1=40, x2=70, y1=10, y2=40),
]
lines = sort_and_group(blocks)
assert len(lines) == 2
# 第一行：A1、A2
assert len(lines[0]) == 2
assert lines[0][0].text == "A1"
assert lines[0][1].text == "A2"
# 第二行：B1
assert len(lines[1]) == 1
assert lines[1][0].text == "B1"
# ═══════════════════════════════════════════════════════════
# 多栏检测
# ═══════════════════════════════════════════════════════════
class TestMultiColumn:
"""多栏检测。"""
def assert_one_column(self, blocks: list, page_width: float = 800) -> None:
cols = _detect_columns(blocks, page_width)
assert len(cols) == 1, f"应检测为单栏，实际 {len(cols)} 栏"
def assert_two_columns(self, blocks: list, page_width: float = 800) -> None:
cols = _detect_columns(blocks, page_width)
assert len(cols) == 2, f"应检测为双栏，实际 {len(cols)} 栏"
assert len(cols[0]) > 0
assert len(cols[1]) > 0
def test_single_column_few_blocks(self) -> None:
"""少于 4 个 block 不触发多栏检测。"""
self.assert_one_column([_make_block("A")])
def test_single_column_no_gap(self) -> None:
"""无间隙 → 单栏。"""
blocks = [
_make_block("左1", x1=0, x2=100, y1=10, y2=30),
_make_block("右1", x1=100, x2=200, y1=10, y2=30),
_make_block("左2", x1=0, x2=100, y1=40, y2=60),
_make_block("右2", x1=100, x2=200, y1=40, y2=60),
]
self.assert_one_column(blocks)
def test_two_columns_clear_gap(self) -> None:
"""有明显水平间隙 → 双栏。"""
blocks = [
_make_block("左栏 A", x1=0, x2=150, y1=10, y2=30),
_make_block("左栏 B", x1=0, x2=150, y1=40, y2=60),
_make_block("右栏 A", x1=500, x2=650, y1=15, y2=35),
_make_block("右栏 B", x1=500, x2=650, y1=50, y2=70),
]
self.assert_two_columns(blocks, page_width=800)
def test_sort_and_group_two_columns(self) -> None:
"""双栏排版：左栏两行，右栏两行。"""
blocks = [
_make_block("右A", x1=500, x2=650, y1=50, y2=70),
_make_block("左B", x1=0, x2=150, y1=40, y2=60),
_make_block("左A", x1=0, x2=150, y1=10, y2=30),
_make_block("右B", x1=500, x2=650, y1=80, y2=100),
]
lines = sort_and_group(blocks, page_width=800)
# 左栏两行 → 右栏两行
assert len(lines) == 4
# 左栏第一行
assert lines[0][0].text == "左A"
# 左栏第二行
assert lines[1][0].text == "左B"
# 右栏第一行
assert lines[2][0].text == "右A"
# 右栏第二行
assert lines[3][0].text == "右B"
# ═══════════════════════════════════════════════════════════
# group_paragraphs
# ═══════════════════════════════════════════════════════════
class TestGroupParagraphs:
"""段落分组。"""
def test_empty(self) -> None:
assert group_paragraphs([]) == []
def test_single_para(self) -> None:
"""多行同段落（用长文本避免触发标题检测）。"""
lines = [
[_make_block("第一行正文内容这里是长文本")],
[_make_block("第二行正文内容再接一些字")],
]
paras = group_paragraphs(lines)
assert len(paras) == 1
assert paras[0]["type"] == "text"
assert paras[0]["type"] == "text"
def test_paragraph_break_by_gap(self) -> None:
"""行间距大 → 分段（需要 3+ 行，单行间距无法超过自身中位值的 1.8 倍）。"""
lines = [
[_make_block("第一段正文在这里第一行啊", y1=10, y2=30)],
[_make_block("第一段正文在这里第二行", y1=40, y2=60)],
[_make_block("第二段正文这里间距很大", y1=200, y2=230)],
]
paras = group_paragraphs(lines, page_width=800)
assert len(paras) == 2
# 第一段：两行
assert paras[0]["type"] == "text"
assert "第一段正文在这里" in paras[0]["text"]
# 第二段：一行
assert "第二段正文" in paras[1]["text"]
def test_heading_paragraph(self) -> None:
"""短行 → 标题段。"""
lines = [[_make_block("第一章")], [_make_block("正文内容")]]
paras = group_paragraphs(lines)
assert len(paras) >= 1
# 标题段
assert paras[0]["type"] == "heading"
def test_list_paragraph(self) -> None:
"""列表项段落。"""
lines = [
[_make_block("- 项目一", y1=10, y2=30)],
[_make_block("- 项目二", y1=40, y2=60)],
]
paras = group_paragraphs(lines)
assert len(paras) == 1
assert paras[0]["type"] == "list_item"
# ═══════════════════════════════════════════════════════════
# build_markdown_text
# ═══════════════════════════════════════════════════════════
class TestBuildMarkdown:
"""完整的 Markdown 转换。"""
def test_empty(self) -> None:
md, lines, paras = build_markdown_text([])
assert md == ""
assert lines == []
assert paras == []
def test_simple_text(self) -> None:
"""简单段落转换。"""
blocks = [
_make_block("你好世界"),
]
md, lines, paras = build_markdown_text(blocks)
assert "你好世界" in md
def test_heading_conversion(self) -> None:
"""标题应加 # 前缀（级别基于字体大小）。"""
# 标题 block 更大 → 推断为 h2
blocks = [
_make_block("第一章", x1=0, x2=100, y1=10, y2=50), # 40px 高
_make_block("正文内容长一些避免被检测为标题", x1=0, x2=200, y1=60, y2=80), # 20px 高
]
md, lines, paras = build_markdown_text(blocks)
assert "### 第一章" in md or "## 第一章" in md
assert "正文内容长一些避免被检测为标题" in md
def test_list_conversion(self) -> None:
"""列表项应加 - 前缀。"""
blocks = [
_make_block("- 苹果"),
_make_block("- 香蕉"),
]
md, lines, paras = build_markdown_text(blocks)
assert "- 苹果" in md
assert "- 香蕉" in md
def test_centered_heading(self) -> None:
"""居中的短行 → heading。"""
# 中心 X 在页面宽度中间
blocks = [
_make_block("欢迎光临", x1=300, x2=500, y1=10, y2=40),
]
md, lines, paras = build_markdown_text(blocks, page_width=800)
assert "欢迎光临" in md
def test_heading_level_h1(self) -> None:
"""居中大字 → h1（# 前缀）。"""
blocks = [
_make_block("年度报告", x1=300, x2=500, y1=10, y2=50),
_make_block("正文内容较小", x1=0, x2=100, y1=60, y2=75),
]
md, lines, paras = build_markdown_text(blocks, page_width=800)
assert "# 年度报告" in md or "## 年度报告" in md
def test_heading_level_h3(self) -> None:
"""等大字体短行 → h3（### 前缀）。"""
blocks = [
_make_block("摘要", x1=0, x2=100, y1=10, y2=30),
_make_block("第一章内容文字", x1=0, x2=100, y1=40, y2=70),
]
md, lines, paras = build_markdown_text(blocks, page_width=800)
# "摘要" ≤ 8 字无标点 → heading, 但字体与正文同大 → h3
assert "### 摘要" in md
def test_separator_conversion(self) -> None:
"""分隔线 → ---。"""
blocks = [
_make_block("---"),
]
md, lines, paras = build_markdown_text(blocks)
assert "---" in md
def test_multi_line_same_para(self) -> None:
"""多行近距离内容合并在一个段落（长文本避免标题检测）。"""
blocks = [
_make_block("第一行长文本避免被检测为标题行", y1=10, y2=30),
_make_block("第二行避免被检测为标题", y1=35, y2=55),
]
md, lines, paras = build_markdown_text(blocks)
assert "第一行长文本避免被检测为标题行" in md
assert "第二行避免被检测为标题" in md
# ═══════════════════════════════════════════════════════════
# 边缘情况
# ═══════════════════════════════════════════════════════════
class TestEdgeCases:
"""边缘情况。"""
def test_single_block_only(self) -> None:
blocks = [_make_block("Hello")]
md, lines, paras = build_markdown_text(blocks)
assert md == "Hello"
def test_blocks_with_same_y(self) -> None:
"""完全同一 Y 位置的多个 blocks（罕见但需处理）。"""
blocks = [
_make_block("A", x1=0, x2=30, y1=10, y2=30),
_make_block("B", x1=40, x2=70, y1=10, y2=30),
]
lines = sort_and_group(blocks)
assert len(lines) == 1
assert len(lines[0]) == 2
def test_blocks_with_negligible_height(self) -> None:
"""高度为 0 的 block 不应崩。"""
block = TextBlock(
text="边缘",
confidence=0.5,
bbox=BBox(points=[[0, 0], [0, 0], [0, 0], [0, 0]]),
)
assert _block_height(block) == 0.0
lines = sort_and_group([block])
assert len(lines) == 1
def test_confidence_range(self) -> None:
"""置信度在 0~1 范围内。"""
block = _make_block("test", confidence=1.0)
assert 0 <= block.confidence <= 1
def test_long_line_content(self) -> None:
"""长行（> 40 字符）→ text 而非 heading。"""
long_text = "A" * 50
blocks = [_make_block(long_text)]
md, lines, paras = build_markdown_text(blocks)
assert long_text in md
assert "##" not in md
def test_empty_text_blocks(self) -> None:
"""空文本 block 不应造成影响。"""
blocks = [
_make_block(""),
_make_block("实际内容"),
]
md, lines, paras = build_markdown_text(blocks)
assert "实际内容" in md
def test_numeric_only(self) -> None:
"""纯数字行 → numbered。"""
blocks = [_make_block("42")]
# 短行数字 → numbered（先于 heading 触发）
md, lines, paras = build_markdown_text(blocks)
assert "42" in md