datastructures.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706
  1. from __future__ import annotations
  2. from collections.abc import ItemsView, Iterable, Iterator, KeysView, Mapping, MutableMapping, Sequence, ValuesView
  3. from shlex import shlex
  4. from typing import (
  5. Any,
  6. BinaryIO,
  7. NamedTuple,
  8. TypeVar,
  9. cast,
  10. )
  11. from urllib.parse import SplitResult, parse_qsl, urlencode, urlsplit
  12. from starlette.concurrency import run_in_threadpool
  13. from starlette.types import Scope
  14. class Address(NamedTuple):
  15. host: str
  16. port: int
  17. _KeyType = TypeVar("_KeyType")
  18. # Mapping keys are invariant but their values are covariant since
  19. # you can only read them
  20. # that is, you can't do `Mapping[str, Animal]()["fido"] = Dog()`
  21. _CovariantValueType = TypeVar("_CovariantValueType", covariant=True)
  22. class URL:
  23. def __init__(
  24. self,
  25. url: str = "",
  26. scope: Scope | None = None,
  27. **components: Any,
  28. ) -> None:
  29. if scope is not None:
  30. assert not url, 'Cannot set both "url" and "scope".'
  31. assert not components, 'Cannot set both "scope" and "**components".'
  32. scheme = scope.get("scheme", "http")
  33. server = scope.get("server", None)
  34. path = scope["path"]
  35. query_string = scope.get("query_string", b"")
  36. host_header = None
  37. for key, value in scope["headers"]:
  38. if key == b"host":
  39. host_header = value.decode("latin-1")
  40. break
  41. if host_header is not None:
  42. url = f"{scheme}://{host_header}{path}"
  43. elif server is None:
  44. url = path
  45. else:
  46. host, port = server
  47. default_port = {"http": 80, "https": 443, "ws": 80, "wss": 443}[scheme]
  48. if port == default_port:
  49. url = f"{scheme}://{host}{path}"
  50. else:
  51. url = f"{scheme}://{host}:{port}{path}"
  52. if query_string:
  53. url += "?" + query_string.decode()
  54. elif components:
  55. assert not url, 'Cannot set both "url" and "**components".'
  56. url = URL("").replace(**components).components.geturl()
  57. self._url = url
  58. @property
  59. def components(self) -> SplitResult:
  60. if not hasattr(self, "_components"):
  61. self._components = urlsplit(self._url)
  62. return self._components
  63. @property
  64. def scheme(self) -> str:
  65. return self.components.scheme
  66. @property
  67. def netloc(self) -> str:
  68. return self.components.netloc
  69. @property
  70. def path(self) -> str:
  71. return self.components.path
  72. @property
  73. def query(self) -> str:
  74. return self.components.query
  75. @property
  76. def fragment(self) -> str:
  77. return self.components.fragment
  78. @property
  79. def username(self) -> None | str:
  80. return self.components.username
  81. @property
  82. def password(self) -> None | str:
  83. return self.components.password
  84. @property
  85. def hostname(self) -> None | str:
  86. return self.components.hostname
  87. @property
  88. def port(self) -> int | None:
  89. return self.components.port
  90. @property
  91. def is_secure(self) -> bool:
  92. return self.scheme in ("https", "wss")
  93. def replace(self, **kwargs: Any) -> URL:
  94. if "username" in kwargs or "password" in kwargs or "hostname" in kwargs or "port" in kwargs:
  95. hostname = kwargs.pop("hostname", None)
  96. port = kwargs.pop("port", self.port)
  97. username = kwargs.pop("username", self.username)
  98. password = kwargs.pop("password", self.password)
  99. if hostname is None:
  100. netloc = self.netloc
  101. _, _, hostname = netloc.rpartition("@")
  102. if hostname[-1] != "]":
  103. hostname = hostname.rsplit(":", 1)[0]
  104. netloc = hostname
  105. if port is not None:
  106. netloc += f":{port}"
  107. if username is not None:
  108. userpass = username
  109. if password is not None:
  110. userpass += f":{password}"
  111. netloc = f"{userpass}@{netloc}"
  112. kwargs["netloc"] = netloc
  113. components = self.components._replace(**kwargs)
  114. return self.__class__(components.geturl())
  115. def include_query_params(self, **kwargs: Any) -> URL:
  116. params = MultiDict(parse_qsl(self.query, keep_blank_values=True))
  117. params.update({str(key): str(value) for key, value in kwargs.items()})
  118. query = urlencode(params.multi_items())
  119. return self.replace(query=query)
  120. def replace_query_params(self, **kwargs: Any) -> URL:
  121. query = urlencode([(str(key), str(value)) for key, value in kwargs.items()])
  122. return self.replace(query=query)
  123. def remove_query_params(self, keys: str | Sequence[str]) -> URL:
  124. if isinstance(keys, str):
  125. keys = [keys]
  126. params = MultiDict(parse_qsl(self.query, keep_blank_values=True))
  127. for key in keys:
  128. params.pop(key, None)
  129. query = urlencode(params.multi_items())
  130. return self.replace(query=query)
  131. def __eq__(self, other: Any) -> bool:
  132. return str(self) == str(other)
  133. def __str__(self) -> str:
  134. return self._url
  135. def __repr__(self) -> str:
  136. url = str(self)
  137. if self.password:
  138. url = str(self.replace(password="********"))
  139. return f"{self.__class__.__name__}({repr(url)})"
  140. class URLPath(str):
  141. """
  142. A URL path string that may also hold an associated protocol and/or host.
  143. Used by the routing to return `url_path_for` matches.
  144. """
  145. def __new__(cls, path: str, protocol: str = "", host: str = "") -> URLPath:
  146. assert protocol in ("http", "websocket", "")
  147. return str.__new__(cls, path)
  148. def __init__(self, path: str, protocol: str = "", host: str = "") -> None:
  149. self.protocol = protocol
  150. self.host = host
  151. def make_absolute_url(self, base_url: str | URL) -> URL:
  152. if isinstance(base_url, str):
  153. base_url = URL(base_url)
  154. if self.protocol:
  155. scheme = {
  156. "http": {True: "https", False: "http"},
  157. "websocket": {True: "wss", False: "ws"},
  158. }[self.protocol][base_url.is_secure]
  159. else:
  160. scheme = base_url.scheme
  161. netloc = self.host or base_url.netloc
  162. path = base_url.path.rstrip("/") + str(self)
  163. return URL(scheme=scheme, netloc=netloc, path=path)
  164. class Secret:
  165. """
  166. Holds a string value that should not be revealed in tracebacks etc.
  167. You should cast the value to `str` at the point it is required.
  168. """
  169. def __init__(self, value: str):
  170. self._value = value
  171. def __repr__(self) -> str:
  172. class_name = self.__class__.__name__
  173. return f"{class_name}('**********')"
  174. def __str__(self) -> str:
  175. return self._value
  176. def __bool__(self) -> bool:
  177. return bool(self._value)
  178. class CommaSeparatedStrings(Sequence[str]):
  179. def __init__(self, value: str | Sequence[str]):
  180. if isinstance(value, str):
  181. splitter = shlex(value, posix=True)
  182. splitter.whitespace = ","
  183. splitter.whitespace_split = True
  184. self._items = [item.strip() for item in splitter]
  185. else:
  186. self._items = list(value)
  187. def __len__(self) -> int:
  188. return len(self._items)
  189. def __getitem__(self, index: int | slice) -> Any:
  190. return self._items[index]
  191. def __iter__(self) -> Iterator[str]:
  192. return iter(self._items)
  193. def __repr__(self) -> str:
  194. class_name = self.__class__.__name__
  195. items = [item for item in self]
  196. return f"{class_name}({items!r})"
  197. def __str__(self) -> str:
  198. return ", ".join(repr(item) for item in self)
  199. class ImmutableMultiDict(Mapping[_KeyType, _CovariantValueType]):
  200. _dict: dict[_KeyType, _CovariantValueType]
  201. def __init__(
  202. self,
  203. *args: ImmutableMultiDict[_KeyType, _CovariantValueType]
  204. | Mapping[_KeyType, _CovariantValueType]
  205. | Iterable[tuple[_KeyType, _CovariantValueType]],
  206. **kwargs: Any,
  207. ) -> None:
  208. assert len(args) < 2, "Too many arguments."
  209. value: Any = args[0] if args else []
  210. if kwargs:
  211. value = ImmutableMultiDict(value).multi_items() + ImmutableMultiDict(kwargs).multi_items()
  212. if not value:
  213. _items: list[tuple[Any, Any]] = []
  214. elif hasattr(value, "multi_items"):
  215. value = cast(ImmutableMultiDict[_KeyType, _CovariantValueType], value)
  216. _items = list(value.multi_items())
  217. elif hasattr(value, "items"):
  218. value = cast(Mapping[_KeyType, _CovariantValueType], value)
  219. _items = list(value.items())
  220. else:
  221. value = cast("list[tuple[Any, Any]]", value)
  222. _items = list(value)
  223. self._dict = {k: v for k, v in _items}
  224. self._list = _items
  225. def getlist(self, key: Any) -> list[_CovariantValueType]:
  226. return [item_value for item_key, item_value in self._list if item_key == key]
  227. def keys(self) -> KeysView[_KeyType]:
  228. return self._dict.keys()
  229. def values(self) -> ValuesView[_CovariantValueType]:
  230. return self._dict.values()
  231. def items(self) -> ItemsView[_KeyType, _CovariantValueType]:
  232. return self._dict.items()
  233. def multi_items(self) -> list[tuple[_KeyType, _CovariantValueType]]:
  234. return list(self._list)
  235. def __getitem__(self, key: _KeyType) -> _CovariantValueType:
  236. return self._dict[key]
  237. def __contains__(self, key: Any) -> bool:
  238. return key in self._dict
  239. def __iter__(self) -> Iterator[_KeyType]:
  240. return iter(self.keys())
  241. def __len__(self) -> int:
  242. return len(self._dict)
  243. def __eq__(self, other: Any) -> bool:
  244. if not isinstance(other, self.__class__):
  245. return False
  246. return sorted(self._list) == sorted(other._list)
  247. def __repr__(self) -> str:
  248. class_name = self.__class__.__name__
  249. items = self.multi_items()
  250. return f"{class_name}({items!r})"
  251. class MultiDict(ImmutableMultiDict[Any, Any]):
  252. def __setitem__(self, key: Any, value: Any) -> None:
  253. self.setlist(key, [value])
  254. def __delitem__(self, key: Any) -> None:
  255. self._list = [(k, v) for k, v in self._list if k != key]
  256. del self._dict[key]
  257. def pop(self, key: Any, default: Any = None) -> Any:
  258. self._list = [(k, v) for k, v in self._list if k != key]
  259. return self._dict.pop(key, default)
  260. def popitem(self) -> tuple[Any, Any]:
  261. key, value = self._dict.popitem()
  262. self._list = [(k, v) for k, v in self._list if k != key]
  263. return key, value
  264. def poplist(self, key: Any) -> list[Any]:
  265. values = [v for k, v in self._list if k == key]
  266. self.pop(key)
  267. return values
  268. def clear(self) -> None:
  269. self._dict.clear()
  270. self._list.clear()
  271. def setdefault(self, key: Any, default: Any = None) -> Any:
  272. if key not in self:
  273. self._dict[key] = default
  274. self._list.append((key, default))
  275. return self[key]
  276. def setlist(self, key: Any, values: list[Any]) -> None:
  277. if not values:
  278. self.pop(key, None)
  279. else:
  280. existing_items = [(k, v) for (k, v) in self._list if k != key]
  281. self._list = existing_items + [(key, value) for value in values]
  282. self._dict[key] = values[-1]
  283. def append(self, key: Any, value: Any) -> None:
  284. self._list.append((key, value))
  285. self._dict[key] = value
  286. def update(
  287. self,
  288. *args: MultiDict | Mapping[Any, Any] | list[tuple[Any, Any]],
  289. **kwargs: Any,
  290. ) -> None:
  291. value = MultiDict(*args, **kwargs)
  292. existing_items = [(k, v) for (k, v) in self._list if k not in value.keys()]
  293. self._list = existing_items + value.multi_items()
  294. self._dict.update(value)
  295. class QueryParams(ImmutableMultiDict[str, str]):
  296. """
  297. An immutable multidict.
  298. """
  299. def __init__(
  300. self,
  301. *args: ImmutableMultiDict[Any, Any] | Mapping[Any, Any] | list[tuple[Any, Any]] | str | bytes,
  302. **kwargs: Any,
  303. ) -> None:
  304. assert len(args) < 2, "Too many arguments."
  305. value = args[0] if args else []
  306. if isinstance(value, str):
  307. super().__init__(parse_qsl(value, keep_blank_values=True), **kwargs)
  308. elif isinstance(value, bytes):
  309. super().__init__(parse_qsl(value.decode("latin-1"), keep_blank_values=True), **kwargs)
  310. else:
  311. super().__init__(*args, **kwargs) # type: ignore[arg-type]
  312. self._list = [(str(k), str(v)) for k, v in self._list]
  313. self._dict = {str(k): str(v) for k, v in self._dict.items()}
  314. def __str__(self) -> str:
  315. return urlencode(self._list)
  316. def __repr__(self) -> str:
  317. class_name = self.__class__.__name__
  318. query_string = str(self)
  319. return f"{class_name}({query_string!r})"
  320. class UploadFile:
  321. """
  322. An uploaded file included as part of the request data.
  323. """
  324. def __init__(
  325. self,
  326. file: BinaryIO,
  327. *,
  328. size: int | None = None,
  329. filename: str | None = None,
  330. headers: Headers | None = None,
  331. ) -> None:
  332. self.filename = filename
  333. self.file = file
  334. self.size = size
  335. self.headers = headers or Headers()
  336. # Capture max size from SpooledTemporaryFile if one is provided. This slightly speeds up future checks.
  337. # Note 0 means unlimited mirroring SpooledTemporaryFile's __init__
  338. self._max_mem_size = getattr(self.file, "_max_size", 0)
  339. @property
  340. def content_type(self) -> str | None:
  341. return self.headers.get("content-type", None)
  342. @property
  343. def _in_memory(self) -> bool:
  344. # check for SpooledTemporaryFile._rolled
  345. rolled_to_disk = getattr(self.file, "_rolled", True)
  346. return not rolled_to_disk
  347. def _will_roll(self, size_to_add: int) -> bool:
  348. # If we're not in_memory then we will always roll
  349. if not self._in_memory:
  350. return True
  351. # Check for SpooledTemporaryFile._max_size
  352. future_size = self.file.tell() + size_to_add
  353. return bool(future_size > self._max_mem_size) if self._max_mem_size else False
  354. async def write(self, data: bytes) -> None:
  355. new_data_len = len(data)
  356. if self.size is not None:
  357. self.size += new_data_len
  358. if self._will_roll(new_data_len):
  359. await run_in_threadpool(self.file.write, data)
  360. else:
  361. self.file.write(data)
  362. async def read(self, size: int = -1) -> bytes:
  363. if self._in_memory:
  364. return self.file.read(size)
  365. return await run_in_threadpool(self.file.read, size)
  366. async def seek(self, offset: int) -> None:
  367. if self._in_memory:
  368. self.file.seek(offset)
  369. else:
  370. await run_in_threadpool(self.file.seek, offset)
  371. async def close(self) -> None:
  372. if self._in_memory:
  373. self.file.close()
  374. else:
  375. await run_in_threadpool(self.file.close)
  376. def __repr__(self) -> str:
  377. return f"{self.__class__.__name__}(filename={self.filename!r}, size={self.size!r}, headers={self.headers!r})"
  378. class FormData(ImmutableMultiDict[str, UploadFile | str]):
  379. """
  380. An immutable multidict, containing both file uploads and text input.
  381. """
  382. def __init__(
  383. self,
  384. *args: FormData | Mapping[str, str | UploadFile] | list[tuple[str, str | UploadFile]],
  385. **kwargs: str | UploadFile,
  386. ) -> None:
  387. super().__init__(*args, **kwargs)
  388. async def close(self) -> None:
  389. for key, value in self.multi_items():
  390. if isinstance(value, UploadFile):
  391. await value.close()
  392. class Headers(Mapping[str, str]):
  393. """
  394. An immutable, case-insensitive multidict.
  395. """
  396. def __init__(
  397. self,
  398. headers: Mapping[str, str] | None = None,
  399. raw: list[tuple[bytes, bytes]] | None = None,
  400. scope: MutableMapping[str, Any] | None = None,
  401. ) -> None:
  402. self._list: list[tuple[bytes, bytes]] = []
  403. if headers is not None:
  404. assert raw is None, 'Cannot set both "headers" and "raw".'
  405. assert scope is None, 'Cannot set both "headers" and "scope".'
  406. self._list = [(key.lower().encode("latin-1"), value.encode("latin-1")) for key, value in headers.items()]
  407. elif raw is not None:
  408. assert scope is None, 'Cannot set both "raw" and "scope".'
  409. self._list = raw
  410. elif scope is not None:
  411. # scope["headers"] isn't necessarily a list
  412. # it might be a tuple or other iterable
  413. self._list = scope["headers"] = list(scope["headers"])
  414. @property
  415. def raw(self) -> list[tuple[bytes, bytes]]:
  416. return list(self._list)
  417. def keys(self) -> list[str]: # type: ignore[override]
  418. return [key.decode("latin-1") for key, value in self._list]
  419. def values(self) -> list[str]: # type: ignore[override]
  420. return [value.decode("latin-1") for key, value in self._list]
  421. def items(self) -> list[tuple[str, str]]: # type: ignore[override]
  422. return [(key.decode("latin-1"), value.decode("latin-1")) for key, value in self._list]
  423. def getlist(self, key: str) -> list[str]:
  424. get_header_key = key.lower().encode("latin-1")
  425. return [item_value.decode("latin-1") for item_key, item_value in self._list if item_key == get_header_key]
  426. def mutablecopy(self) -> MutableHeaders:
  427. return MutableHeaders(raw=self._list[:])
  428. def __getitem__(self, key: str) -> str:
  429. get_header_key = key.lower().encode("latin-1")
  430. for header_key, header_value in self._list:
  431. if header_key == get_header_key:
  432. return header_value.decode("latin-1")
  433. raise KeyError(key)
  434. def __contains__(self, key: Any) -> bool:
  435. get_header_key = key.lower().encode("latin-1")
  436. for header_key, header_value in self._list:
  437. if header_key == get_header_key:
  438. return True
  439. return False
  440. def __iter__(self) -> Iterator[Any]:
  441. return iter(self.keys())
  442. def __len__(self) -> int:
  443. return len(self._list)
  444. def __eq__(self, other: Any) -> bool:
  445. if not isinstance(other, Headers):
  446. return False
  447. return sorted(self._list) == sorted(other._list)
  448. def __repr__(self) -> str:
  449. class_name = self.__class__.__name__
  450. as_dict = dict(self.items())
  451. if len(as_dict) == len(self):
  452. return f"{class_name}({as_dict!r})"
  453. return f"{class_name}(raw={self.raw!r})"
  454. class MutableHeaders(Headers):
  455. def __setitem__(self, key: str, value: str) -> None:
  456. """
  457. Set the header `key` to `value`, removing any duplicate entries.
  458. Retains insertion order.
  459. """
  460. set_key = key.lower().encode("latin-1")
  461. set_value = value.encode("latin-1")
  462. found_indexes: list[int] = []
  463. for idx, (item_key, item_value) in enumerate(self._list):
  464. if item_key == set_key:
  465. found_indexes.append(idx)
  466. for idx in reversed(found_indexes[1:]):
  467. del self._list[idx]
  468. if found_indexes:
  469. idx = found_indexes[0]
  470. self._list[idx] = (set_key, set_value)
  471. else:
  472. self._list.append((set_key, set_value))
  473. def __delitem__(self, key: str) -> None:
  474. """
  475. Remove the header `key`.
  476. """
  477. del_key = key.lower().encode("latin-1")
  478. pop_indexes: list[int] = []
  479. for idx, (item_key, item_value) in enumerate(self._list):
  480. if item_key == del_key:
  481. pop_indexes.append(idx)
  482. for idx in reversed(pop_indexes):
  483. del self._list[idx]
  484. def __ior__(self, other: Mapping[str, str]) -> MutableHeaders:
  485. if not isinstance(other, Mapping):
  486. raise TypeError(f"Expected a mapping but got {other.__class__.__name__}")
  487. self.update(other)
  488. return self
  489. def __or__(self, other: Mapping[str, str]) -> MutableHeaders:
  490. if not isinstance(other, Mapping):
  491. raise TypeError(f"Expected a mapping but got {other.__class__.__name__}")
  492. new = self.mutablecopy()
  493. new.update(other)
  494. return new
  495. @property
  496. def raw(self) -> list[tuple[bytes, bytes]]:
  497. return self._list
  498. def setdefault(self, key: str, value: str) -> str:
  499. """
  500. If the header `key` does not exist, then set it to `value`.
  501. Returns the header value.
  502. """
  503. set_key = key.lower().encode("latin-1")
  504. set_value = value.encode("latin-1")
  505. for idx, (item_key, item_value) in enumerate(self._list):
  506. if item_key == set_key:
  507. return item_value.decode("latin-1")
  508. self._list.append((set_key, set_value))
  509. return value
  510. def update(self, other: Mapping[str, str]) -> None:
  511. for key, val in other.items():
  512. self[key] = val
  513. def append(self, key: str, value: str) -> None:
  514. """
  515. Append a header, preserving any duplicate entries.
  516. """
  517. append_key = key.lower().encode("latin-1")
  518. append_value = value.encode("latin-1")
  519. self._list.append((append_key, append_value))
  520. def add_vary_header(self, vary: str) -> None:
  521. existing = self.get("vary")
  522. if existing is not None:
  523. vary = ", ".join([existing, vary])
  524. self["vary"] = vary
  525. class State:
  526. """
  527. An object that can be used to store arbitrary state.
  528. Used for `request.state` and `app.state`.
  529. """
  530. _state: dict[str, Any]
  531. def __init__(self, state: dict[str, Any] | None = None):
  532. if state is None:
  533. state = {}
  534. super().__setattr__("_state", state)
  535. def __setattr__(self, key: Any, value: Any) -> None:
  536. self._state[key] = value
  537. def __getattr__(self, key: Any) -> Any:
  538. try:
  539. return self._state[key]
  540. except KeyError:
  541. message = "'{}' object has no attribute '{}'"
  542. raise AttributeError(message.format(self.__class__.__name__, key))
  543. def __delattr__(self, key: Any) -> None:
  544. del self._state[key]
  545. def __getitem__(self, key: str) -> Any:
  546. return self._state[key]
  547. def __setitem__(self, key: str, value: Any) -> None:
  548. self._state[key] = value
  549. def __delitem__(self, key: str) -> None:
  550. del self._state[key]
  551. def __iter__(self) -> Iterator[str]:
  552. return iter(self._state)
  553. def __len__(self) -> int:
  554. return len(self._state)