codekingpro/portable-devtools

1from __future__ import annotations2 3import inspect4import warnings5from json import dumps as json_dumps6from typing import (7    Any,8    AsyncIterable,9    AsyncIterator,10    Iterable,11    Iterator,12    Mapping,13)14from urllib.parse import urlencode15 16from ._exceptions import StreamClosed, StreamConsumed17from ._multipart import MultipartStream18from ._types import (19    AsyncByteStream,20    RequestContent,21    RequestData,22    RequestFiles,23    ResponseContent,24    SyncByteStream,25)26from ._utils import peek_filelike_length, primitive_value_to_str27 28__all__ = ["ByteStream"]29 30 31class ByteStream(AsyncByteStream, SyncByteStream):32    def __init__(self, stream: bytes) -> None:33        self._stream = stream34 35    def __iter__(self) -> Iterator[bytes]:36        yield self._stream37 38    async def __aiter__(self) -> AsyncIterator[bytes]:39        yield self._stream40 41 42class IteratorByteStream(SyncByteStream):43    CHUNK_SIZE = 65_53644 45    def __init__(self, stream: Iterable[bytes]) -> None:46        self._stream = stream47        self._is_stream_consumed = False48        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 = True55        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 chunk60                chunk = self._stream.read(self.CHUNK_SIZE)61        else:62            # Otherwise iterate.63            for part in self._stream:64                yield part65 66 67class AsyncIteratorByteStream(AsyncByteStream):68    CHUNK_SIZE = 65_53669 70    def __init__(self, stream: AsyncIterable[bytes]) -> None:71        self._stream = stream72        self._is_stream_consumed = False73        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 = True80        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 chunk85                chunk = await self._stream.aread(self.CHUNK_SIZE)86        else:87            # Otherwise iterate.88            async for part in self._stream:89                yield part90 91 92class UnattachedStream(AsyncByteStream, SyncByteStream):93    """94    If a request or response is serialized using pickle, then it is no longer95    attached to a stream for I/O purposes. Any stream operations should result96    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 cover105 106 107def 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 content112        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 it118        # catches a case that's easy for users to make in error, and would119        # 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: ignore128 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 136def 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 152def encode_multipart_data(153    data: RequestData, files: RequestFiles, boundary: bytes | None154) -> tuple[dict[str, str], MultipartStream]:155    multipart = MultipartStream(data=data, files=files, boundary=boundary)156    headers = multipart.get_headers()157    return headers, multipart158 159 160def 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 168def 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 176def encode_json(json: Any) -> tuple[dict[str, str], ByteStream]:177    body = json_dumps(178        json, ensure_ascii=False, separators=(",", ":"), allow_nan=False179    ).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 186def 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 or200        # multipart form content.201        #202        # However for compat with requests, we *do* still support203        # `data=<bytes...>` usages. We deal with that case here, treating it204        # 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 221def 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 of229    (<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