jakubroztocil/httpie

import errno
import requests
from typing import Any, Dict, IO, Optional, TextIO, Tuple, Type, Union

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


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


def write_message(
    requests_message: RequestsMessage,
    env: Environment,
    output_options: OutputOptions,
    processing_options: ProcessingOptions,
    extra_stream_kwargs: Optional[Dict[str, Any]] = None
):
    if not output_options.any():
        return
    write_stream_kwargs = {
        'stream': build_output_stream_for_message(
            env=env,
            requests_message=requests_message,
            output_options=output_options,
            processing_options=processing_options,
            extra_stream_kwargs=extra_stream_kwargs
        ),
        # NOTE: `env.stdout` will in fact be `stderr` with `--download`
        'outfile': env.stdout,
        'flush': env.stdout_isatty or processing_options.stream
    }
    try:
        if env.is_windows and 'colors' in processing_options.get_prettify(env):
            write_stream_with_colors_win(**write_stream_kwargs)
        else:
            write_stream(**write_stream_kwargs)
    except OSError as e:
        if processing_options.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 write_raw_data(
    env: Environment,
    data: Any,
    *,
    processing_options: Optional[ProcessingOptions] = None,
    headers: Optional[HTTPHeadersDict] = None,
    stream_kwargs: Optional[Dict[str, Any]] = None
):
    msg = requests.PreparedRequest()
    msg.is_body_upload_chunk = True
    msg.body = data
    msg.headers = headers or HTTPHeadersDict()
    msg_output_options = OutputOptions.from_message(msg, body=True, headers=False)
    return write_message(
        requests_message=msg,
        env=env,
        output_options=msg_output_options,
        processing_options=processing_options or ProcessingOptions(),
        extra_stream_kwargs=stream_kwargs
    )


def build_output_stream_for_message(
    env: Environment,
    requests_message: RequestsMessage,
    output_options: OutputOptions,
    processing_options: ProcessingOptions,
    extra_stream_kwargs: Optional[Dict[str, Any]] = None
):
    message_type = {
        RequestsMessageKind.REQUEST: HTTPRequest,
        RequestsMessageKind.RESPONSE: HTTPResponse,
    }[output_options.kind]
    stream_class, stream_kwargs = get_stream_type_and_kwargs(
        env=env,
        processing_options=processing_options,
        message_type=message_type,
        headers=requests_message.headers
    )
    if extra_stream_kwargs:
        stream_kwargs.update(extra_stream_kwargs)
    yield from stream_class(
        msg=message_type(requests_message),
        output_options=output_options,
        **stream_kwargs,
    )
    if (env.stdout_isatty and output_options.body and not output_options.meta
            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,
    processing_options: ProcessingOptions,
    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 = processing_options.stream
    prettify_groups = processing_options.get_prettify(env)
    if not is_stream and message_type is HTTPResponse:
        # If this is a response, then check the headers for determining
        # auto-streaming.
        raw_content_type_header = headers.get('Content-Type', None)
        if raw_content_type_header:
            content_type_header, _ = parse_content_type_header(raw_content_type_header)
            is_stream = (content_type_header == 'text/event-stream')

    if not env.stdout_isatty and not prettify_groups:
        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': processing_options.response_mime,
                'encoding_overwrite': processing_options.response_charset,
            })
        if prettify_groups:
            stream_class = PrettyStream if is_stream else BufferedPrettyStream
            stream_kwargs.update({
                'conversion': Conversion(),
                'formatting': Formatting(
                    env=env,
                    groups=prettify_groups,
                    color_scheme=processing_options.style,
                    explicit_json=processing_options.json,
                    format_options=processing_options.format_options,
                )
            })

    return stream_class, stream_kwargs