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