cli.py

import argparse
import importlib
import os
import sys
import logging
import warnings
import contextlib
from pathlib import Path

from common import dumps_pretty, run_ocr, run_structure


@contextlib.contextmanager
def suppress_stdout_stderr_fd():
    """A context manager that redirects stdout and stderr to devnull"""
    null_fd = os.open(os.devnull, os.O_RDWR)
    old_stdout, old_stderr = sys.stdout, sys.stderr
    sys.stdout = open(os.devnull, 'w')
    sys.stderr = open(os.devnull, 'w')
    
    try:
        save_fd_out = os.dup(1)
        save_fd_err = os.dup(2)
        os.dup2(null_fd, 1)
        os.dup2(null_fd, 2)
    except OSError:
        save_fd_out = save_fd_err = None

    try:
        yield
    finally:
        sys.stdout.flush()
        sys.stderr.flush()
        old_stdout.flush()
        old_stderr.flush()
        if save_fd_out is not None:
            os.dup2(save_fd_out, 1)
            os.close(save_fd_out)
        if save_fd_err is not None:
            os.dup2(save_fd_err, 2)
            os.close(save_fd_err)
        os.close(null_fd)
        sys.stdout.close()
        sys.stderr.close()
        sys.stdout, sys.stderr = old_stdout, old_stderr

def suppress_logs() -> None:
    os.environ["GLOG_minloglevel"] = "2"
    warnings.filterwarnings("ignore")
    logging.getLogger("ppocr").setLevel(logging.ERROR)
    # Also attempt to silence root logger in case other paddle modules log there
    logging.getLogger().setLevel(logging.ERROR)
    
    try:
        from common import ensure_runtime_env
        ensure_runtime_env()
        px_logging = importlib.import_module("paddlex.utils.logging")
        px_logging.setup_logging("ERROR")
    except ImportError:
        pass

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Offline PaddleOCR CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    for name in ("ocr", "structure"):
        sub = subparsers.add_parser(name)
        sub.add_argument("-i", "--input", required=True, help="Input image or PDF path")
        sub.add_argument(
            "--device",
            default=None,
            help="Device such as cpu or gpu:0. Defaults to OCR_DEVICE env or cpu.",
        )
        sub.add_argument("-o", "--output", default=None, help="Output file path")
        sub.add_argument("--json", action="store_true", help="Output raw JSON instead of plain text")
        sub.add_argument("--debug", action="store_true", help="Enable debug output and paddle logs")

    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if not args.debug:
        suppress_logs()

    # Apply fd-level suppression for the OCR processing unless debugging is enabled
    with suppress_stdout_stderr_fd() if not args.debug else contextlib.nullcontext():
        if args.command == "ocr":
            payload = run_ocr(args.input, device=args.device)
            default_text = payload.get("text", "")
        else:
            payload = run_structure(args.input, device=args.device)
            default_text = dumps_pretty(payload)

    text = dumps_pretty(payload) if args.json else default_text
    
    if args.output:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(text, encoding="utf-8")
        print(f"Saved result to: {output_path}")
    else:
        print(text)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())