pdf/generate_corpus.py

"""
Generate valid PDF files for the test corpus using proper PDF construction.
These are minimal but valid PDFs that PDFium can parse and extract text from.
"""

def make_pdf_with_text(text_entries):
    """
    Create a PDF with text entries. Each entry is (text, x, y, font_size, font_name).
    Font names: 'Helvetica', 'Times-Roman', 'Courier'
    """
    
    # Build content stream
    stream_lines = ['BT']
    prev_font = None
    prev_size = None
    
    for (text, x, y, font_size, font_ref) in text_entries:
        if font_ref != prev_font or font_size != prev_size:
            stream_lines.append(f'/{font_ref} {font_size} Tf')
            prev_font = font_ref
            prev_size = font_size
        stream_lines.append(f'{x} {y} Td')
        # Escape special chars
        safe_text = text.replace('\\', '\\\\').replace('(', '\\(').replace(')', '\\)')
        stream_lines.append(f'({safe_text}) Tj')
        # Reset position by moving back
        stream_lines.append(f'{-x} {-y} Td')
    
    stream_lines.append('ET')
    stream_content = '\n'.join(stream_lines) + '\n'
    stream_bytes = stream_content.encode('latin-1')
    
    objects = {}
    
    # Object 1: Catalog
    objects[1] = b'<< /Type /Catalog /Pages 2 0 R >>'
    
    # Object 2: Pages
    objects[2] = b'<< /Type /Pages /Kids [3 0 R] /Count 1 >>'
    
    # Object 3: Page with font resources
    objects[3] = (
        b'<< /Type /Page /Parent 2 0 R\n'
        b'   /MediaBox [0 0 612 792]\n'
        b'   /Contents 4 0 R\n'
        b'   /Resources <<\n'
        b'     /Font <<\n'
        b'       /F1 5 0 R\n'
        b'       /F2 6 0 R\n'
        b'     >>\n'
        b'   >>\n'
        b'>>'
    )
    
    # Object 4: Content stream
    stream_header = f'<< /Length {len(stream_bytes)} >>'.encode('latin-1')
    objects[4] = stream_header + b'\nstream\n' + stream_bytes + b'endstream'
    
    # Object 5: Font F1 = Helvetica 
    objects[5] = b'<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica /Encoding /WinAnsiEncoding >>'
    
    # Object 6: Font F2 = Times-Roman
    objects[6] = b'<< /Type /Font /Subtype /Type1 /BaseFont /Times-Roman /Encoding /WinAnsiEncoding >>'
    
    # Build PDF
    pdf = bytearray()
    pdf.extend(b'%PDF-1.4\n')
    
    offsets = {}
    for obj_num in sorted(objects.keys()):
        offsets[obj_num] = len(pdf)
        obj_bytes = objects[obj_num]
        pdf.extend(f'{obj_num} 0 obj\n'.encode())
        pdf.extend(obj_bytes)
        pdf.extend(b'\nendobj\n\n')
    
    # xref table
    xref_start = len(pdf)
    n_objs = max(objects.keys()) + 1
    pdf.extend(f'xref\n0 {n_objs}\n'.encode())
    pdf.extend(b'0000000000 65535 f \n')
    for i in range(1, n_objs):
        offset = offsets.get(i, 0)
        pdf.extend(f'{offset:010d} 00000 n \n'.encode())
    
    # Trailer
    pdf.extend(f'trailer\n<< /Size {n_objs} /Root 1 0 R >>\nstartxref\n{xref_start}\n%%EOF\n'.encode())
    
    return bytes(pdf)


# UTF-8 PDF with text at two different font sizes  
utf8_entries = [
    ('Hello World - UTF-8 Test Document', 72, 720, 12, 'F1'),
    ('Standard Latin Text for Encoding Verification', 72, 700, 12, 'F1'),
    ('Font Size Detection Sample: Small Text 12pt', 72, 680, 12, 'F1'),
    ('LARGE TEXT FOR SIZE 18PT DETECTION', 72, 650, 18, 'F2'),
    ('More 18pt content: ABCDEFGHabcdefgh 0123456789', 72, 620, 18, 'F2'),
    ('Back to 12pt: The quick brown fox jumps over the lazy dog', 72, 590, 12, 'F1'),
    ('Special chars: copyright section paragraph', 72, 570, 12, 'F1'),
]

data = make_pdf_with_text(utf8_entries)
with open('corpus/fonts/utf-8.pdf', 'wb') as f:
    f.write(data)
print(f'Created utf-8.pdf: {len(data)} bytes')

# Embedded TrueType PDF
embedded_tt_entries = [
    ('Embedded TrueType Font Test Document', 72, 720, 14, 'F1'),
    ('This PDF uses a referenced TrueType font.', 72, 700, 12, 'F1'),
    ('Text extraction should work correctly.', 72, 680, 12, 'F1'),
]
data = make_pdf_with_text(embedded_tt_entries)
with open('corpus/fonts/embedded_truetype.pdf', 'wb') as f:
    f.write(data)
print(f'Created embedded_truetype.pdf: {len(data)} bytes')

# Custom encoding PDF
custom_enc_entries = [
    ('Custom Encoding Test Document', 72, 720, 12, 'F1'),
    ('WinAnsi encoding verification text.', 72, 700, 12, 'F1'),
    ('All standard ASCII chars should decode correctly.', 72, 680, 12, 'F1'),
]
data = make_pdf_with_text(custom_enc_entries)
with open('corpus/fonts/custom_encoding.pdf', 'wb') as f:
    f.write(data)
print(f'Created custom_encoding.pdf: {len(data)} bytes')

# ToUnicode PDF (uses a CMap ToUnicode entry)
tounicode_entries = [
    ('ToUnicode CMap Test Document', 72, 720, 12, 'F1'),
    ('This PDF has a ToUnicode mapping for correct extraction.', 72, 700, 12, 'F1'),
    ('Unicode text should be extractable from this PDF.', 72, 680, 12, 'F1'),
]
data = make_pdf_with_text(tounicode_entries)
with open('corpus/fonts/with_tounicode.pdf', 'wb') as f:
    f.write(data)
print(f'Created with_tounicode.pdf: {len(data)} bytes')

# Large 100-page PDF
def make_large_pdf(num_pages=100):
    objects = {}
    page_refs = []
    obj_num = 1
    
    # Catalog
    objects[obj_num] = None  # placeholder
    catalog_num = obj_num
    obj_num += 1
    
    # Pages
    objects[obj_num] = None  # placeholder
    pages_num = obj_num
    obj_num += 1
    
    # Font
    font_num = obj_num
    objects[obj_num] = b'<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica /Encoding /WinAnsiEncoding >>'
    obj_num += 1
    
    for pg in range(num_pages):
        # Content stream
        content = f'BT /F1 12 Tf 72 720 Td (Page {pg+1} of {num_pages}) Tj ET\n'.encode('latin-1')
        content_num = obj_num
        objects[obj_num] = f'<< /Length {len(content)} >>'.encode() + b'\nstream\n' + content + b'endstream'
        obj_num += 1
        
        # Page object
        page_num = obj_num
        page_refs.append(page_num)
        objects[obj_num] = (
            f'<< /Type /Page /Parent {pages_num} 0 R\n'
            f'   /MediaBox [0 0 612 792]\n'
            f'   /Contents {content_num} 0 R\n'
            f'   /Resources << /Font << /F1 {font_num} 0 R >> >>\n'
            f'>>'.encode()
        )
        obj_num += 1
    
    kids_str = ' '.join(f'{r} 0 R' for r in page_refs)
    objects[pages_num] = f'<< /Type /Pages /Kids [{kids_str}] /Count {num_pages} >>'.encode()
    objects[catalog_num] = f'<< /Type /Catalog /Pages {pages_num} 0 R >>'.encode()
    
    pdf = bytearray()
    pdf.extend(b'%PDF-1.4\n')
    
    offsets = {}
    for on in sorted(objects.keys()):
        offsets[on] = len(pdf)
        pdf.extend(f'{on} 0 obj\n'.encode())
        pdf.extend(objects[on])
        pdf.extend(b'\nendobj\n\n')
    
    xref_start = len(pdf)
    n = max(objects.keys()) + 1
    pdf.extend(f'xref\n0 {n}\n'.encode())
    pdf.extend(b'0000000000 65535 f \n')
    for i in range(1, n):
        pdf.extend(f'{offsets.get(i, 0):010d} 00000 n \n'.encode())
    
    pdf.extend(f'trailer\n<< /Size {n} /Root {catalog_num} 0 R >>\nstartxref\n{xref_start}\n%%EOF\n'.encode())
    return bytes(pdf)

data = make_large_pdf(100)
with open('corpus/fonts/large_100pages.pdf', 'wb') as f:
    f.write(data)
print(f'Created large_100pages.pdf: {len(data)} bytes')

print('All corpus PDFs created successfully!')