Source code for litestar.response.base

from __future__ import annotations

import itertools
import re
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar, overload

from litestar.datastructures.cookie import Cookie
from litestar.datastructures.headers import ETag, MutableScopeHeaders
from litestar.enums import MediaType, OpenAPIMediaType
from litestar.exceptions import ImproperlyConfiguredException
from litestar.serialization import default_serializer, encode_json, encode_msgpack, get_serializer
from litestar.status_codes import HTTP_200_OK, HTTP_204_NO_CONTENT, HTTP_304_NOT_MODIFIED
from litestar.types.empty import Empty
from litestar.utils.helpers import get_enum_string_value

if TYPE_CHECKING:
    from typing import Optional

    from litestar.background_tasks import BackgroundTask, BackgroundTasks
    from litestar.connection import Request
    from litestar.types import (
        HTTPResponseBodyEvent,
        HTTPResponseStartEvent,
        Receive,
        ResponseCookies,
        ResponseHeaders,
        Scope,
        Send,
        Serializer,
        TypeEncodersMap,
    )

__all__ = ("ASGIResponse", "Response")

T = TypeVar("T")

MEDIA_TYPE_APPLICATION_JSON_PATTERN = re.compile(r"^application/(?:.+\+)?json")


[docs] class ASGIResponse: """A low-level ASGI response class.""" __slots__ = ( "_encoded_cookies", "background", "body", "content_length", "encoding", "headers", "is_head_response", "status_code", ) _should_set_content_length: ClassVar[bool] = True """A flag to indicate whether the content-length header should be set by default or not."""
[docs] def __init__( self, *, background: BackgroundTask | BackgroundTasks | None = None, body: bytes | str = b"", content_length: int | None = None, cookies: Iterable[Cookie] | None = None, encoding: str = "utf-8", headers: dict[str, Any] | Iterable[tuple[str, str]] | None = None, is_head_response: bool = False, media_type: MediaType | str | None = None, status_code: int | None = None, ) -> None: """A low-level ASGI response class. Args: background: A background task or a list of background tasks to be executed after the response is sent. body: encoded content to send in the response body. content_length: The response content length. cookies: The response cookies. encoding: The response encoding. headers: The response headers. is_head_response: A boolean indicating if the response is a HEAD response. media_type: The response media type. status_code: The response status code. """ body = body.encode() if isinstance(body, str) else body status_code = status_code or HTTP_200_OK self.headers = MutableScopeHeaders() if headers is not None: for k, v in headers.items() if isinstance(headers, dict) else headers: self.headers.add(k, v) # pyright: ignore media_type = get_enum_string_value(media_type or MediaType.JSON) status_allows_body = ( status_code not in {HTTP_204_NO_CONTENT, HTTP_304_NOT_MODIFIED} and status_code >= HTTP_200_OK ) if content_length is None: content_length = len(body) if not status_allows_body or is_head_response: if body and body != b"null": raise ImproperlyConfiguredException( "response content is not supported for HEAD responses and responses with a status code " "that does not allow content (304, 204, < 200)" ) body = b"" else: self.headers.setdefault( "content-type", (f"{media_type}; charset={encoding}" if media_type.startswith("text/") else media_type) ) if self._should_set_content_length: self.headers.setdefault("content-length", str(content_length)) self.background = background self.body = body self.content_length = content_length self._encoded_cookies = tuple( cookie.to_encoded_header() for cookie in (cookies or ()) if not cookie.documentation_only ) self.encoding = encoding self.is_head_response = is_head_response self.status_code: int = status_code
[docs] def encode_headers(self) -> list[tuple[bytes, bytes]]: """Return a list of headers for this response. The list contains both headers and encoded cookies, as tuples, where each tuple represents a single header key-value pair encoded as bytes. """ return [*self.headers.headers, *self._encoded_cookies]
[docs] async def after_response(self) -> None: """Execute after the response is sent. Returns: None """ if self.background is not None: await self.background()
[docs] async def start_response(self, send: Send) -> None: """Emit the start event of the response. This event includes the headers and status codes. Args: send: The ASGI send function. Returns: None """ event: HTTPResponseStartEvent = { "type": "http.response.start", "status": self.status_code, "headers": self.encode_headers(), } await send(event)
[docs] async def send_body(self, send: Send, receive: Receive) -> None: """Emit the response body. Args: send: The ASGI send function. receive: The ASGI receive function. Notes: - Response subclasses should customize this method if there is a need to customize sending data. Returns: None """ event: HTTPResponseBodyEvent = {"type": "http.response.body", "body": self.body, "more_body": False} await send(event)
[docs] async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: """ASGI callable of the ``Response``. Args: scope: The ASGI connection scope. receive: The ASGI receive function. send: The ASGI send function. Returns: None """ await self.start_response(send=send) if self.is_head_response: event: HTTPResponseBodyEvent = {"type": "http.response.body", "body": b"", "more_body": False} await send(event) else: await self.send_body(send=send, receive=receive) await self.after_response()
[docs] class Response(Generic[T]): """Base Litestar HTTP response class, used as the basis for all other response classes.""" __slots__ = ( "background", "content", "cookies", "encoding", "headers", "media_type", "response_type_encoders", "status_code", ) content: T type_encoders: Optional[TypeEncodersMap] = None # noqa: UP045
[docs] def __init__( self, content: T, *, background: BackgroundTask | BackgroundTasks | None = None, cookies: ResponseCookies | None = None, encoding: str = "utf-8", headers: ResponseHeaders | None = None, media_type: MediaType | OpenAPIMediaType | str | None = None, status_code: int | None = None, type_encoders: TypeEncodersMap | None = None, ) -> None: """Initialize the response. Args: content: A value for the response body that will be rendered into bytes string. status_code: An HTTP status code. media_type: A value for the response ``Content-Type`` header. background: A :class:`BackgroundTask <.background_tasks.BackgroundTask>` instance or :class:`BackgroundTasks <.background_tasks.BackgroundTasks>` to execute after the response is finished. Defaults to ``None``. headers: A string keyed dictionary of response headers. Header keys are insensitive. cookies: A list of :class:`Cookie <.datastructures.Cookie>` instances to be set under the response ``Set-Cookie`` header. encoding: The encoding to be used for the response headers. type_encoders: A mapping of types to callables that transform them into types supported for serialization. """ self.content = content self.background = background self.cookies: list[Cookie] = ( [Cookie(key=key, value=value) for key, value in cookies.items()] if isinstance(cookies, Mapping) else list(cookies or []) ) self.encoding = encoding self.headers: dict[str, Any] = ( dict(headers) if isinstance(headers, Mapping) else {h.name: h.value for h in headers or {}} ) self.media_type = media_type self.status_code = status_code self.response_type_encoders = {**(self.type_encoders or {}), **(type_encoders or {})}
@overload def set_cookie(self, /, cookie: Cookie) -> None: ... @overload def set_cookie( self, key: str, value: str | None = None, max_age: int | None = None, expires: int | None = None, path: str = "/", domain: str | None = None, secure: bool = False, httponly: bool = False, samesite: Literal["lax", "strict", "none"] = "lax", ) -> None: ...
[docs] def set_header(self, key: str, value: Any) -> None: """Set a header on the response. Args: key: Header key. value: Header value. Returns: None. """ self.headers[key] = value
[docs] def set_etag(self, etag: str | ETag) -> None: """Set an etag header. Args: etag: An etag value. Returns: None """ self.headers["etag"] = etag.to_header() if isinstance(etag, ETag) else etag
[docs] def render(self, content: Any, media_type: str, enc_hook: Serializer = default_serializer) -> bytes: """Handle the rendering of content into a bytes string. Returns: An encoded bytes string """ if isinstance(content, bytes): return content if content is Empty: raise RuntimeError("The `Empty` sentinel cannot be used as response content") try: if media_type.startswith("text/") and not content: return b"" if isinstance(content, str): return content.encode(self.encoding) if media_type == MediaType.MESSAGEPACK: return encode_msgpack(content, enc_hook) if MEDIA_TYPE_APPLICATION_JSON_PATTERN.match( media_type, ): return encode_json(content, enc_hook) raise ImproperlyConfiguredException(f"unsupported media_type {media_type} for content {content!r}") except (AttributeError, ValueError, TypeError) as e: raise ImproperlyConfiguredException("Unable to serialize response content") from e
[docs] def to_asgi_response( self, request: Request, *, background: BackgroundTask | BackgroundTasks | None = None, cookies: Iterable[Cookie] | None = None, headers: dict[str, str] | None = None, is_head_response: bool = False, media_type: MediaType | str | None = None, status_code: int | None = None, type_encoders: TypeEncodersMap | None = None, ) -> ASGIResponse: """Create an ASGIResponse from a Response instance. Args: background: Background task(s) to be executed after the response is sent. cookies: A list of cookies to be set on the response. headers: Additional headers to be merged with the response headers. Response headers take precedence. is_head_response: Whether the response is a HEAD response. media_type: Media type for the response. If ``media_type`` is already set on the response, this is ignored. request: The :class:`Request <.connection.Request>` instance. status_code: Status code for the response. If ``status_code`` is already set on the response, this is type_encoders: A dictionary of type encoders to use for encoding the response content. Returns: An ASGIResponse instance. """ headers = {**headers, **self.headers} if headers is not None else self.headers cookies = self.cookies if cookies is None else itertools.chain(self.cookies, cookies) if type_encoders: type_encoders = {**type_encoders, **(self.response_type_encoders or {})} else: type_encoders = self.response_type_encoders media_type = get_enum_string_value(self.media_type or media_type or MediaType.JSON) return ASGIResponse( background=self.background or background, body=self.render(self.content, media_type, get_serializer(type_encoders)), cookies=cookies, encoding=self.encoding, headers=headers, is_head_response=is_head_response, media_type=media_type, status_code=self.status_code or status_code, )