mirror of
https://github.com/httpie/cli.git
synced 2024-11-22 15:53:13 +01:00
255 lines
8.0 KiB
Python
255 lines
8.0 KiB
Python
from abc import ABCMeta, abstractmethod
|
|
from itertools import chain
|
|
from typing import Callable, Iterable, Optional, Union
|
|
|
|
from .processing import Conversion, Formatting
|
|
from ..context import Environment
|
|
from ..encoding import smart_decode, smart_encode, UTF8
|
|
from ..models import HTTPMessage, OutputOptions
|
|
from ..utils import parse_content_type_header
|
|
|
|
|
|
BINARY_SUPPRESSED_NOTICE = (
|
|
b'\n'
|
|
b'+-----------------------------------------+\n'
|
|
b'| NOTE: binary data not shown in terminal |\n'
|
|
b'+-----------------------------------------+'
|
|
)
|
|
|
|
|
|
class DataSuppressedError(Exception):
|
|
message = None
|
|
|
|
|
|
class BinarySuppressedError(DataSuppressedError):
|
|
"""An error indicating that the body is binary and won't be written,
|
|
e.g., for terminal output)."""
|
|
message = BINARY_SUPPRESSED_NOTICE
|
|
|
|
|
|
class BaseStream(metaclass=ABCMeta):
|
|
"""Base HTTP message output stream class."""
|
|
|
|
def __init__(
|
|
self,
|
|
msg: HTTPMessage,
|
|
output_options: OutputOptions,
|
|
on_body_chunk_downloaded: Callable[[bytes], None] = None,
|
|
**kwargs
|
|
):
|
|
"""
|
|
:param msg: a :class:`models.HTTPMessage` subclass
|
|
:param output_options: a :class:`OutputOptions` instance to represent
|
|
which parts of the message is printed.
|
|
"""
|
|
assert output_options.any()
|
|
self.msg = msg
|
|
self.output_options = output_options
|
|
self.on_body_chunk_downloaded = on_body_chunk_downloaded
|
|
self.extra_options = kwargs
|
|
|
|
def get_headers(self) -> bytes:
|
|
"""Return the headers' bytes."""
|
|
return self.msg.headers.encode()
|
|
|
|
def get_metadata(self) -> bytes:
|
|
"""Return the message metadata."""
|
|
return self.msg.metadata.encode()
|
|
|
|
@abstractmethod
|
|
def iter_body(self) -> Iterable[bytes]:
|
|
"""Return an iterator over the message body."""
|
|
|
|
def __iter__(self) -> Iterable[bytes]:
|
|
"""Return an iterator over `self.msg`."""
|
|
if self.output_options.headers:
|
|
yield self.get_headers()
|
|
yield b'\r\n\r\n'
|
|
|
|
if self.output_options.body:
|
|
try:
|
|
for chunk in self.iter_body():
|
|
yield chunk
|
|
if self.on_body_chunk_downloaded:
|
|
self.on_body_chunk_downloaded(chunk)
|
|
except DataSuppressedError as e:
|
|
if self.output_options.headers:
|
|
yield b'\n'
|
|
yield e.message
|
|
|
|
if self.output_options.meta:
|
|
if self.output_options.body:
|
|
yield b'\n\n'
|
|
|
|
yield self.get_metadata()
|
|
yield b'\n\n'
|
|
|
|
|
|
class RawStream(BaseStream):
|
|
"""The message is streamed in chunks with no processing."""
|
|
|
|
CHUNK_SIZE = 1024 * 100
|
|
CHUNK_SIZE_BY_LINE = 1
|
|
|
|
def __init__(self, chunk_size=CHUNK_SIZE, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.chunk_size = chunk_size
|
|
|
|
def iter_body(self) -> Iterable[bytes]:
|
|
return self.msg.iter_body(self.chunk_size)
|
|
|
|
|
|
ENCODING_GUESS_THRESHOLD = 3
|
|
|
|
|
|
class EncodedStream(BaseStream):
|
|
"""Encoded HTTP message stream.
|
|
|
|
The message bytes are converted to an encoding suitable for
|
|
`self.env.stdout`. Unicode errors are replaced and binary data
|
|
is suppressed. The body is always streamed by line.
|
|
|
|
"""
|
|
CHUNK_SIZE = 1
|
|
|
|
def __init__(
|
|
self,
|
|
env=Environment(),
|
|
mime_overwrite: str = None,
|
|
encoding_overwrite: str = None,
|
|
**kwargs
|
|
):
|
|
super().__init__(**kwargs)
|
|
if mime_overwrite:
|
|
self.mime = mime_overwrite
|
|
else:
|
|
self.mime, _ = parse_content_type_header(self.msg.content_type)
|
|
self._encoding = encoding_overwrite or self.msg.encoding
|
|
self._encoding_guesses = []
|
|
if env.stdout_isatty:
|
|
# Use the encoding supported by the terminal.
|
|
output_encoding = env.stdout_encoding
|
|
else:
|
|
# Preserve the message encoding.
|
|
output_encoding = self.msg.encoding
|
|
# Default to UTF-8 when unsure.
|
|
self.output_encoding = output_encoding or UTF8
|
|
|
|
def iter_body(self) -> Iterable[bytes]:
|
|
for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):
|
|
if b'\0' in line:
|
|
raise BinarySuppressedError()
|
|
line = self.decode_chunk(line)
|
|
yield smart_encode(line, self.output_encoding) + lf
|
|
|
|
def decode_chunk(self, raw_chunk: str) -> str:
|
|
chunk, guessed_encoding = smart_decode(raw_chunk, self.encoding)
|
|
self._encoding_guesses.append(guessed_encoding)
|
|
return chunk
|
|
|
|
@property
|
|
def encoding(self) -> Optional[str]:
|
|
if self._encoding:
|
|
return self._encoding
|
|
|
|
# If we find a reliable (used consecutively) encoding, than
|
|
# use it for the next iterations.
|
|
if len(self._encoding_guesses) < ENCODING_GUESS_THRESHOLD:
|
|
return None
|
|
|
|
guess_1, guess_2 = self._encoding_guesses[-2:]
|
|
if guess_1 == guess_2:
|
|
self._encoding = guess_1
|
|
return guess_1
|
|
|
|
@encoding.setter
|
|
def encoding(self, value) -> None:
|
|
self._encoding = value
|
|
|
|
|
|
class PrettyStream(EncodedStream):
|
|
"""In addition to :class:`EncodedStream` behaviour, this stream applies
|
|
content processing.
|
|
|
|
Useful for long-lived HTTP responses that stream by lines
|
|
such as the Twitter streaming API.
|
|
|
|
"""
|
|
|
|
CHUNK_SIZE = 1
|
|
|
|
def __init__(
|
|
self, conversion: Conversion,
|
|
formatting: Formatting,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.formatting = formatting
|
|
self.conversion = conversion
|
|
|
|
def get_headers(self) -> bytes:
|
|
return self.formatting.format_headers(
|
|
self.msg.headers).encode(self.output_encoding)
|
|
|
|
def get_metadata(self) -> bytes:
|
|
return self.formatting.format_metadata(
|
|
self.msg.metadata).encode(self.output_encoding)
|
|
|
|
def iter_body(self) -> Iterable[bytes]:
|
|
first_chunk = True
|
|
iter_lines = self.msg.iter_lines(self.CHUNK_SIZE)
|
|
for line, lf in iter_lines:
|
|
if b'\0' in line:
|
|
if first_chunk:
|
|
converter = self.conversion.get_converter(self.mime)
|
|
if converter:
|
|
body = bytearray()
|
|
# noinspection PyAssignmentToLoopOrWithParameter
|
|
for line, lf in chain([(line, lf)], iter_lines):
|
|
body.extend(line)
|
|
body.extend(lf)
|
|
self.mime, body = converter.convert(body)
|
|
assert isinstance(body, str)
|
|
yield self.process_body(body)
|
|
return
|
|
raise BinarySuppressedError()
|
|
yield self.process_body(line) + lf
|
|
first_chunk = False
|
|
|
|
def process_body(self, chunk: Union[str, bytes]) -> bytes:
|
|
if not isinstance(chunk, str):
|
|
# Text when a converter has been used,
|
|
# otherwise it will always be bytes.
|
|
chunk = self.decode_chunk(chunk)
|
|
chunk = self.formatting.format_body(content=chunk, mime=self.mime)
|
|
return smart_encode(chunk, self.output_encoding)
|
|
|
|
|
|
class BufferedPrettyStream(PrettyStream):
|
|
"""The same as :class:`PrettyStream` except that the body is fully
|
|
fetched before it's processed.
|
|
|
|
Suitable regular HTTP responses.
|
|
|
|
"""
|
|
|
|
CHUNK_SIZE = 1024 * 10
|
|
|
|
def iter_body(self) -> Iterable[bytes]:
|
|
# Read the whole body before prettifying it,
|
|
# but bail out immediately if the body is binary.
|
|
converter = None
|
|
body = bytearray()
|
|
|
|
for chunk in self.msg.iter_body(self.CHUNK_SIZE):
|
|
if not converter and b'\0' in chunk:
|
|
converter = self.conversion.get_converter(self.mime)
|
|
if not converter:
|
|
raise BinarySuppressedError()
|
|
body.extend(chunk)
|
|
|
|
if converter:
|
|
self.mime, body = converter.convert(body)
|
|
|
|
yield self.process_body(body)
|