|
from __future__ import annotations |
|
|
|
import io |
|
import typing |
|
from base64 import b64encode |
|
from enum import Enum |
|
|
|
from ..exceptions import UnrewindableBodyError |
|
from .util import to_bytes |
|
|
|
if typing.TYPE_CHECKING: |
|
from typing import Final |
|
|
|
|
|
|
|
|
|
|
|
SKIP_HEADER = "@@@SKIP_HEADER@@@" |
|
SKIPPABLE_HEADERS = frozenset(["accept-encoding", "host", "user-agent"]) |
|
|
|
ACCEPT_ENCODING = "gzip,deflate" |
|
try: |
|
try: |
|
import brotlicffi as _unused_module_brotli |
|
except ImportError: |
|
import brotli as _unused_module_brotli |
|
except ImportError: |
|
pass |
|
else: |
|
ACCEPT_ENCODING += ",br" |
|
try: |
|
import zstandard as _unused_module_zstd |
|
except ImportError: |
|
pass |
|
else: |
|
ACCEPT_ENCODING += ",zstd" |
|
|
|
|
|
class _TYPE_FAILEDTELL(Enum): |
|
token = 0 |
|
|
|
|
|
_FAILEDTELL: Final[_TYPE_FAILEDTELL] = _TYPE_FAILEDTELL.token |
|
|
|
_TYPE_BODY_POSITION = typing.Union[int, _TYPE_FAILEDTELL] |
|
|
|
|
|
|
|
|
|
|
|
|
|
_METHODS_NOT_EXPECTING_BODY = {"GET", "HEAD", "DELETE", "TRACE", "OPTIONS", "CONNECT"} |
|
|
|
|
|
def make_headers( |
|
keep_alive: bool | None = None, |
|
accept_encoding: bool | list[str] | str | None = None, |
|
user_agent: str | None = None, |
|
basic_auth: str | None = None, |
|
proxy_basic_auth: str | None = None, |
|
disable_cache: bool | None = None, |
|
) -> dict[str, str]: |
|
""" |
|
Shortcuts for generating request headers. |
|
|
|
:param keep_alive: |
|
If ``True``, adds 'connection: keep-alive' header. |
|
|
|
:param accept_encoding: |
|
Can be a boolean, list, or string. |
|
``True`` translates to 'gzip,deflate'. If either the ``brotli`` or |
|
``brotlicffi`` package is installed 'gzip,deflate,br' is used instead. |
|
List will get joined by comma. |
|
String will be used as provided. |
|
|
|
:param user_agent: |
|
String representing the user-agent you want, such as |
|
"python-urllib3/0.6" |
|
|
|
:param basic_auth: |
|
Colon-separated username:password string for 'authorization: basic ...' |
|
auth header. |
|
|
|
:param proxy_basic_auth: |
|
Colon-separated username:password string for 'proxy-authorization: basic ...' |
|
auth header. |
|
|
|
:param disable_cache: |
|
If ``True``, adds 'cache-control: no-cache' header. |
|
|
|
Example: |
|
|
|
.. code-block:: python |
|
|
|
import urllib3 |
|
|
|
print(urllib3.util.make_headers(keep_alive=True, user_agent="Batman/1.0")) |
|
# {'connection': 'keep-alive', 'user-agent': 'Batman/1.0'} |
|
print(urllib3.util.make_headers(accept_encoding=True)) |
|
# {'accept-encoding': 'gzip,deflate'} |
|
""" |
|
headers: dict[str, str] = {} |
|
if accept_encoding: |
|
if isinstance(accept_encoding, str): |
|
pass |
|
elif isinstance(accept_encoding, list): |
|
accept_encoding = ",".join(accept_encoding) |
|
else: |
|
accept_encoding = ACCEPT_ENCODING |
|
headers["accept-encoding"] = accept_encoding |
|
|
|
if user_agent: |
|
headers["user-agent"] = user_agent |
|
|
|
if keep_alive: |
|
headers["connection"] = "keep-alive" |
|
|
|
if basic_auth: |
|
headers[ |
|
"authorization" |
|
] = f"Basic {b64encode(basic_auth.encode('latin-1')).decode()}" |
|
|
|
if proxy_basic_auth: |
|
headers[ |
|
"proxy-authorization" |
|
] = f"Basic {b64encode(proxy_basic_auth.encode('latin-1')).decode()}" |
|
|
|
if disable_cache: |
|
headers["cache-control"] = "no-cache" |
|
|
|
return headers |
|
|
|
|
|
def set_file_position( |
|
body: typing.Any, pos: _TYPE_BODY_POSITION | None |
|
) -> _TYPE_BODY_POSITION | None: |
|
""" |
|
If a position is provided, move file to that point. |
|
Otherwise, we'll attempt to record a position for future use. |
|
""" |
|
if pos is not None: |
|
rewind_body(body, pos) |
|
elif getattr(body, "tell", None) is not None: |
|
try: |
|
pos = body.tell() |
|
except OSError: |
|
|
|
|
|
pos = _FAILEDTELL |
|
|
|
return pos |
|
|
|
|
|
def rewind_body(body: typing.IO[typing.AnyStr], body_pos: _TYPE_BODY_POSITION) -> None: |
|
""" |
|
Attempt to rewind body to a certain position. |
|
Primarily used for request redirects and retries. |
|
|
|
:param body: |
|
File-like object that supports seek. |
|
|
|
:param int pos: |
|
Position to seek to in file. |
|
""" |
|
body_seek = getattr(body, "seek", None) |
|
if body_seek is not None and isinstance(body_pos, int): |
|
try: |
|
body_seek(body_pos) |
|
except OSError as e: |
|
raise UnrewindableBodyError( |
|
"An error occurred when rewinding request body for redirect/retry." |
|
) from e |
|
elif body_pos is _FAILEDTELL: |
|
raise UnrewindableBodyError( |
|
"Unable to record file position for rewinding " |
|
"request body during a redirect/retry." |
|
) |
|
else: |
|
raise ValueError( |
|
f"body_pos must be of type integer, instead it was {type(body_pos)}." |
|
) |
|
|
|
|
|
class ChunksAndContentLength(typing.NamedTuple): |
|
chunks: typing.Iterable[bytes] | None |
|
content_length: int | None |
|
|
|
|
|
def body_to_chunks( |
|
body: typing.Any | None, method: str, blocksize: int |
|
) -> ChunksAndContentLength: |
|
"""Takes the HTTP request method, body, and blocksize and |
|
transforms them into an iterable of chunks to pass to |
|
socket.sendall() and an optional 'Content-Length' header. |
|
|
|
A 'Content-Length' of 'None' indicates the length of the body |
|
can't be determined so should use 'Transfer-Encoding: chunked' |
|
for framing instead. |
|
""" |
|
|
|
chunks: typing.Iterable[bytes] | None |
|
content_length: int | None |
|
|
|
|
|
|
|
|
|
if body is None: |
|
chunks = None |
|
if method.upper() not in _METHODS_NOT_EXPECTING_BODY: |
|
content_length = 0 |
|
else: |
|
content_length = None |
|
|
|
|
|
elif isinstance(body, (str, bytes)): |
|
chunks = (to_bytes(body),) |
|
content_length = len(chunks[0]) |
|
|
|
|
|
elif hasattr(body, "read"): |
|
|
|
def chunk_readable() -> typing.Iterable[bytes]: |
|
nonlocal body, blocksize |
|
encode = isinstance(body, io.TextIOBase) |
|
while True: |
|
datablock = body.read(blocksize) |
|
if not datablock: |
|
break |
|
if encode: |
|
datablock = datablock.encode("iso-8859-1") |
|
yield datablock |
|
|
|
chunks = chunk_readable() |
|
content_length = None |
|
|
|
|
|
else: |
|
try: |
|
|
|
mv = memoryview(body) |
|
except TypeError: |
|
try: |
|
|
|
chunks = iter(body) |
|
content_length = None |
|
except TypeError: |
|
raise TypeError( |
|
f"'body' must be a bytes-like object, file-like " |
|
f"object, or iterable. Instead was {body!r}" |
|
) from None |
|
else: |
|
|
|
chunks = (body,) |
|
content_length = mv.nbytes |
|
|
|
return ChunksAndContentLength(chunks=chunks, content_length=content_length) |
|
|