_content.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  1. import inspect
  2. import warnings
  3. from json import dumps as json_dumps
  4. from typing import (
  5. Any,
  6. AsyncIterable,
  7. AsyncIterator,
  8. Dict,
  9. Iterable,
  10. Iterator,
  11. Mapping,
  12. Optional,
  13. Tuple,
  14. Union,
  15. )
  16. from urllib.parse import urlencode
  17. from ._exceptions import StreamClosed, StreamConsumed
  18. from ._multipart import MultipartStream
  19. from ._types import (
  20. AsyncByteStream,
  21. RequestContent,
  22. RequestData,
  23. RequestFiles,
  24. ResponseContent,
  25. SyncByteStream,
  26. )
  27. from ._utils import peek_filelike_length, primitive_value_to_str
  28. class ByteStream(AsyncByteStream, SyncByteStream):
  29. def __init__(self, stream: bytes) -> None:
  30. self._stream = stream
  31. def __iter__(self) -> Iterator[bytes]:
  32. yield self._stream
  33. async def __aiter__(self) -> AsyncIterator[bytes]:
  34. yield self._stream
  35. class IteratorByteStream(SyncByteStream):
  36. CHUNK_SIZE = 65_536
  37. def __init__(self, stream: Iterable[bytes]) -> None:
  38. self._stream = stream
  39. self._is_stream_consumed = False
  40. self._is_generator = inspect.isgenerator(stream)
  41. def __iter__(self) -> Iterator[bytes]:
  42. if self._is_stream_consumed and self._is_generator:
  43. raise StreamConsumed()
  44. self._is_stream_consumed = True
  45. if hasattr(self._stream, "read"):
  46. # File-like interfaces should use 'read' directly.
  47. chunk = self._stream.read(self.CHUNK_SIZE)
  48. while chunk:
  49. yield chunk
  50. chunk = self._stream.read(self.CHUNK_SIZE)
  51. else:
  52. # Otherwise iterate.
  53. for part in self._stream:
  54. yield part
  55. class AsyncIteratorByteStream(AsyncByteStream):
  56. CHUNK_SIZE = 65_536
  57. def __init__(self, stream: AsyncIterable[bytes]) -> None:
  58. self._stream = stream
  59. self._is_stream_consumed = False
  60. self._is_generator = inspect.isasyncgen(stream)
  61. async def __aiter__(self) -> AsyncIterator[bytes]:
  62. if self._is_stream_consumed and self._is_generator:
  63. raise StreamConsumed()
  64. self._is_stream_consumed = True
  65. if hasattr(self._stream, "aread"):
  66. # File-like interfaces should use 'aread' directly.
  67. chunk = await self._stream.aread(self.CHUNK_SIZE)
  68. while chunk:
  69. yield chunk
  70. chunk = await self._stream.aread(self.CHUNK_SIZE)
  71. else:
  72. # Otherwise iterate.
  73. async for part in self._stream:
  74. yield part
  75. class UnattachedStream(AsyncByteStream, SyncByteStream):
  76. """
  77. If a request or response is serialized using pickle, then it is no longer
  78. attached to a stream for I/O purposes. Any stream operations should result
  79. in `httpx.StreamClosed`.
  80. """
  81. def __iter__(self) -> Iterator[bytes]:
  82. raise StreamClosed()
  83. async def __aiter__(self) -> AsyncIterator[bytes]:
  84. raise StreamClosed()
  85. yield b"" # pragma: no cover
  86. def encode_content(
  87. content: Union[str, bytes, Iterable[bytes], AsyncIterable[bytes]],
  88. ) -> Tuple[Dict[str, str], Union[SyncByteStream, AsyncByteStream]]:
  89. if isinstance(content, (bytes, str)):
  90. body = content.encode("utf-8") if isinstance(content, str) else content
  91. content_length = len(body)
  92. headers = {"Content-Length": str(content_length)} if body else {}
  93. return headers, ByteStream(body)
  94. elif isinstance(content, Iterable) and not isinstance(content, dict):
  95. # `not isinstance(content, dict)` is a bit oddly specific, but it
  96. # catches a case that's easy for users to make in error, and would
  97. # otherwise pass through here, like any other bytes-iterable,
  98. # because `dict` happens to be iterable. See issue #2491.
  99. content_length_or_none = peek_filelike_length(content)
  100. if content_length_or_none is None:
  101. headers = {"Transfer-Encoding": "chunked"}
  102. else:
  103. headers = {"Content-Length": str(content_length_or_none)}
  104. return headers, IteratorByteStream(content) # type: ignore
  105. elif isinstance(content, AsyncIterable):
  106. headers = {"Transfer-Encoding": "chunked"}
  107. return headers, AsyncIteratorByteStream(content)
  108. raise TypeError(f"Unexpected type for 'content', {type(content)!r}")
  109. def encode_urlencoded_data(
  110. data: RequestData,
  111. ) -> Tuple[Dict[str, str], ByteStream]:
  112. plain_data = []
  113. for key, value in data.items():
  114. if isinstance(value, (list, tuple)):
  115. plain_data.extend([(key, primitive_value_to_str(item)) for item in value])
  116. else:
  117. plain_data.append((key, primitive_value_to_str(value)))
  118. body = urlencode(plain_data, doseq=True).encode("utf-8")
  119. content_length = str(len(body))
  120. content_type = "application/x-www-form-urlencoded"
  121. headers = {"Content-Length": content_length, "Content-Type": content_type}
  122. return headers, ByteStream(body)
  123. def encode_multipart_data(
  124. data: RequestData, files: RequestFiles, boundary: Optional[bytes]
  125. ) -> Tuple[Dict[str, str], MultipartStream]:
  126. multipart = MultipartStream(data=data, files=files, boundary=boundary)
  127. headers = multipart.get_headers()
  128. return headers, multipart
  129. def encode_text(text: str) -> Tuple[Dict[str, str], ByteStream]:
  130. body = text.encode("utf-8")
  131. content_length = str(len(body))
  132. content_type = "text/plain; charset=utf-8"
  133. headers = {"Content-Length": content_length, "Content-Type": content_type}
  134. return headers, ByteStream(body)
  135. def encode_html(html: str) -> Tuple[Dict[str, str], ByteStream]:
  136. body = html.encode("utf-8")
  137. content_length = str(len(body))
  138. content_type = "text/html; charset=utf-8"
  139. headers = {"Content-Length": content_length, "Content-Type": content_type}
  140. return headers, ByteStream(body)
  141. def encode_json(json: Any) -> Tuple[Dict[str, str], ByteStream]:
  142. body = json_dumps(json).encode("utf-8")
  143. content_length = str(len(body))
  144. content_type = "application/json"
  145. headers = {"Content-Length": content_length, "Content-Type": content_type}
  146. return headers, ByteStream(body)
  147. def encode_request(
  148. content: Optional[RequestContent] = None,
  149. data: Optional[RequestData] = None,
  150. files: Optional[RequestFiles] = None,
  151. json: Optional[Any] = None,
  152. boundary: Optional[bytes] = None,
  153. ) -> Tuple[Dict[str, str], Union[SyncByteStream, AsyncByteStream]]:
  154. """
  155. Handles encoding the given `content`, `data`, `files`, and `json`,
  156. returning a two-tuple of (<headers>, <stream>).
  157. """
  158. if data is not None and not isinstance(data, Mapping):
  159. # We prefer to separate `content=<bytes|str|byte iterator|bytes aiterator>`
  160. # for raw request content, and `data=<form data>` for url encoded or
  161. # multipart form content.
  162. #
  163. # However for compat with requests, we *do* still support
  164. # `data=<bytes...>` usages. We deal with that case here, treating it
  165. # as if `content=<...>` had been supplied instead.
  166. message = "Use 'content=<...>' to upload raw bytes/text content."
  167. warnings.warn(message, DeprecationWarning)
  168. return encode_content(data)
  169. if content is not None:
  170. return encode_content(content)
  171. elif files:
  172. return encode_multipart_data(data or {}, files, boundary)
  173. elif data:
  174. return encode_urlencoded_data(data)
  175. elif json is not None:
  176. return encode_json(json)
  177. return {}, ByteStream(b"")
  178. def encode_response(
  179. content: Optional[ResponseContent] = None,
  180. text: Optional[str] = None,
  181. html: Optional[str] = None,
  182. json: Optional[Any] = None,
  183. ) -> Tuple[Dict[str, str], Union[SyncByteStream, AsyncByteStream]]:
  184. """
  185. Handles encoding the given `content`, returning a two-tuple of
  186. (<headers>, <stream>).
  187. """
  188. if content is not None:
  189. return encode_content(content)
  190. elif text is not None:
  191. return encode_text(text)
  192. elif html is not None:
  193. return encode_html(html)
  194. elif json is not None:
  195. return encode_json(json)
  196. return {}, ByteStream(b"")