No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.
 
 
 
 

331 líneas
11 KiB

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