from __future__ import annotations
import hashlib
import re
from collections.abc import Iterator
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import fitz
from pdf_jsonl_smith.records import make_text_unit_record
from pdf_jsonl_smith.writer import write_jsonl
TextBlock = tuple[float, float, float, float, str]
[docs]
def iter_text_unit_records(
*,
input_path: str | Path,
source_uri: str | None = None,
schema_version: str = "0.1",
) -> Iterator[dict[str, Any]]:
"""Yield one flat record for each merged text unit."""
path = Path(input_path)
doc_id = _sha256_doc_id(path)
extracted_at = _utc_now_iso()
with fitz.open(path) as document:
for page_index, page in enumerate(document):
blocks = _merge_adjacent_text_blocks(_iter_page_text_blocks(page))
for unit_index, block in enumerate(blocks):
x0, y0, x1, y1, text = block
yield make_text_unit_record(
schema_version=schema_version,
doc_id=doc_id,
source_uri=source_uri,
filename=path.name,
page_number=page_index + 1,
unit_type="paragraph",
unit_index=unit_index,
text=text,
bbox=(x0, y0, x1, y1),
extracted_at=extracted_at,
)
def _iter_page_text_blocks(page: fitz.Page) -> Iterator[TextBlock]:
for block in page.get_text("blocks", sort=True):
x0, y0, x1, y1, text, *_ = block
block_type = block[6] if len(block) > 6 else 0
if block_type != 0:
continue
text = _normalize_extracted_text(text)
if not text:
continue
yield (float(x0), float(y0), float(x1), float(y1), text)
def _merge_adjacent_text_blocks(blocks: Iterator[TextBlock]) -> list[TextBlock]:
merged: list[TextBlock] = []
current: TextBlock | None = None
last_block: TextBlock | None = None
for block in blocks:
if current is None:
current = block
last_block = block
continue
if last_block is not None and _should_merge_blocks(last_block, block):
current = _combine_blocks(current, block)
last_block = block
continue
merged.append(current)
current = block
last_block = block
if current is not None:
merged.append(current)
return merged
def _should_merge_blocks(previous: TextBlock, current: TextBlock) -> bool:
prev_x0, prev_y0, _prev_x1, prev_y1, _prev_text = previous
curr_x0, curr_y0, _curr_x1, curr_y1, _curr_text = current
vertical_gap = curr_y0 - prev_y1
previous_height = prev_y1 - prev_y0
current_height = curr_y1 - curr_y0
if vertical_gap < -1:
return False
if vertical_gap > max(12.0, previous_height * 0.8):
return False
if previous_height > current_height * 1.25:
return False
left_delta = curr_x0 - prev_x0
same_left_edge = abs(left_delta) <= 3.0
outdent_continuation = -24.0 <= left_delta < -3.0
lowercase_continuation = _looks_like_lowercase_continuation(
previous_text=_prev_text,
current_text=_curr_text,
)
return same_left_edge or outdent_continuation or lowercase_continuation
def _combine_blocks(previous: TextBlock, current: TextBlock) -> TextBlock:
prev_x0, prev_y0, prev_x1, prev_y1, prev_text = previous
curr_x0, curr_y0, curr_x1, curr_y1, curr_text = current
return (
min(prev_x0, curr_x0),
min(prev_y0, curr_y0),
max(prev_x1, curr_x1),
max(prev_y1, curr_y1),
_join_text(prev_text, curr_text),
)
def _join_text(previous: str, current: str) -> str:
previous = previous.strip()
current = current.strip()
if not previous:
return current
if not current:
return previous
if _is_cjk(previous[-1]) and _is_cjk(current[0]):
return f"{previous}{current}"
return f"{previous} {current}"
def _normalize_extracted_text(text: str) -> str:
text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", text)
parts = [part.strip() for part in text.splitlines() if part.strip()]
if not parts:
return ""
normalized = parts[0]
for part in parts[1:]:
normalized = _join_text(normalized, part)
return re.sub(r"[ \t\f\v]+", " ", normalized).strip()
def _looks_like_lowercase_continuation(
*,
previous_text: str,
current_text: str,
) -> bool:
previous_text = previous_text.strip()
current_text = current_text.strip()
if not previous_text or not current_text:
return False
if previous_text.endswith((".", "!", "?", ":", ";", "。", "！", "？")):
return False
return current_text[0].islower()
def _is_cjk(character: str) -> bool:
codepoint = ord(character)
return (
0x3000 <= codepoint <= 0x303F
or 0x3040 <= codepoint <= 0x30FF
or 0x3400 <= codepoint <= 0x4DBF
or 0x4E00 <= codepoint <= 0x9FFF
or 0xF900 <= codepoint <= 0xFAFF
or 0xFF00 <= codepoint <= 0xFFEF
)
def _sha256_doc_id(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as input_file:
for chunk in iter(lambda: input_file.read(1024 * 1024), b""):
digest.update(chunk)
return f"sha256:{digest.hexdigest()}"
def _utc_now_iso() -> str:
return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace(
"+00:00", "Z"
)