146 lines
8.2 KiB
Python
146 lines
8.2 KiB
Python
"""Lightweight import: selection, provenance, OCR routing, reflow and EPUB links."""
|
|
import json
|
|
import zipfile
|
|
from pathlib import Path
|
|
import pytest
|
|
from src.pdf_import.converter import convert, page_selection, validate_lines
|
|
from src.pdf_import.layout import paragraphs, repeated_furniture, join_text
|
|
from src.qa.epub_audit import audit_epub
|
|
|
|
|
|
@pytest.fixture
|
|
def native_pdf(tmp_path):
|
|
pytest.importorskip('pypdfium2')
|
|
# Minimal native-text PDF fixture; no external PDF authoring dependency.
|
|
def write(name='input.pdf', image=False):
|
|
streams=[b'BT /F1 18 Tf 60 740 Td (Chapter 1: Start) Tj 0 -40 Td /F1 12 Tf (A paragraph with enough native English text to use without OCR.) Tj ET',b'BT /F1 12 Tf 60 740 Td (This is the second page with more native English prose.) Tj ET']
|
|
if image:streams=[b'q 600 0 0 800 0 0 cm /Im1 Do Q '+s for s in streams]
|
|
resources=b'<< /Font << /F1 7 0 R >>'+(b' /XObject << /Im1 8 0 R >>' if image else b'')+b' >>'
|
|
objects=[b'<< /Type /Catalog /Pages 2 0 R >>',b'<< /Type /Pages /Kids [3 0 R 4 0 R] /Count 2 >>']
|
|
objects.extend(b'<< /Type /Page /Parent 2 0 R /MediaBox [0 0 600 800] /Resources '+resources+b' /Contents '+str(n).encode()+b' 0 R >>' for n in [5,6])
|
|
objects.extend(b'<< /Length '+str(len(s)).encode()+b' >>\nstream\n'+s+b'\nendstream' for s in streams)
|
|
objects.append(b'<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>')
|
|
if image:objects.append(b'<< /Type /XObject /Subtype /Image /Width 1 /Height 1 /ColorSpace /DeviceRGB /BitsPerComponent 8 /Length 3 >>\nstream\n\xff\xff\xff\nendstream')
|
|
raw=b'%PDF-1.4\n';offsets=[0]
|
|
for i,obj in enumerate(objects):offsets.append(len(raw));raw+=f'{i+1} 0 obj\n'.encode()+obj+b'\nendobj\n'
|
|
xref=len(raw);raw+=f'xref\n0 {len(objects)+1}\n0000000000 65535 f \n'.encode()
|
|
raw+=b''.join(f'{n:010d} 00000 n \n'.encode() for n in offsets[1:])
|
|
raw+=f'trailer\n<< /Size {len(objects)+1} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF'.encode()
|
|
path=tmp_path/name;path.write_bytes(raw);return path
|
|
return write
|
|
|
|
|
|
def test_page_ranges_do_not_silently_drop_pages():
|
|
assert page_selection(None,3)==[0,1,2]
|
|
assert page_selection('3,1-2,2',3)==[0,1,2]
|
|
for value in ['0','3-2','1-4','-1','1,','abc']:
|
|
with pytest.raises(ValueError):page_selection(value,3)
|
|
|
|
|
|
def test_folio_ocr_variants_are_removed_but_body_years_survive():
|
|
lines=[{'id':'header','text':'002 | 中国的世界系','bbox':[.15,.04,.4,.06]}, {'id':'body','text':'2019年发生的事。','bbox':[.15,.12,.8,.14]}]
|
|
blocks,removed=paragraphs({'number':5,'lines':lines},'zh-CN',set())
|
|
assert [b['text'] for b in blocks]==['2019年发生的事。']
|
|
assert removed[0]['id']=='header'
|
|
assert join_text('跨页','继续。','zh-CN')=='跨页继续。'
|
|
assert join_text('Some','words','en')=='Some words'
|
|
|
|
|
|
def test_native_english_import_is_independent_and_complete(native_pdf,tmp_path):
|
|
source=native_pdf();out=tmp_path/'book.epub';work=tmp_path/'work'
|
|
report=convert(source,out,work,language='en',engine='native')
|
|
assert report['engines']=={'native':2} and not report['partial']
|
|
assert audit_epub(out)['issues']==[]
|
|
with zipfile.ZipFile(out) as z:
|
|
chapters=[n for n in z.namelist() if '/chapters/' in n]
|
|
assert len(chapters)==1 # Consecutive pages flow inside one reading document.
|
|
text=z.read(chapters[0]).decode()
|
|
assert 'native English' in text and 'second page' in text
|
|
assert 'pdf-page-1' in text and 'pdf-page-2' in text
|
|
assert '../styles/book.css' in text and 'lang="en"' in text
|
|
from src.preprocessing.epub_cleaner import EpubCleaner
|
|
from src.preprocessing.text_extractor import FineGrainedExtractor
|
|
from src.common.data_model import BookStructure
|
|
structure=BookStructure.load(EpubCleaner(out,tmp_path/'translation-input').clean())
|
|
units=FineGrainedExtractor(complete_book=True).extract(structure)
|
|
assert any('native English' in unit.original_text for unit in units)
|
|
ir=json.loads((work/'book.json').read_text())
|
|
assert len(ir['pages'])==2 and ir['pages'][0]['lines'][0]['bbox']
|
|
partial=convert(source,tmp_path/'partial.epub',work,language='en',engine='native',pages='2')
|
|
assert partial['partial'] and partial['processed_pages']==1
|
|
with zipfile.ZipFile(tmp_path/'partial.epub') as z:
|
|
opf=z.read('EPUB/content.opf').decode();assert '部分样例' in opf
|
|
with pytest.raises(ValueError,match='another source'):
|
|
convert(native_pdf('another.pdf',image=True),tmp_path/'other.epub',work,engine='native')
|
|
|
|
|
|
def test_scanned_text_layer_is_replaced_and_ocr_cache_is_reused(native_pdf,tmp_path,monkeypatch):
|
|
from src.pdf_import import ocr
|
|
import src.pdf_import.converter as converter
|
|
called=[]
|
|
class FakeOCR:
|
|
identity='fixture-v1'
|
|
def __init__(self,*args):pass
|
|
def recognize(self,image):
|
|
called.append(image)
|
|
return [{'text':'这是来自扫描图像的新识别结果，必须保留完整正文内容，不能使用已有错误文本层。','bbox':[.1,.2,.9,.25],'confidence':.9}]
|
|
def close(self):pass
|
|
monkeypatch.setattr(ocr,'VisionOCR',FakeOCR)
|
|
monkeypatch.setattr(ocr,'TesseractOCR',FakeOCR)
|
|
source=native_pdf(image=True);out=tmp_path/'scan.epub';work=tmp_path/'work'
|
|
convert(source,out,work,language='zh-CN',engine='auto')
|
|
assert len(called)==2
|
|
convert(source,out,work,language='zh-CN',engine='auto')
|
|
assert len(called)==2
|
|
with zipfile.ZipFile(out) as z:
|
|
body=b''.join(z.read(n) for n in z.namelist() if '/chapters/' in n).decode()
|
|
assert '新识别结果' in body and 'native English' not in body
|
|
assert 'lang="zh-CN"' in body
|
|
FakeOCR.identity='fixture-v2'
|
|
convert(source,out,work,language='zh-CN',engine='auto')
|
|
assert len(called)==4
|
|
|
|
|
|
def test_invalid_ocr_geometry_fails_closed():
|
|
with pytest.raises(ValueError):validate_lines([{'text':'bad','bbox':[.8,.2,.1,.4]}])
|
|
with pytest.raises(ValueError):validate_lines([{'text':'bad','bbox':[0,0,float('nan'),1]}])
|
|
|
|
|
|
def test_ocr_failure_leaves_page_checkpoint_and_no_partial_final(native_pdf,tmp_path,monkeypatch):
|
|
from src.pdf_import import ocr
|
|
called=[]
|
|
class FakeOCR:
|
|
identity='resume-fixture'
|
|
fail=True
|
|
def __init__(self,*args):pass
|
|
def recognize(self,image):
|
|
called.append(Path(image).name)
|
|
if Path(image).stem=='0002' and self.fail:raise RuntimeError('fixture OCR failure')
|
|
return [{'text':'A sufficiently long recognized paragraph to remain searchable in the EPUB.', 'bbox':[.1,.2,.9,.3], 'confidence':1}]
|
|
def close(self):pass
|
|
monkeypatch.setattr(ocr,'VisionOCR',FakeOCR);monkeypatch.setattr(ocr,'TesseractOCR',FakeOCR)
|
|
source=native_pdf(image=True);out=tmp_path/'output.epub';work=tmp_path/'work'
|
|
with pytest.raises(RuntimeError,match='fixture OCR failure'):convert(source,out,work,language='en')
|
|
assert not out.exists() and (work/'pages/0001.json').exists() and not (work/'pages/0002.json').exists()
|
|
FakeOCR.fail=False
|
|
convert(source,out,work,language='en')
|
|
assert called==['0001.jpg','0002.jpg','0002.jpg']
|
|
assert audit_epub(out)['issues']==[]
|
|
|
|
|
|
def test_consistently_indented_quote_lines_stay_one_paragraph():
|
|
lines=[{'id':'a','text':'这是正文段落。','bbox':[.1,.11,.9,.13]}, {'id':'b','text':'一段尚未结束的引文','bbox':[.2,.16,.8,.18]}, {'id':'c','text':'继续在相同缩进位置。','bbox':[.2,.19,.8,.21]}]
|
|
blocks,_=paragraphs({'number':1,'lines':lines},'zh-CN',set())
|
|
assert len(blocks)==2
|
|
assert blocks[1]['text']=='一段尚未结束的引文继续在相同缩进位置。'
|
|
|
|
|
|
def test_corrections_require_source_and_exact_expected_text():
|
|
from src.pdf_import.converter import apply_corrections
|
|
pages=[{'lines':[{'id':'p0001-l0001','text':'牲富'}]}]
|
|
patch={'source_sha256':'source','edits':[{'line_id':'p0001-l0001','expected':'牲富','replacement':'牲畜'}]}
|
|
with pytest.raises(ValueError):apply_corrections(pages,patch,'different-source')
|
|
applied=apply_corrections(pages,patch,'source')
|
|
assert pages[0]['lines'][0]['text']=='牲畜' and len(applied)==1
|
|
with pytest.raises(ValueError,match='Stale'):apply_corrections(pages,patch,'source')
|