httpie-cli/httpie/output/writer.py

import argparse
import errno
from typing import IO, TextIO, Tuple, Type, Union

from ..cli.dicts import HTTPHeadersDict
from ..context import Environment
from ..models import (
    HTTPRequest,
    HTTPResponse,
    HTTPMessage,
    RequestsMessage,
    RequestsMessageKind,
    infer_requests_message_kind
)
from .processing import Conversion, Formatting
from .streams import (
    BaseStream, BufferedPrettyStream, EncodedStream, PrettyStream, RawStream,
)


MESSAGE_SEPARATOR = '\n\n'
MESSAGE_SEPARATOR_BYTES = MESSAGE_SEPARATOR.encode()


def write_message(
    requests_message: RequestsMessage,
    env: Environment,
    args: argparse.Namespace,
    with_headers=False,
    with_body=False,
):
    if not (with_body or with_headers):
        return
    write_stream_kwargs = {
        'stream': build_output_stream_for_message(
            args=args,
            env=env,
            requests_message=requests_message,
            with_body=with_body,
            with_headers=with_headers,
        ),
        # NOTE: `env.stdout` will in fact be `stderr` with `--download`
        'outfile': env.stdout,
        'flush': env.stdout_isatty or args.stream
    }
    try:
        if env.is_windows and 'colors' in args.prettify:
            write_stream_with_colors_win(**write_stream_kwargs)
        else:
            write_stream(**write_stream_kwargs)
    except OSError as e:
        show_traceback = args.debug or args.traceback
        if not show_traceback and e.errno == errno.EPIPE:
            # Ignore broken pipes unless --traceback.
            env.stderr.write('\n')
        else:
            raise


def write_stream(
    stream: BaseStream,
    outfile: Union[IO, TextIO],
    flush: bool
):
    """Write the output stream."""
    try:
        # Writing bytes so we use the buffer interface.
        buf = outfile.buffer
    except AttributeError:
        buf = outfile

    for chunk in stream:
        buf.write(chunk)
        if flush:
            outfile.flush()


def write_stream_with_colors_win(
    stream: 'BaseStream',
    outfile: TextIO,
    flush: bool
):
    """Like `write`, but colorized chunks are written as text
    directly to `outfile` to ensure it gets processed by colorama.
    Applies only to Windows and colorized terminal output.

    """
    color = b'\x1b['
    encoding = outfile.encoding
    for chunk in stream:
        if color in chunk:
            outfile.write(chunk.decode(encoding))
        else:
            outfile.buffer.write(chunk)
        if flush:
            outfile.flush()


def build_output_stream_for_message(
    args: argparse.Namespace,
    env: Environment,
    requests_message: RequestsMessage,
    with_headers: bool,
    with_body: bool,
):
    message_type = {
        RequestsMessageKind.REQUEST: HTTPRequest,
        RequestsMessageKind.RESPONSE: HTTPResponse,
    }[infer_requests_message_kind(requests_message)]
    stream_class, stream_kwargs = get_stream_type_and_kwargs(
        env=env,
        args=args,
        message_type=message_type,
        headers=requests_message.headers
    )
    yield from stream_class(
        msg=message_type(requests_message),
        with_headers=with_headers,
        with_body=with_body,
        **stream_kwargs,
    )
    if (env.stdout_isatty and with_body
            and not getattr(requests_message, 'is_body_upload_chunk', False)):
        # Ensure a blank line after the response body.
        # For terminal output only.
        yield MESSAGE_SEPARATOR_BYTES


def get_stream_type_and_kwargs(
    env: Environment,
    args: argparse.Namespace,
    message_type: Type[HTTPMessage],
    headers: HTTPHeadersDict,
) -> Tuple[Type['BaseStream'], dict]:
    """Pick the right stream type and kwargs for it based on `env` and `args`.

    """
    is_stream = args.stream
    if not is_stream and message_type is HTTPResponse:
        # If this is a response, then check the headers for determining
        # auto-streaming.
        is_stream = headers.get('Content-Type') == 'text/event-stream'

    if not env.stdout_isatty and not args.prettify:
        stream_class = RawStream
        stream_kwargs = {
            'chunk_size': (
                RawStream.CHUNK_SIZE_BY_LINE
                if is_stream
                else RawStream.CHUNK_SIZE
            )
        }
    else:
        stream_class = EncodedStream
        stream_kwargs = {
            'env': env,
        }
        if message_type is HTTPResponse:
            stream_kwargs.update({
                'mime_overwrite': args.response_mime,
                'encoding_overwrite': args.response_charset,
            })
        if args.prettify:
            stream_class = PrettyStream if is_stream else BufferedPrettyStream
            stream_kwargs.update({
                'conversion': Conversion(),
                'formatting': Formatting(
                    env=env,
                    groups=args.prettify,
                    color_scheme=args.style,
                    explicit_json=args.json,
                    format_options=args.format_options,
                )
            })

    return stream_class, stream_kwargs