codekingpro/portable-devtools

1from __future__ import annotations2 3import io4import sys5import typing6from base64 import b64encode7from enum import Enum8 9from ..exceptions import UnrewindableBodyError10from .util import to_bytes11 12if typing.TYPE_CHECKING:13    from typing import Final14 15# Pass as a value within ``headers`` to skip16# emitting some HTTP headers that are added automatically.17# The only headers that are supported are ``Accept-Encoding``,18# ``Host``, and ``User-Agent``.19SKIP_HEADER = "@@@SKIP_HEADER@@@"20SKIPPABLE_HEADERS = frozenset(["accept-encoding", "host", "user-agent"])21 22ACCEPT_ENCODING = "gzip,deflate"23try:24    try:25        import brotlicffi as _unused_module_brotli  # type: ignore[import-not-found] # noqa: F40126    except ImportError:27        import brotli as _unused_module_brotli  # type: ignore[import-not-found] # noqa: F40128except ImportError:29    pass30else:31    ACCEPT_ENCODING += ",br"32 33try:34    if sys.version_info >= (3, 14):35        from compression import zstd as _unused_module_zstd  # noqa: F40136    else:37        from backports import zstd as _unused_module_zstd  # noqa: F40138except ImportError:39    pass40else:41    ACCEPT_ENCODING += ",zstd"42 43 44class _TYPE_FAILEDTELL(Enum):45    token = 046 47 48_FAILEDTELL: Final[_TYPE_FAILEDTELL] = _TYPE_FAILEDTELL.token49 50_TYPE_BODY_POSITION = typing.Union[int, _TYPE_FAILEDTELL]51 52# When sending a request with these methods we aren't expecting53# a body so don't need to set an explicit 'Content-Length: 0'54# The reason we do this in the negative instead of tracking methods55# which 'should' have a body is because unknown methods should be56# treated as if they were 'POST' which *does* expect a body.57_METHODS_NOT_EXPECTING_BODY = {"GET", "HEAD", "DELETE", "TRACE", "OPTIONS", "CONNECT"}58 59 60def make_headers(61    keep_alive: bool | None = None,62    accept_encoding: bool | list[str] | str | None = None,63    user_agent: str | None = None,64    basic_auth: str | None = None,65    proxy_basic_auth: str | None = None,66    disable_cache: bool | None = None,67) -> dict[str, str]:68    """69    Shortcuts for generating request headers.70 71    :param keep_alive:72        If ``True``, adds 'connection: keep-alive' header.73 74    :param accept_encoding:75        Can be a boolean, list, or string.76        ``True`` translates to 'gzip,deflate'.  If the dependencies for77        Brotli (either the ``brotli`` or ``brotlicffi`` package) and/or78        Zstandard (the ``backports.zstd`` package for Python before 3.14)79        algorithms are installed, then their encodings are80        included in the string ('br' and 'zstd', respectively).81        List will get joined by comma.82        String will be used as provided.83 84    :param user_agent:85        String representing the user-agent you want, such as86        "python-urllib3/0.6"87 88    :param basic_auth:89        Colon-separated username:password string for 'authorization: basic ...'90        auth header.91 92    :param proxy_basic_auth:93        Colon-separated username:password string for 'proxy-authorization: basic ...'94        auth header.95 96    :param disable_cache:97        If ``True``, adds 'cache-control: no-cache' header.98 99    Example:100 101    .. code-block:: python102 103        import urllib3104 105        print(urllib3.util.make_headers(keep_alive=True, user_agent="Batman/1.0"))106        # {'connection': 'keep-alive', 'user-agent': 'Batman/1.0'}107        print(urllib3.util.make_headers(accept_encoding=True))108        # {'accept-encoding': 'gzip,deflate'}109    """110    headers: dict[str, str] = {}111    if accept_encoding:112        if isinstance(accept_encoding, str):113            pass114        elif isinstance(accept_encoding, list):115            accept_encoding = ",".join(accept_encoding)116        else:117            accept_encoding = ACCEPT_ENCODING118        headers["accept-encoding"] = accept_encoding119 120    if user_agent:121        headers["user-agent"] = user_agent122 123    if keep_alive:124        headers["connection"] = "keep-alive"125 126    if basic_auth:127        headers["authorization"] = (128            f"Basic {b64encode(basic_auth.encode('latin-1')).decode()}"129        )130 131    if proxy_basic_auth:132        headers["proxy-authorization"] = (133            f"Basic {b64encode(proxy_basic_auth.encode('latin-1')).decode()}"134        )135 136    if disable_cache:137        headers["cache-control"] = "no-cache"138 139    return headers140 141 142def set_file_position(143    body: typing.Any, pos: _TYPE_BODY_POSITION | None144) -> _TYPE_BODY_POSITION | None:145    """146    If a position is provided, move file to that point.147    Otherwise, we'll attempt to record a position for future use.148    """149    if pos is not None:150        rewind_body(body, pos)151    elif getattr(body, "tell", None) is not None:152        try:153            pos = body.tell()154        except OSError:155            # This differentiates from None, allowing us to catch156            # a failed `tell()` later when trying to rewind the body.157            pos = _FAILEDTELL158 159    return pos160 161 162def rewind_body(body: typing.IO[typing.AnyStr], body_pos: _TYPE_BODY_POSITION) -> None:163    """164    Attempt to rewind body to a certain position.165    Primarily used for request redirects and retries.166 167    :param body:168        File-like object that supports seek.169 170    :param int pos:171        Position to seek to in file.172    """173    body_seek = getattr(body, "seek", None)174    if body_seek is not None and isinstance(body_pos, int):175        try:176            body_seek(body_pos)177        except OSError as e:178            raise UnrewindableBodyError(179                "An error occurred when rewinding request body for redirect/retry."180            ) from e181    elif body_pos is _FAILEDTELL:182        raise UnrewindableBodyError(183            "Unable to record file position for rewinding "184            "request body during a redirect/retry."185        )186    else:187        raise ValueError(188            f"body_pos must be of type integer, instead it was {type(body_pos)}."189        )190 191 192class ChunksAndContentLength(typing.NamedTuple):193    chunks: typing.Iterable[bytes] | None194    content_length: int | None195 196 197def body_to_chunks(198    body: typing.Any | None, method: str, blocksize: int199) -> ChunksAndContentLength:200    """Takes the HTTP request method, body, and blocksize and201    transforms them into an iterable of chunks to pass to202    socket.sendall() and an optional 'Content-Length' header.203 204    A 'Content-Length' of 'None' indicates the length of the body205    can't be determined so should use 'Transfer-Encoding: chunked'206    for framing instead.207    """208 209    chunks: typing.Iterable[bytes] | None210    content_length: int | None211 212    # No body, we need to make a recommendation on 'Content-Length'213    # based on whether that request method is expected to have214    # a body or not.215    if body is None:216        chunks = None217        if method.upper() not in _METHODS_NOT_EXPECTING_BODY:218            content_length = 0219        else:220            content_length = None221 222    # Bytes or strings become bytes223    elif isinstance(body, (str, bytes)):224        chunks = (to_bytes(body),)225        content_length = len(chunks[0])226 227    # File-like object, TODO: use seek() and tell() for length?228    elif hasattr(body, "read"):229 230        def chunk_readable() -> typing.Iterable[bytes]:231            encode = isinstance(body, io.TextIOBase)232            while True:233                datablock = body.read(blocksize)234                if not datablock:235                    break236                if encode:237                    datablock = datablock.encode("utf-8")238                yield datablock239 240        chunks = chunk_readable()241        content_length = None242 243    # Otherwise we need to start checking via duck-typing.244    else:245        try:246            # Check if the body implements the buffer API.247            mv = memoryview(body)248        except TypeError:249            try:250                # Check if the body is an iterable251                chunks = iter(body)252                content_length = None253            except TypeError:254                raise TypeError(255                    f"'body' must be a bytes-like object, file-like "256                    f"object, or iterable. Instead was {body!r}"257                ) from None258        else:259            # Since it implements the buffer API can be passed directly to socket.sendall()260            chunks = (body,)261            content_length = mv.nbytes262 263    return ChunksAndContentLength(chunks=chunks, content_length=content_length)264