cn-py-serv/server/response.py

import logging

log = logging.getLogger(__name__)

STATUS_BY_CODE = {
    200: "OK",

    400: "Bad Request",
    404: "Not Found",
    405: "Method Not Allowed",
    413: "Payload Too Large",

    500: "Internal Server Error",
    505: "HTTP Version Not Supported",
}


class Response:
    def __iter__(self):
        return self

    def __init__(self, status_code=200, additional_headers={}, data=None):
        self.status_code = status_code
        self.additional_headers = additional_headers
        self.data = data
        self.first = True

    def __next__(self):
        if self.first:
            log.debug("First packet generation of response")

            resp = "HTTP/1.1 {} {}\r\n"
            resp += "Server: cnserv\r\n"
            resp += "Connection: close\r\n"
            for key, val in self.additional_headers.items():
                resp += "{}: {}\r\n".format(key, val)
            resp += "\r\n"

            resp = resp.format(self.status_code, STATUS_BY_CODE[self.status_code])
            resp = resp.encode("UTF-8")

            if type(self.data) == bytes:
                resp += self.data

            self.first = False

            return resp

        else:
            if type(self.data) == bytes:
                raise StopIteration

            elif hasattr(self.data, "__iter__"):
                log.debug("Data is iterator, just generate next part")
                try:
                    return next(self.data)
                except StopIteration:
                    raise StopIteration