codekingpro/portable-devtools

1import sys2import os3import zlib4import functools5import threading6from typing import Any, Callable, IO, Iterable, Optional, Tuple, Union, TYPE_CHECKING7from urllib.parse import urlencode8 9import requests10from requests.utils import super_len11 12if TYPE_CHECKING:13    from requests_toolbelt import MultipartEncoder14 15from .context import Environment16from .cli.dicts import MultipartRequestDataDict, RequestDataDict17from .compat import is_windows18 19 20class ChunkedStream:21    def __iter__(self) -> Iterable[Union[str, bytes]]:22        raise NotImplementedError23 24 25class ChunkedUploadStream(ChunkedStream):26    def __init__(27        self,28        stream: Iterable,29        callback: Callable,30        event: Optional[threading.Event] = None31    ) -> None:32        self.callback = callback33        self.stream = stream34        self.event = event35 36    def __iter__(self) -> Iterable[Union[str, bytes]]:37        for chunk in self.stream:38            if self.event:39                self.event.set()40            self.callback(chunk)41            yield chunk42 43 44class ChunkedMultipartUploadStream(ChunkedStream):45    chunk_size = 100 * 102446 47    def __init__(48        self,49        encoder: 'MultipartEncoder',50        event: Optional[threading.Event] = None51    ) -> None:52        self.encoder = encoder53        self.event = event54 55    def __iter__(self) -> Iterable[Union[str, bytes]]:56        while True:57            chunk = self.encoder.read(self.chunk_size)58            if self.event:59                self.event.set()60            if not chunk:61                break62            yield chunk63 64 65def as_bytes(data: Union[str, bytes]) -> bytes:66    if isinstance(data, str):67        return data.encode()68    else:69        return data70 71 72CallbackT = Callable[[bytes], bytes]73 74 75def _wrap_function_with_callback(76    func: Callable[..., Any],77    callback: CallbackT78) -> Callable[..., Any]:79    @functools.wraps(func)80    def wrapped(*args, **kwargs):81        chunk = func(*args, **kwargs)82        callback(chunk)83        return chunk84    return wrapped85 86 87def is_stdin(file: IO) -> bool:88    try:89        file_no = file.fileno()90    except Exception:91        return False92    else:93        return file_no == sys.stdin.fileno()94 95 96READ_THRESHOLD = float(os.getenv('HTTPIE_STDIN_READ_WARN_THRESHOLD', 10.0))97 98 99def observe_stdin_for_data_thread(env: Environment, file: IO, read_event: threading.Event) -> None:100    # Windows unfortunately does not support select() operation101    # on regular files, like stdin in our use case.102    # https://docs.python.org/3/library/select.html#select.select103    if is_windows:104        return None105 106    # If the user configures READ_THRESHOLD to be 0, then107    # disable this warning.108    if READ_THRESHOLD == 0:109        return None110 111    def worker(event: threading.Event) -> None:112        if not event.wait(timeout=READ_THRESHOLD):113            env.stderr.write(114                f'> warning: no stdin data read in {READ_THRESHOLD}s '115                f'(perhaps you want to --ignore-stdin)\n'116                f'> See: https://httpie.io/docs/cli/best-practices\n'117            )118 119    # Making it a daemon ensures that if the user exits from the main program120    # (e.g. either regularly or with Ctrl-C), the thread will not121    # block them.122    thread = threading.Thread(123        target=worker,124        args=(read_event,),125        daemon=True126    )127    thread.start()128 129 130def _read_file_with_selectors(file: IO, read_event: threading.Event) -> bytes:131    if is_windows or not is_stdin(file):132        return as_bytes(file.read())133 134    import select135 136    # Try checking whether there is any incoming data for READ_THRESHOLD137    # seconds. If there isn't anything in the given period, issue138    # a warning about a misusage.139    read_selectors, _, _ = select.select([file], [], [], READ_THRESHOLD)140    if read_selectors:141        read_event.set()142 143    return as_bytes(file.read())144 145 146def _prepare_file_for_upload(147    env: Environment,148    file: Union[IO, 'MultipartEncoder'],149    callback: CallbackT,150    chunked: bool = False,151    content_length_header_value: Optional[int] = None,152) -> Union[bytes, IO, ChunkedStream]:153    read_event = threading.Event()154    if not super_len(file):155        if is_stdin(file):156            observe_stdin_for_data_thread(env, file, read_event)157 158        # Zero-length -> assume stdin.159        if content_length_header_value is None and not chunked:160            # Read the whole stdin to determine `Content-Length`.161            #162            # TODO: Instead of opt-in --chunked, consider making163            #   `Transfer-Encoding: chunked` for STDIN opt-out via164            #   something like --no-chunked.165            #   This would be backwards-incompatible so wait until v3.0.0.166            #167            file = _read_file_with_selectors(file, read_event)168    else:169        file.read = _wrap_function_with_callback(170            file.read,171            callback172        )173 174    if chunked:175        from requests_toolbelt import MultipartEncoder176        if isinstance(file, MultipartEncoder):177            return ChunkedMultipartUploadStream(178                encoder=file,179                event=read_event,180            )181        else:182            return ChunkedUploadStream(183                stream=file,184                callback=callback,185                event=read_event186            )187    else:188        return file189 190 191def prepare_request_body(192    env: Environment,193    raw_body: Union[str, bytes, IO, 'MultipartEncoder', RequestDataDict],194    body_read_callback: CallbackT,195    offline: bool = False,196    chunked: bool = False,197    content_length_header_value: Optional[int] = None,198) -> Union[bytes, IO, 'MultipartEncoder', ChunkedStream]:199    is_file_like = hasattr(raw_body, 'read')200    if isinstance(raw_body, (bytes, str)):201        body = as_bytes(raw_body)202    elif isinstance(raw_body, RequestDataDict):203        body = as_bytes(urlencode(raw_body, doseq=True))204    else:205        body = raw_body206 207    if offline:208        if is_file_like:209            return as_bytes(raw_body.read())210        else:211            return body212 213    if is_file_like:214        return _prepare_file_for_upload(215            env,216            body,217            chunked=chunked,218            callback=body_read_callback,219            content_length_header_value=content_length_header_value220        )221    elif chunked:222        return ChunkedUploadStream(223            stream=iter([body]),224            callback=body_read_callback225        )226    else:227        return body228 229 230def get_multipart_data_and_content_type(231    data: MultipartRequestDataDict,232    boundary: str = None,233    content_type: str = None,234) -> Tuple['MultipartEncoder', str]:235    from requests_toolbelt import MultipartEncoder236 237    encoder = MultipartEncoder(238        fields=data.items(),239        boundary=boundary,240    )241    if content_type:242        content_type = content_type.strip()243        if 'boundary=' not in content_type:244            content_type = f'{content_type}; boundary={encoder.boundary_value}'245    else:246        content_type = encoder.content_type247 248    data = encoder249    return data, content_type250 251 252def compress_request(253    request: requests.PreparedRequest,254    always: bool,255):256    deflater = zlib.compressobj()257    if isinstance(request.body, str):258        body_bytes = request.body.encode()259    elif hasattr(request.body, 'read'):260        body_bytes = request.body.read()261    else:262        body_bytes = request.body263    deflated_data = deflater.compress(body_bytes)264    deflated_data += deflater.flush()265    is_economical = len(deflated_data) < len(body_bytes)266    if is_economical or always:267        request.body = deflated_data268        request.headers['Content-Encoding'] = 'deflate'269        request.headers['Content-Length'] = str(len(deflated_data))270