path: root/server.py

import inspect
import json
import re
from typing import Union, get_type_hints

from pydantic import ConfigDict, TypeAdapter, ValidationError, create_model


CONTENT_TEXT = b"text/plain; charset=utf-8"
JSON_SCHEMA_DRAFT = "https://json-schema.org/draft/2020-12/schema"


def reject_non_json_constant(value):
    raise ValueError(f"Invalid JSON value: {value}")


def schema_document(schema):
    return {"$schema": JSON_SCHEMA_DRAFT, **schema}


def input_schema_document(adapter, model):
    schema = adapter.json_schema()
    definitions = schema.get("$defs", {})
    for index, choice in enumerate(schema.get("anyOf", [])):
        reference = choice.get("$ref", "")
        key = reference.removeprefix("#/$defs/")
        if (
            reference.startswith("#/$defs/")
            and definitions.get(key, {}).get("title") == model.__name__
        ):
            definition = definitions.pop(key)
            definition.pop("title", None)
            schema["anyOf"][index] = definition
            if not definitions:
                schema.pop("$defs")
            break
    if schema.get("title") == model.__name__:
        schema.pop("title")
    return schema_document(schema)


class Route:
    def __init__(self, name):
        self.name = name
        self.source = None
        self.function = None
        self.input_model = None
        self.input_type = None
        self.value_type = None
        self.output_type = None
        self.schema = None

    def patch(self, body):
        try:
            source = body.decode("utf-8")
            namespace = {"__name__": "__main__"}
            exec(compile(source, "<function>", "exec"), namespace)
            function = namespace.get(self.name)
            if not inspect.isfunction(function) or inspect.iscoroutinefunction(
                function
            ):
                raise ValueError(f"Expected a synchronous function named {self.name}")
            parameters = list(inspect.signature(function).parameters.values())
            if not parameters or any(
                parameter.kind
                not in (
                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
                    inspect.Parameter.KEYWORD_ONLY,
                )
                for parameter in parameters
            ):
                raise ValueError("Function must have one or more named parameters")
            hints = get_type_hints(function)
            if "return" not in hints or any(
                parameter.name not in hints for parameter in parameters
            ):
                raise ValueError(
                    "All parameters and the return value must be annotated"
                )
            fields = {
                parameter.name: (
                    hints[parameter.name],
                    ...
                    if parameter.default is inspect.Parameter.empty
                    else parameter.default,
                )
                for parameter in parameters
            }
            input_model = create_model(
                f"{self.name}Input",
                __config__=ConfigDict(extra="forbid", validate_default=True),
                **fields,
            )
            input_type = TypeAdapter(
                Union[input_model, hints[parameters[0].name]]
                if len(parameters) == 1
                else input_model
            )
            value_type = (
                TypeAdapter(hints[parameters[0].name]) if len(parameters) == 1 else None
            )
            output_type = TypeAdapter(hints["return"])
            schema = {
                "input": input_schema_document(input_type, input_model),
                "output": schema_document(output_type.json_schema()),
            }
        except Exception as error:
            return 400, str(error).encode(), CONTENT_TEXT

        self.source = source
        self.function = function
        self.input_model = input_model
        self.input_type = input_type
        self.value_type = value_type
        self.output_type = output_type
        self.schema = schema
        return 200, b"", CONTENT_TEXT

    def get(self):
        if self.source is None:
            return 404, b"Function not found\n", CONTENT_TEXT
        return 200, json.dumps(self.describe()).encode(), b"application/json"

    def describe(self):
        return {"name": self.name, "source": self.source, "schema": self.schema}

    def post(self, body):
        if (
            self.function is None
            or self.input_model is None
            or self.input_type is None
            or self.output_type is None
        ):
            return 404, b"Function not found\n", CONTENT_TEXT
        try:
            data = json.loads(body, parse_constant=reject_non_json_constant)
            if self.value_type is None:
                values = self.input_model.model_validate(data, strict=True)
            elif isinstance(data, dict):
                try:
                    values = self.input_model.model_validate(data, strict=True)
                except ValidationError:
                    values = self.value_type.validate_python(data, strict=True)
            else:
                values = self.value_type.validate_python(data, strict=True)
        except (ValueError, UnicodeError, ValidationError) as error:
            return 400, str(error).encode(), CONTENT_TEXT
        try:
            if isinstance(values, self.input_model):
                arguments = {
                    name: getattr(values, name)
                    for name in self.input_model.model_fields
                }
            else:
                arguments = {next(iter(self.input_model.model_fields)): values}
            result = self.function(**arguments)
            result = self.output_type.validate_python(result, strict=True)
            response = self.output_type.dump_json(result)
        except Exception as error:
            return 500, str(error).encode(), CONTENT_TEXT
        return 200, response, b"application/json"


class Routes:
    def __init__(self):
        self.routes = {}

    def route(self, name):
        if not re.fullmatch(r"[A-Za-z_]\w*", name):
            return None
        return self.routes.setdefault(name, Route(name))

    def patch(self, name, body):
        route = self.route(name)
        if route is None:
            return 400, b"Invalid function name\n", CONTENT_TEXT
        return route.patch(body)

    def get(self, name):
        route = self.route(name)
        if route is None:
            return 400, b"Invalid function name\n", CONTENT_TEXT
        return route.get()

    def post(self, name, body):
        route = self.route(name)
        if route is None:
            return 400, b"Invalid function name\n", CONTENT_TEXT
        return route.post(body)


class Server:
    def __init__(self, routes=None):
        self.routes = routes or Routes()

    async def read_body(self, receive):
        chunks = []
        while True:
            message = await receive()
            if message["type"] == "http.disconnect":
                break
            chunks.append(message.get("body", b""))
            if not message.get("more_body", False):
                break
        return b"".join(chunks)

    def function_route(self, method, name, body):
        if method == "PATCH":
            return self.routes.patch(name, body)
        if method == "GET":
            return self.routes.get(name)
        if method == "POST":
            return self.routes.post(name, body)
        return 404, b"Function not found\n", CONTENT_TEXT

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return

        path = scope["path"]
        method = scope["method"]
        body = await self.read_body(receive) if method in {"POST", "PATCH"} else b""

        if method == "GET" and path == "/":
            functions = [
                route.describe()
                for name, route in sorted(self.routes.routes.items())
                if route.source is not None
            ]
            status, response_body, content_type = (
                200,
                json.dumps({"functions": functions}).encode(),
                b"application/json",
            )
        elif path == "/":
            status, response_body, content_type = 200, b"", CONTENT_TEXT
        elif (
            method in {"GET", "POST", "PATCH"} and path.count("/") == 1 and path != "/"
        ):
            status, response_body, content_type = self.function_route(
                method, path[1:], body
            )
        else:
            status, response_body, content_type = 404, b"Not found\n", CONTENT_TEXT

        await send(
            {
                "type": "http.response.start",
                "status": status,
                "headers": [
                    (b"content-length", str(len(response_body)).encode()),
                    (b"content-type", content_type),
                ],
            }
        )
        await send({"type": "http.response.body", "body": response_body})


app = Server()


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("server:app", host="0.0.0.0", port=8000)