from __future__ import annotations import http.cookies import json import os import stat import typing import warnings from datetime import datetime from email.utils import format_datetime, formatdate from functools import partial from mimetypes import guess_type from urllib.parse import quote import anyio import anyio.to_thread from starlette._compat import md5_hexdigest from starlette.background import BackgroundTask from starlette.concurrency import iterate_in_threadpool from starlette.datastructures import URL, MutableHeaders from starlette.types import Receive, Scope, Send class Response: media_type = None charset = "utf-8" def __init__( self, content: typing.Any = None, status_code: int = 200, headers: typing.Mapping[str, str] | None = None, media_type: str | None = None, background: BackgroundTask | None = None, ) -> None: self.status_code = status_code if media_type is not None: self.media_type = media_type self.background = background self.body = self.render(content) self.init_headers(headers) def render(self, content: typing.Any) -> bytes: if content is None: return b"" if isinstance(content, bytes): return content return content.encode(self.charset) # type: ignore def init_headers(self, headers: typing.Mapping[str, str] | None = None) -> None: if headers is None: raw_headers: list[tuple[bytes, bytes]] = [] populate_content_length = True populate_content_type = True else: raw_headers = [ (k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers.items() ] keys = [h[0] for h in raw_headers] populate_content_length = b"content-length" not in keys populate_content_type = b"content-type" not in keys body = getattr(self, "body", None) if ( body is not None and populate_content_length and not (self.status_code < 200 or self.status_code in (204, 304)) ): content_length = str(len(body)) raw_headers.append((b"content-length", content_length.encode("latin-1"))) content_type = self.media_type if content_type is not None and populate_content_type: if ( content_type.startswith("text/") and "charset=" not in content_type.lower() ): content_type += "; charset=" + self.charset raw_headers.append((b"content-type", content_type.encode("latin-1"))) self.raw_headers = raw_headers @property def headers(self) -> MutableHeaders: if not hasattr(self, "_headers"): self._headers = MutableHeaders(raw=self.raw_headers) return self._headers def set_cookie( self, key: str, value: str = "", max_age: int | None = None, expires: datetime | str | int | None = None, path: str = "/", domain: str | None = None, secure: bool = False, httponly: bool = False, samesite: typing.Literal["lax", "strict", "none"] | None = "lax", ) -> None: cookie: http.cookies.BaseCookie[str] = http.cookies.SimpleCookie() cookie[key] = value if max_age is not None: cookie[key]["max-age"] = max_age if expires is not None: if isinstance(expires, datetime): cookie[key]["expires"] = format_datetime(expires, usegmt=True) else: cookie[key]["expires"] = expires if path is not None: cookie[key]["path"] = path if domain is not None: cookie[key]["domain"] = domain if secure: cookie[key]["secure"] = True if httponly: cookie[key]["httponly"] = True if samesite is not None: assert samesite.lower() in [ "strict", "lax", "none", ], "samesite must be either 'strict', 'lax' or 'none'" cookie[key]["samesite"] = samesite cookie_val = cookie.output(header="").strip() self.raw_headers.append((b"set-cookie", cookie_val.encode("latin-1"))) def delete_cookie( self, key: str, path: str = "/", domain: str | None = None, secure: bool = False, httponly: bool = False, samesite: typing.Literal["lax", "strict", "none"] | None = "lax", ) -> None: self.set_cookie( key, max_age=0, expires=0, path=path, domain=domain, secure=secure, httponly=httponly, samesite=samesite, ) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: prefix = "websocket." if scope["type"] == "websocket" else "" await send( { "type": prefix + "http.response.start", "status": self.status_code, "headers": self.raw_headers, } ) await send({"type": prefix + "http.response.body", "body": self.body}) if self.background is not None: await self.background() class HTMLResponse(Response): media_type = "text/html" class PlainTextResponse(Response): media_type = "text/plain" class JSONResponse(Response): media_type = "application/json" def __init__( self, content: typing.Any, status_code: int = 200, headers: typing.Mapping[str, str] | None = None, media_type: str | None = None, background: BackgroundTask | None = None, ) -> None: super().__init__(content, status_code, headers, media_type, background) def render(self, content: typing.Any) -> bytes: return json.dumps( content, ensure_ascii=False, allow_nan=False, indent=None, separators=(",", ":"), ).encode("utf-8") class RedirectResponse(Response): def __init__( self, url: str | URL, status_code: int = 307, headers: typing.Mapping[str, str] | None = None, background: BackgroundTask | None = None, ) -> None: super().__init__( content=b"", status_code=status_code, headers=headers, background=background ) self.headers["location"] = quote(str(url), safe=":/%#?=@[]!$&'()*+,;") Content = typing.Union[str, bytes] SyncContentStream = typing.Iterable[Content] AsyncContentStream = typing.AsyncIterable[Content] ContentStream = typing.Union[AsyncContentStream, SyncContentStream] class StreamingResponse(Response): body_iterator: AsyncContentStream def __init__( self, content: ContentStream, status_code: int = 200, headers: typing.Mapping[str, str] | None = None, media_type: str | None = None, background: BackgroundTask | None = None, ) -> None: if isinstance(content, typing.AsyncIterable): self.body_iterator = content else: self.body_iterator = iterate_in_threadpool(content) self.status_code = status_code self.media_type = self.media_type if media_type is None else media_type self.background = background self.init_headers(headers) async def listen_for_disconnect(self, receive: Receive) -> None: while True: message = await receive() if message["type"] == "http.disconnect": break async def stream_response(self, send: Send) -> None: await send( { "type": "http.response.start", "status": self.status_code, "headers": self.raw_headers, } ) async for chunk in self.body_iterator: if not isinstance(chunk, bytes): chunk = chunk.encode(self.charset) await send({"type": "http.response.body", "body": chunk, "more_body": True}) await send({"type": "http.response.body", "body": b"", "more_body": False}) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: async with anyio.create_task_group() as task_group: async def wrap(func: typing.Callable[[], typing.Awaitable[None]]) -> None: await func() task_group.cancel_scope.cancel() task_group.start_soon(wrap, partial(self.stream_response, send)) await wrap(partial(self.listen_for_disconnect, receive)) if self.background is not None: await self.background() class FileResponse(Response): chunk_size = 64 * 1024 def __init__( self, path: str | os.PathLike[str], status_code: int = 200, headers: typing.Mapping[str, str] | None = None, media_type: str | None = None, background: BackgroundTask | None = None, filename: str | None = None, stat_result: os.stat_result | None = None, method: str | None = None, content_disposition_type: str = "attachment", ) -> None: self.path = path self.status_code = status_code self.filename = filename if method is not None: warnings.warn( "The 'method' parameter is not used, and it will be removed.", DeprecationWarning, ) if media_type is None: media_type = guess_type(filename or path)[0] or "text/plain" self.media_type = media_type self.background = background self.init_headers(headers) if self.filename is not None: content_disposition_filename = quote(self.filename) if content_disposition_filename != self.filename: content_disposition = "{}; filename*=utf-8''{}".format( content_disposition_type, content_disposition_filename ) else: content_disposition = '{}; filename="{}"'.format( content_disposition_type, self.filename ) self.headers.setdefault("content-disposition", content_disposition) self.stat_result = stat_result if stat_result is not None: self.set_stat_headers(stat_result) def set_stat_headers(self, stat_result: os.stat_result) -> None: content_length = str(stat_result.st_size) last_modified = formatdate(stat_result.st_mtime, usegmt=True) etag_base = str(stat_result.st_mtime) + "-" + str(stat_result.st_size) etag = f'"{md5_hexdigest(etag_base.encode(), usedforsecurity=False)}"' self.headers.setdefault("content-length", content_length) self.headers.setdefault("last-modified", last_modified) self.headers.setdefault("etag", etag) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if self.stat_result is None: try: stat_result = await anyio.to_thread.run_sync(os.stat, self.path) self.set_stat_headers(stat_result) except FileNotFoundError: raise RuntimeError(f"File at path {self.path} does not exist.") else: mode = stat_result.st_mode if not stat.S_ISREG(mode): raise RuntimeError(f"File at path {self.path} is not a file.") await send( { "type": "http.response.start", "status": self.status_code, "headers": self.raw_headers, } ) if scope["method"].upper() == "HEAD": await send({"type": "http.response.body", "body": b"", "more_body": False}) elif "extensions" in scope and "http.response.pathsend" in scope["extensions"]: await send({"type": "http.response.pathsend", "path": str(self.path)}) else: async with await anyio.open_file(self.path, mode="rb") as file: more_body = True while more_body: chunk = await file.read(self.chunk_size) more_body = len(chunk) == self.chunk_size await send( { "type": "http.response.body", "body": chunk, "more_body": more_body, } ) if self.background is not None: await self.background()