codekingpro/portable-devtools

1import errno2import requests3from typing import Any, Dict, IO, Optional, TextIO, Tuple, Type, Union4 5from ..cli.dicts import HTTPHeadersDict6from ..context import Environment7from ..models import (8    HTTPRequest,9    HTTPResponse,10    HTTPMessage,11    RequestsMessage,12    RequestsMessageKind,13    OutputOptions,14)15from .models import ProcessingOptions16from .processing import Conversion, Formatting17from .streams import (18    BaseStream, BufferedPrettyStream, EncodedStream, PrettyStream, RawStream,19)20from ..utils import parse_content_type_header21 22 23MESSAGE_SEPARATOR = '\n\n'24MESSAGE_SEPARATOR_BYTES = MESSAGE_SEPARATOR.encode()25 26 27def write_message(28    requests_message: RequestsMessage,29    env: Environment,30    output_options: OutputOptions,31    processing_options: ProcessingOptions,32    extra_stream_kwargs: Optional[Dict[str, Any]] = None33):34    if not output_options.any():35        return36    write_stream_kwargs = {37        'stream': build_output_stream_for_message(38            env=env,39            requests_message=requests_message,40            output_options=output_options,41            processing_options=processing_options,42            extra_stream_kwargs=extra_stream_kwargs43        ),44        # NOTE: `env.stdout` will in fact be `stderr` with `--download`45        'outfile': env.stdout,46        'flush': env.stdout_isatty or processing_options.stream47    }48    try:49        if env.is_windows and 'colors' in processing_options.get_prettify(env):50            write_stream_with_colors_win(**write_stream_kwargs)51        else:52            write_stream(**write_stream_kwargs)53    except OSError as e:54        if processing_options.show_traceback and e.errno == errno.EPIPE:55            # Ignore broken pipes unless --traceback.56            env.stderr.write('\n')57        else:58            raise59 60 61def write_stream(62    stream: BaseStream,63    outfile: Union[IO, TextIO],64    flush: bool65):66    """Write the output stream."""67    try:68        # Writing bytes so we use the buffer interface.69        buf = outfile.buffer70    except AttributeError:71        buf = outfile72 73    for chunk in stream:74        buf.write(chunk)75        if flush:76            outfile.flush()77 78 79def write_stream_with_colors_win(80    stream: 'BaseStream',81    outfile: TextIO,82    flush: bool83):84    """Like `write`, but colorized chunks are written as text85    directly to `outfile` to ensure it gets processed by colorama.86    Applies only to Windows and colorized terminal output.87 88    """89    color = b'\x1b['90    encoding = outfile.encoding91    for chunk in stream:92        if color in chunk:93            outfile.write(chunk.decode(encoding))94        else:95            outfile.buffer.write(chunk)96        if flush:97            outfile.flush()98 99 100def write_raw_data(101    env: Environment,102    data: Any,103    *,104    processing_options: Optional[ProcessingOptions] = None,105    headers: Optional[HTTPHeadersDict] = None,106    stream_kwargs: Optional[Dict[str, Any]] = None107):108    msg = requests.PreparedRequest()109    msg.is_body_upload_chunk = True110    msg.body = data111    msg.headers = headers or HTTPHeadersDict()112    msg_output_options = OutputOptions.from_message(msg, body=True, headers=False)113    return write_message(114        requests_message=msg,115        env=env,116        output_options=msg_output_options,117        processing_options=processing_options or ProcessingOptions(),118        extra_stream_kwargs=stream_kwargs119    )120 121 122def build_output_stream_for_message(123    env: Environment,124    requests_message: RequestsMessage,125    output_options: OutputOptions,126    processing_options: ProcessingOptions,127    extra_stream_kwargs: Optional[Dict[str, Any]] = None128):129    message_type = {130        RequestsMessageKind.REQUEST: HTTPRequest,131        RequestsMessageKind.RESPONSE: HTTPResponse,132    }[output_options.kind]133    stream_class, stream_kwargs = get_stream_type_and_kwargs(134        env=env,135        processing_options=processing_options,136        message_type=message_type,137        headers=requests_message.headers138    )139    if extra_stream_kwargs:140        stream_kwargs.update(extra_stream_kwargs)141    yield from stream_class(142        msg=message_type(requests_message),143        output_options=output_options,144        **stream_kwargs,145    )146    if (env.stdout_isatty and output_options.body and not output_options.meta147            and not getattr(requests_message, 'is_body_upload_chunk', False)):148        # Ensure a blank line after the response body.149        # For terminal output only.150        yield MESSAGE_SEPARATOR_BYTES151 152 153def get_stream_type_and_kwargs(154    env: Environment,155    processing_options: ProcessingOptions,156    message_type: Type[HTTPMessage],157    headers: HTTPHeadersDict,158) -> Tuple[Type['BaseStream'], dict]:159    """Pick the right stream type and kwargs for it based on `env` and `args`.160 161    """162    is_stream = processing_options.stream163    prettify_groups = processing_options.get_prettify(env)164    if not is_stream and message_type is HTTPResponse:165        # If this is a response, then check the headers for determining166        # auto-streaming.167        raw_content_type_header = headers.get('Content-Type', None)168        if raw_content_type_header:169            content_type_header, _ = parse_content_type_header(raw_content_type_header)170            is_stream = (content_type_header == 'text/event-stream')171 172    if not env.stdout_isatty and not prettify_groups:173        stream_class = RawStream174        stream_kwargs = {175            'chunk_size': (176                RawStream.CHUNK_SIZE_BY_LINE177                if is_stream178                else RawStream.CHUNK_SIZE179            )180        }181    else:182        stream_class = EncodedStream183        stream_kwargs = {184            'env': env,185        }186        if message_type is HTTPResponse:187            stream_kwargs.update({188                'mime_overwrite': processing_options.response_mime,189                'encoding_overwrite': processing_options.response_charset,190            })191        if prettify_groups:192            stream_class = PrettyStream if is_stream else BufferedPrettyStream193            stream_kwargs.update({194                'conversion': Conversion(),195                'formatting': Formatting(196                    env=env,197                    groups=prettify_groups,198                    color_scheme=processing_options.style,199                    explicit_json=processing_options.json,200                    format_options=processing_options.format_options,201                )202            })203 204    return stream_class, stream_kwargs205