responses.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354
  1. import http.cookies
  2. import json
  3. import os
  4. import stat
  5. import typing
  6. import warnings
  7. from datetime import datetime
  8. from email.utils import format_datetime, formatdate
  9. from functools import partial
  10. from mimetypes import guess_type
  11. from urllib.parse import quote
  12. import anyio
  13. import anyio.to_thread
  14. from starlette._compat import md5_hexdigest
  15. from starlette.background import BackgroundTask
  16. from starlette.concurrency import iterate_in_threadpool
  17. from starlette.datastructures import URL, MutableHeaders
  18. from starlette.types import Receive, Scope, Send
  19. class Response:
  20. media_type = None
  21. charset = "utf-8"
  22. def __init__(
  23. self,
  24. content: typing.Any = None,
  25. status_code: int = 200,
  26. headers: typing.Optional[typing.Mapping[str, str]] = None,
  27. media_type: typing.Optional[str] = None,
  28. background: typing.Optional[BackgroundTask] = None,
  29. ) -> None:
  30. self.status_code = status_code
  31. if media_type is not None:
  32. self.media_type = media_type
  33. self.background = background
  34. self.body = self.render(content)
  35. self.init_headers(headers)
  36. def render(self, content: typing.Any) -> bytes:
  37. if content is None:
  38. return b""
  39. if isinstance(content, bytes):
  40. return content
  41. return content.encode(self.charset) # type: ignore
  42. def init_headers(
  43. self, headers: typing.Optional[typing.Mapping[str, str]] = None
  44. ) -> None:
  45. if headers is None:
  46. raw_headers: typing.List[typing.Tuple[bytes, bytes]] = []
  47. populate_content_length = True
  48. populate_content_type = True
  49. else:
  50. raw_headers = [
  51. (k.lower().encode("latin-1"), v.encode("latin-1"))
  52. for k, v in headers.items()
  53. ]
  54. keys = [h[0] for h in raw_headers]
  55. populate_content_length = b"content-length" not in keys
  56. populate_content_type = b"content-type" not in keys
  57. body = getattr(self, "body", None)
  58. if (
  59. body is not None
  60. and populate_content_length
  61. and not (self.status_code < 200 or self.status_code in (204, 304))
  62. ):
  63. content_length = str(len(body))
  64. raw_headers.append((b"content-length", content_length.encode("latin-1")))
  65. content_type = self.media_type
  66. if content_type is not None and populate_content_type:
  67. if content_type.startswith("text/"):
  68. content_type += "; charset=" + self.charset
  69. raw_headers.append((b"content-type", content_type.encode("latin-1")))
  70. self.raw_headers = raw_headers
  71. @property
  72. def headers(self) -> MutableHeaders:
  73. if not hasattr(self, "_headers"):
  74. self._headers = MutableHeaders(raw=self.raw_headers)
  75. return self._headers
  76. def set_cookie(
  77. self,
  78. key: str,
  79. value: str = "",
  80. max_age: typing.Optional[int] = None,
  81. expires: typing.Optional[typing.Union[datetime, str, int]] = None,
  82. path: str = "/",
  83. domain: typing.Optional[str] = None,
  84. secure: bool = False,
  85. httponly: bool = False,
  86. samesite: typing.Optional[typing.Literal["lax", "strict", "none"]] = "lax",
  87. ) -> None:
  88. cookie: "http.cookies.BaseCookie[str]" = http.cookies.SimpleCookie()
  89. cookie[key] = value
  90. if max_age is not None:
  91. cookie[key]["max-age"] = max_age
  92. if expires is not None:
  93. if isinstance(expires, datetime):
  94. cookie[key]["expires"] = format_datetime(expires, usegmt=True)
  95. else:
  96. cookie[key]["expires"] = expires
  97. if path is not None:
  98. cookie[key]["path"] = path
  99. if domain is not None:
  100. cookie[key]["domain"] = domain
  101. if secure:
  102. cookie[key]["secure"] = True
  103. if httponly:
  104. cookie[key]["httponly"] = True
  105. if samesite is not None:
  106. assert samesite.lower() in [
  107. "strict",
  108. "lax",
  109. "none",
  110. ], "samesite must be either 'strict', 'lax' or 'none'"
  111. cookie[key]["samesite"] = samesite
  112. cookie_val = cookie.output(header="").strip()
  113. self.raw_headers.append((b"set-cookie", cookie_val.encode("latin-1")))
  114. def delete_cookie(
  115. self,
  116. key: str,
  117. path: str = "/",
  118. domain: typing.Optional[str] = None,
  119. secure: bool = False,
  120. httponly: bool = False,
  121. samesite: typing.Optional[typing.Literal["lax", "strict", "none"]] = "lax",
  122. ) -> None:
  123. self.set_cookie(
  124. key,
  125. max_age=0,
  126. expires=0,
  127. path=path,
  128. domain=domain,
  129. secure=secure,
  130. httponly=httponly,
  131. samesite=samesite,
  132. )
  133. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  134. await send(
  135. {
  136. "type": "http.response.start",
  137. "status": self.status_code,
  138. "headers": self.raw_headers,
  139. }
  140. )
  141. await send({"type": "http.response.body", "body": self.body})
  142. if self.background is not None:
  143. await self.background()
  144. class HTMLResponse(Response):
  145. media_type = "text/html"
  146. class PlainTextResponse(Response):
  147. media_type = "text/plain"
  148. class JSONResponse(Response):
  149. media_type = "application/json"
  150. def __init__(
  151. self,
  152. content: typing.Any,
  153. status_code: int = 200,
  154. headers: typing.Optional[typing.Mapping[str, str]] = None,
  155. media_type: typing.Optional[str] = None,
  156. background: typing.Optional[BackgroundTask] = None,
  157. ) -> None:
  158. super().__init__(content, status_code, headers, media_type, background)
  159. def render(self, content: typing.Any) -> bytes:
  160. return json.dumps(
  161. content,
  162. ensure_ascii=False,
  163. allow_nan=False,
  164. indent=None,
  165. separators=(",", ":"),
  166. ).encode("utf-8")
  167. class RedirectResponse(Response):
  168. def __init__(
  169. self,
  170. url: typing.Union[str, URL],
  171. status_code: int = 307,
  172. headers: typing.Optional[typing.Mapping[str, str]] = None,
  173. background: typing.Optional[BackgroundTask] = None,
  174. ) -> None:
  175. super().__init__(
  176. content=b"", status_code=status_code, headers=headers, background=background
  177. )
  178. self.headers["location"] = quote(str(url), safe=":/%#?=@[]!$&'()*+,;")
  179. Content = typing.Union[str, bytes]
  180. SyncContentStream = typing.Iterable[Content]
  181. AsyncContentStream = typing.AsyncIterable[Content]
  182. ContentStream = typing.Union[AsyncContentStream, SyncContentStream]
  183. class StreamingResponse(Response):
  184. body_iterator: AsyncContentStream
  185. def __init__(
  186. self,
  187. content: ContentStream,
  188. status_code: int = 200,
  189. headers: typing.Optional[typing.Mapping[str, str]] = None,
  190. media_type: typing.Optional[str] = None,
  191. background: typing.Optional[BackgroundTask] = None,
  192. ) -> None:
  193. if isinstance(content, typing.AsyncIterable):
  194. self.body_iterator = content
  195. else:
  196. self.body_iterator = iterate_in_threadpool(content)
  197. self.status_code = status_code
  198. self.media_type = self.media_type if media_type is None else media_type
  199. self.background = background
  200. self.init_headers(headers)
  201. async def listen_for_disconnect(self, receive: Receive) -> None:
  202. while True:
  203. message = await receive()
  204. if message["type"] == "http.disconnect":
  205. break
  206. async def stream_response(self, send: Send) -> None:
  207. await send(
  208. {
  209. "type": "http.response.start",
  210. "status": self.status_code,
  211. "headers": self.raw_headers,
  212. }
  213. )
  214. async for chunk in self.body_iterator:
  215. if not isinstance(chunk, bytes):
  216. chunk = chunk.encode(self.charset)
  217. await send({"type": "http.response.body", "body": chunk, "more_body": True})
  218. await send({"type": "http.response.body", "body": b"", "more_body": False})
  219. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  220. async with anyio.create_task_group() as task_group:
  221. async def wrap(func: "typing.Callable[[], typing.Awaitable[None]]") -> None:
  222. await func()
  223. task_group.cancel_scope.cancel()
  224. task_group.start_soon(wrap, partial(self.stream_response, send))
  225. await wrap(partial(self.listen_for_disconnect, receive))
  226. if self.background is not None:
  227. await self.background()
  228. class FileResponse(Response):
  229. chunk_size = 64 * 1024
  230. def __init__(
  231. self,
  232. path: typing.Union[str, "os.PathLike[str]"],
  233. status_code: int = 200,
  234. headers: typing.Optional[typing.Mapping[str, str]] = None,
  235. media_type: typing.Optional[str] = None,
  236. background: typing.Optional[BackgroundTask] = None,
  237. filename: typing.Optional[str] = None,
  238. stat_result: typing.Optional[os.stat_result] = None,
  239. method: typing.Optional[str] = None,
  240. content_disposition_type: str = "attachment",
  241. ) -> None:
  242. self.path = path
  243. self.status_code = status_code
  244. self.filename = filename
  245. if method is not None:
  246. warnings.warn(
  247. "The 'method' parameter is not used, and it will be removed.",
  248. DeprecationWarning,
  249. )
  250. if media_type is None:
  251. media_type = guess_type(filename or path)[0] or "text/plain"
  252. self.media_type = media_type
  253. self.background = background
  254. self.init_headers(headers)
  255. if self.filename is not None:
  256. content_disposition_filename = quote(self.filename)
  257. if content_disposition_filename != self.filename:
  258. content_disposition = "{}; filename*=utf-8''{}".format(
  259. content_disposition_type, content_disposition_filename
  260. )
  261. else:
  262. content_disposition = '{}; filename="{}"'.format(
  263. content_disposition_type, self.filename
  264. )
  265. self.headers.setdefault("content-disposition", content_disposition)
  266. self.stat_result = stat_result
  267. if stat_result is not None:
  268. self.set_stat_headers(stat_result)
  269. def set_stat_headers(self, stat_result: os.stat_result) -> None:
  270. content_length = str(stat_result.st_size)
  271. last_modified = formatdate(stat_result.st_mtime, usegmt=True)
  272. etag_base = str(stat_result.st_mtime) + "-" + str(stat_result.st_size)
  273. etag = f'"{md5_hexdigest(etag_base.encode(), usedforsecurity=False)}"'
  274. self.headers.setdefault("content-length", content_length)
  275. self.headers.setdefault("last-modified", last_modified)
  276. self.headers.setdefault("etag", etag)
  277. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  278. if self.stat_result is None:
  279. try:
  280. stat_result = await anyio.to_thread.run_sync(os.stat, self.path)
  281. self.set_stat_headers(stat_result)
  282. except FileNotFoundError:
  283. raise RuntimeError(f"File at path {self.path} does not exist.")
  284. else:
  285. mode = stat_result.st_mode
  286. if not stat.S_ISREG(mode):
  287. raise RuntimeError(f"File at path {self.path} is not a file.")
  288. await send(
  289. {
  290. "type": "http.response.start",
  291. "status": self.status_code,
  292. "headers": self.raw_headers,
  293. }
  294. )
  295. if scope["method"].upper() == "HEAD":
  296. await send({"type": "http.response.body", "body": b"", "more_body": False})
  297. else:
  298. async with await anyio.open_file(self.path, mode="rb") as file:
  299. more_body = True
  300. while more_body:
  301. chunk = await file.read(self.chunk_size)
  302. more_body = len(chunk) == self.chunk_size
  303. await send(
  304. {
  305. "type": "http.response.body",
  306. "body": chunk,
  307. "more_body": more_body,
  308. }
  309. )
  310. if self.background is not None:
  311. await self.background()