Python PDF to Markdown

import base64
import hashlib
import json
import os
import sys
import tempfile
import time
import uuid
from pathlib import Path
import requests

# Usage: python batch.py selected-pdfs batch-state [--resume]
# Keep batch-state: it contains private PDF snapshots and saved identities.
ENDPOINT = "https://pdftomarkdown.dev/v1/convert"
MAX_FILES = 5
OPTIONS = {"max_pages": 3, "include_raw": False}
RECOVERY = ("Stopped. Keep batch-state and its immutable inputs/options/identities. "
            "Read /docs/#errors before explicitly rerunning with --resume. "
            "Never delete state to retry; after replay expiry, reconcile manually.")


def durable_write(path, data, replace=False):
    # A complete file becomes visible atomically; new outputs never overwrite.
    fd, temporary = tempfile.mkstemp(dir=path.parent)
    try:
        with os.fdopen(fd, "wb") as output:
            output.write(data)
            output.flush()
            os.fsync(output.fileno())
        if replace:
            os.replace(temporary, path)
        else:
            os.link(temporary, path)
        directory = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def main():
    if len(sys.argv) not in (3, 4) or (len(sys.argv) == 4 and sys.argv[3] != "--resume"):
        sys.exit("Usage: python batch.py selected-pdfs batch-state [--resume]")
    key = os.environ.get("PDFTOMARKDOWN_API_KEY")
    if not key or key == "demo_public_key":
        sys.exit("Set PDFTOMARKDOWN_API_KEY to your account key")
    selected, state = (Path(value).resolve() for value in sys.argv[1:3])
    resume = len(sys.argv) == 4
    state.mkdir(mode=0o700, parents=True, exist_ok=True)
    lock = state / "running.lock"
    try:
        lock.mkdir()
    except FileExistsError:
        sys.exit("Batch locked. Confirm no process is running before removing running.lock. " + RECOVERY)
    try:
        manifest = state / "manifest.json"
        if manifest.exists():
            batch = json.loads(manifest.read_text(encoding="utf-8"))
            if batch["selected"] != str(selected):
                sys.exit("This state belongs to another selected directory. " + RECOVERY)
        else:
            if resume:
                sys.exit("Recovery state is missing. " + RECOVERY)
            files = sorted(selected.glob("*.pdf"))
            if not selected.is_dir() or not 1 <= len(files) <= MAX_FILES:
                sys.exit("Select a directory containing 1 to 5 finished .pdf files")
            jobs = []
            for pdf in files:
                if not pdf.is_file() or not 0 < pdf.stat().st_size <= 50 * 1024 * 1024:
                    sys.exit("Each selected PDF must be between 1 byte and 50 MiB")
                data = pdf.read_bytes()
                jobs.append({"output": pdf.name + ".md", "identity": str(uuid.uuid4()),
                             "sha256": hashlib.sha256(data).hexdigest(),
                             "input": {"pdf_base64": base64.b64encode(data).decode("ascii"), **OPTIONS},
                             "attempted_at": None, "result": None})
            batch = {"selected": str(selected), "jobs": jobs}
            durable_write(manifest, json.dumps(batch).encode("utf-8"))

        def save():
            durable_write(manifest, json.dumps(batch).encode("utf-8"), replace=True)

        # Check every destination and snapshot before submitting any work.
        for job in batch["jobs"]:
            data = base64.b64decode(job["input"]["pdf_base64"], validate=True)
            if hashlib.sha256(data).hexdigest() != job["sha256"]:
                sys.exit("Saved input is damaged. " + RECOVERY)
            output = state / job["output"]
            if output.exists() and (job["result"] is None or
                    output.read_bytes() != job["result"]["markdown"].encode("utf-8")):
                sys.exit("Output already exists; no file was overwritten. " + RECOVERY)

        with requests.Session() as session:
            for job in batch["jobs"]:
                output = state / job["output"]
                if job["result"] is None:
                    if job["attempted_at"] is not None:
                        age = time.time() - job["attempted_at"]
                        if not resume or age < 0 or age >= 24 * 60 * 60:
                            sys.exit("Explicit recovery within the 24-hour window is required. " + RECOVERY)
                    else:
                        job["attempted_at"] = time.time()
                        save()  # Durable identity, exact input/options and attempt time before POST.
                    response = session.post(
                        ENDPOINT,
                        headers={"Authorization": f"Bearer {key}",
                                 "Idempotency-Key": job["identity"]},
                        json={"input": job["input"]}, timeout=(10, 660),
                        allow_redirects=False,
                    )
                    if response.status_code != 200:
                        sys.exit(f"HTTP {response.status_code}. " + RECOVERY)
                    result = response.json()
                    if not (isinstance(result, dict) and result.get("complete") is True
                            and isinstance(result.get("markdown"), str)
                            and type(result.get("pages")) is int and result["pages"] >= 0
                            and isinstance(result.get("request_id"), str) and result["request_id"].strip()):
                        sys.exit("Incomplete or invalid conversion response. " + RECOVERY)
                    job["result"] = result
                    save()  # Recovery can finish writing without another conversion.
                if not output.exists():
                    durable_write(output, job["result"]["markdown"].encode("utf-8"))
        print("Batch complete. Markdown saved in batch-state; retain its manifest.")
    finally:
        lock.rmdir()


try:
    main()
except (requests.RequestException, ValueError, OSError, KeyError, TypeError):
    sys.exit(RECOVERY)