onnx-web/api/onnx_web/server/utils.py

from functools import partial, update_wrapper
from os import makedirs, path
from typing import Callable, Dict, List, Tuple

from flask import Flask

from ..utils import base_join
from ..worker.pool import DevicePoolExecutor
from .context import ServerContext


def check_paths(server: ServerContext) -> None:
    if not path.exists(server.model_path):
        raise RuntimeError("model path must exist")

    if not path.exists(server.output_path):
        makedirs(server.output_path)


def get_model_path(server: ServerContext, model: str):
    return base_join(server.model_path, model)


def register_routes(
    app: Flask,
    server: ServerContext,
    pool: DevicePoolExecutor,
    routes: List[Tuple[str, Dict, Callable]],
):
    for route, kwargs, method in routes:
        app.route(route, **kwargs)(wrap_route(method, server, pool=pool))


def wrap_route(func, *args, **kwargs):
    """
    From http://louistiao.me/posts/adding-__name__-and-__doc__-attributes-to-functoolspartial-objects/
    """
    partial_func = partial(func, *args, **kwargs)
    update_wrapper(partial_func, func)
    return partial_func