File size: 2,782 Bytes
fd4e724 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import asyncio
import base64
import json
import sys
from pathlib import Path
from .schemas import RasterPage, SafeError
async def parse(data, filename, options, settings):
# Only non-secret limits cross the parser boundary; credentials are never inherited.
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()
|