| import asyncio |
| import base64 |
| import json |
| import sys |
| from pathlib import Path |
|
|
| from .schemas import RasterPage, SafeError |
|
|
|
|
| async def parse(data, filename, options, settings): |
| |
| limits = { |
| name: getattr(settings, name) |
| for name in ("max_upload", "max_pages", "max_pixels", "total_pixels") |
| } |
| payload = json.dumps( |
| dict( |
| data=base64.b64encode(data).decode(), |
| filename=filename, |
| options=options.model_dump(), |
| limits=limits, |
| ) |
| ).encode() |
| proc = await asyncio.create_subprocess_exec( |
| sys.executable, |
| "-m", |
| "ocr.parse_worker", |
| stdin=asyncio.subprocess.PIPE, |
| stdout=asyncio.subprocess.PIPE, |
| stderr=asyncio.subprocess.DEVNULL, |
| cwd=str(Path(__file__).resolve().parents[1]), |
| env={ |
| "PATH": "/usr/bin:/bin", |
| "LANG": "C.UTF-8", |
| "OPENBLAS_NUM_THREADS": "1", |
| "OMP_NUM_THREADS": "1", |
| }, |
| ) |
|
|
| async def transfer(): |
| async def write(): |
| try: |
| proc.stdin.write(payload) |
| await proc.stdin.drain() |
| proc.stdin.close() |
| except (BrokenPipeError, ConnectionResetError): |
| pass |
|
|
| task = asyncio.create_task(write()) |
| try: |
| out = bytearray() |
| while chunk := await proc.stdout.read(64 * 1024): |
| out.extend(chunk) |
| if len(out) > 128 * 1024 * 1024: |
| raise SafeError("parser_output_limit") |
| await task |
| await proc.wait() |
| if proc.returncode: |
| raise SafeError("parser_resource_limit") |
| try: |
| result = json.loads(out) |
| if "error" in result: |
| raise SafeError(result["error"]) |
| pages = [ |
| RasterPage(**{**p, "image": base64.b64decode(p["image"], validate=True)}) |
| for p in result["pages"] |
| ] |
| if not 1 <= len(pages) <= min(options.max_pages, settings.max_pages): |
| raise SafeError("page_limit") |
| return pages |
| except SafeError: |
| raise |
| except Exception: |
| raise SafeError("invalid_document") from None |
| finally: |
| task.cancel() |
| await asyncio.gather(task, return_exceptions=True) |
|
|
| try: |
| return await asyncio.wait_for(transfer(), 40) |
| except asyncio.TimeoutError: |
| raise SafeError("parser_timeout") from None |
| finally: |
| if proc.returncode is None: |
| proc.kill() |
| await proc.wait() |
|
|
