SlowAPI/slow/responses.py

import http.client
import json
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .slow import App


class Response:
    status: int
    headers: list[str]
    body: bytes

    def __init__(self, status: int, headers: list[str], body: bytes):
        self.status = status
        self.headers = headers
        self.body = body
        self.no_header = False

    def header(self, app: "App") -> bytes:
        header = f"HTTP/1.1 {self.status} {http.client.responses.get(self.status, 'Unkown Status Code')}\r\n"
        header += f"Content-Length: {len(self.body)}\r\n"
        for h in self.headers:
            header += h + "\r\n"

        header += "\r\n"

        return header.encode(encoding="utf-8")

    def render(self, app: "App") -> bytes:
        if self.no_header:
            return self.header(app)
        return self.header(app) + self.body


def HTTPResponse(
    content: str | bytes,
    status=200,
    content_type="text/plain; charset=utf-8",
    headers=[],
) -> Response:
    if isinstance(content, str):
        content_bytes = content.encode(encoding="utf-8")
    else:
        content_bytes = content
    return Response(
        status, [f"Content-Type: {content_type}", *headers], body=content_bytes
    )


def JSONResponse(d: dict, status=200) -> Response:
    return HTTPResponse(
        json.dumps(d), status=status, content_type="text/json; charset=utf-8"
    )