mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 17:00:23 +00:00
89 lines
3.8 KiB
Python
89 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
|
|
def test_analyze_document_structure_flat_backmatter() -> None:
|
|
from paperforge.worker.ocr_document import DocumentStructure, PagePosition, analyze_document_structure
|
|
|
|
blocks = [
|
|
{"page": 1, "role": "body_paragraph", "text": "Intro"},
|
|
{"page": 2, "role": "body_paragraph", "text": "Methods"},
|
|
{"page": 3, "role": "body_paragraph", "text": "Results"},
|
|
{"page": 3, "role": "section_heading", "text": "Discussion"},
|
|
{"page": 4, "role": "body_paragraph", "text": "Discussion text"},
|
|
{"page": 5, "role": "backmatter_heading", "text": "Acknowledgments"},
|
|
{"page": 5, "role": "backmatter_body", "text": "Thanks."},
|
|
{"page": 6, "role": "reference_heading", "text": "References"},
|
|
{"page": 6, "role": "reference_item", "text": "[1] Ref A"},
|
|
]
|
|
|
|
doc = analyze_document_structure(blocks)
|
|
|
|
assert isinstance(doc, DocumentStructure)
|
|
assert doc.body_end_page == 4
|
|
assert doc.spread_start == 5
|
|
assert doc.spread_end == 6
|
|
assert doc.backmatter_form == "flat"
|
|
|
|
|
|
def test_analyze_document_structure_container_backmatter() -> None:
|
|
from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure
|
|
|
|
blocks = [
|
|
{"page": 1, "role": "body_paragraph", "text": "Intro"},
|
|
{"page": 2, "role": "section_heading", "text": "Methods"},
|
|
{"page": 3, "role": "body_paragraph", "text": "Results"},
|
|
{"page": 3, "role": "section_heading", "text": "Discussion"},
|
|
{"page": 4, "role": "body_paragraph", "text": "Discussion text"},
|
|
{"page": 5, "role": "backmatter_boundary_heading", "text": "ADDITIONAL INFORMATION AND DECLARATIONS"},
|
|
{"page": 5, "role": "backmatter_heading", "text": "Funding"},
|
|
{"page": 5, "role": "backmatter_body", "text": "Funded by NIH."},
|
|
{"page": 6, "role": "backmatter_heading", "text": "Acknowledgments"},
|
|
{"page": 6, "role": "backmatter_body", "text": "Thanks."},
|
|
{"page": 6, "role": "backmatter_heading", "text": "Author Contributions"},
|
|
{"page": 6, "role": "backmatter_body", "text": "AB did X."},
|
|
{"page": 7, "role": "reference_heading", "text": "References"},
|
|
{"page": 7, "role": "reference_item", "text": "[1] Ref A"},
|
|
]
|
|
|
|
doc = analyze_document_structure(blocks)
|
|
|
|
assert isinstance(doc, DocumentStructure)
|
|
assert doc.backmatter_form == "container"
|
|
|
|
|
|
def test_analyze_document_structure_no_tail() -> None:
|
|
from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure
|
|
|
|
blocks = [
|
|
{"page": 1, "role": "body_paragraph", "text": "Only body"},
|
|
{"page": 2, "role": "body_paragraph", "text": "Still body"},
|
|
]
|
|
|
|
doc = analyze_document_structure(blocks)
|
|
|
|
assert isinstance(doc, DocumentStructure)
|
|
assert doc.body_end_page is None
|
|
assert doc.spread_start is None
|
|
assert doc.spread_end is None
|
|
assert doc.backmatter_form == "flat"
|
|
|
|
|
|
def test_analyze_document_structure_mixed_tail_spread() -> None:
|
|
from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure
|
|
|
|
blocks = [
|
|
{"page": 1, "role": "body_paragraph", "text": "Intro"},
|
|
{"page": 2, "role": "body_paragraph", "text": "Results"},
|
|
{"page": 3, "role": "body_paragraph", "text": "Continued results"},
|
|
{"page": 3, "role": "backmatter_heading", "text": "Acknowledgments"},
|
|
{"page": 3, "role": "backmatter_body", "text": "Thanks."},
|
|
{"page": 4, "role": "body_paragraph", "text": "More results interleaved"},
|
|
{"page": 4, "role": "reference_heading", "text": "References"},
|
|
{"page": 4, "role": "reference_item", "text": "[1] Ref A"},
|
|
]
|
|
|
|
doc = analyze_document_structure(blocks)
|
|
|
|
assert isinstance(doc, DocumentStructure)
|
|
assert doc.spread_start is not None
|
|
assert doc.spread_end is not None
|