mirror of
https://github.com/httpie/cli.git
synced 2024-12-04 09:43:15 +02:00
271 lines
7.7 KiB
Python
271 lines
7.7 KiB
Python
from itertools import chain
|
|
from functools import partial
|
|
|
|
from httpie.context import Environment
|
|
from httpie.models import HTTPRequest, HTTPResponse
|
|
from httpie.input import (OUT_REQ_BODY, OUT_REQ_HEAD,
|
|
OUT_RESP_HEAD, OUT_RESP_BODY)
|
|
from httpie.output.processors import ProcessorManager
|
|
|
|
|
|
BINARY_SUPPRESSED_NOTICE = (
|
|
b'\n'
|
|
b'+-----------------------------------------+\n'
|
|
b'| NOTE: binary data not shown in terminal |\n'
|
|
b'+-----------------------------------------+'
|
|
)
|
|
|
|
|
|
class BinarySuppressedError(Exception):
|
|
"""An error indicating that the body is binary and won't be written,
|
|
e.g., for terminal output)."""
|
|
|
|
message = BINARY_SUPPRESSED_NOTICE
|
|
|
|
|
|
def write(stream, outfile, flush):
|
|
"""Write the output stream."""
|
|
try:
|
|
# Writing bytes so we use the buffer interface (Python 3).
|
|
buf = outfile.buffer
|
|
except AttributeError:
|
|
buf = outfile
|
|
|
|
for chunk in stream:
|
|
buf.write(chunk)
|
|
if flush:
|
|
outfile.flush()
|
|
|
|
|
|
def write_with_colors_win_py3(stream, outfile, flush):
|
|
"""Like `write`, but colorized chunks are written as text
|
|
directly to `outfile` to ensure it gets processed by colorama.
|
|
Applies only to Windows with Python 3 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(args, env, request, response):
|
|
"""Build and return a chain of iterators over the `request`-`response`
|
|
exchange each of which yields `bytes` chunks.
|
|
|
|
"""
|
|
|
|
req_h = OUT_REQ_HEAD in args.output_options
|
|
req_b = OUT_REQ_BODY in args.output_options
|
|
resp_h = OUT_RESP_HEAD in args.output_options
|
|
resp_b = OUT_RESP_BODY in args.output_options
|
|
req = req_h or req_b
|
|
resp = resp_h or resp_b
|
|
|
|
output = []
|
|
Stream = get_stream_type(env, args)
|
|
|
|
if req:
|
|
output.append(Stream(
|
|
msg=HTTPRequest(request),
|
|
with_headers=req_h,
|
|
with_body=req_b))
|
|
|
|
if req_b and resp:
|
|
# Request/Response separator.
|
|
output.append([b'\n\n'])
|
|
|
|
if resp:
|
|
output.append(Stream(
|
|
msg=HTTPResponse(response),
|
|
with_headers=resp_h,
|
|
with_body=resp_b))
|
|
|
|
if env.stdout_isatty and resp_b:
|
|
# Ensure a blank line after the response body.
|
|
# For terminal output only.
|
|
output.append([b'\n\n'])
|
|
|
|
return chain(*output)
|
|
|
|
|
|
def get_stream_type(env, args):
|
|
"""Pick the right stream type based on `env` and `args`.
|
|
Wrap it in a partial with the type-specific args so that
|
|
we don't need to think what stream we are dealing with.
|
|
|
|
"""
|
|
if not env.stdout_isatty and not args.prettify:
|
|
Stream = partial(
|
|
RawStream,
|
|
chunk_size=RawStream.CHUNK_SIZE_BY_LINE
|
|
if args.stream
|
|
else RawStream.CHUNK_SIZE
|
|
)
|
|
elif args.prettify:
|
|
Stream = partial(
|
|
PrettyStream if args.stream else BufferedPrettyStream,
|
|
env=env,
|
|
processor=ProcessorManager(
|
|
env=env,
|
|
groups=args.prettify,
|
|
pygments_style=args.style
|
|
),
|
|
)
|
|
else:
|
|
Stream = partial(EncodedStream, env=env)
|
|
|
|
return Stream
|
|
|
|
|
|
class BaseStream(object):
|
|
"""Base HTTP message output stream class."""
|
|
|
|
def __init__(self, msg, with_headers=True, with_body=True,
|
|
on_body_chunk_downloaded=None):
|
|
"""
|
|
:param msg: a :class:`models.HTTPMessage` subclass
|
|
:param with_headers: if `True`, headers will be included
|
|
:param with_body: if `True`, body will be included
|
|
|
|
"""
|
|
assert with_headers or with_body
|
|
self.msg = msg
|
|
self.with_headers = with_headers
|
|
self.with_body = with_body
|
|
self.on_body_chunk_downloaded = on_body_chunk_downloaded
|
|
|
|
def _get_headers(self):
|
|
"""Return the headers' bytes."""
|
|
return self.msg.headers.encode('utf8')
|
|
|
|
def _iter_body(self):
|
|
"""Return an iterator over the message body."""
|
|
raise NotImplementedError()
|
|
|
|
def __iter__(self):
|
|
"""Return an iterator over `self.msg`."""
|
|
if self.with_headers:
|
|
yield self._get_headers()
|
|
yield b'\r\n\r\n'
|
|
|
|
if self.with_body:
|
|
try:
|
|
for chunk in self._iter_body():
|
|
yield chunk
|
|
if self.on_body_chunk_downloaded:
|
|
self.on_body_chunk_downloaded(chunk)
|
|
except BinarySuppressedError as e:
|
|
if self.with_headers:
|
|
yield b'\n'
|
|
yield e.message
|
|
|
|
|
|
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(RawStream, self).__init__(**kwargs)
|
|
self.chunk_size = chunk_size
|
|
|
|
def _iter_body(self):
|
|
return self.msg.iter_body(self.chunk_size)
|
|
|
|
|
|
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(), **kwargs):
|
|
|
|
super(EncodedStream, self).__init__(**kwargs)
|
|
|
|
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 utf8 when unsure.
|
|
self.output_encoding = output_encoding or 'utf8'
|
|
|
|
def _iter_body(self):
|
|
|
|
for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):
|
|
|
|
if b'\0' in line:
|
|
raise BinarySuppressedError()
|
|
|
|
yield line.decode(self.msg.encoding) \
|
|
.encode(self.output_encoding, 'replace') + lf
|
|
|
|
|
|
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, processor, **kwargs):
|
|
super(PrettyStream, self).__init__(**kwargs)
|
|
self.processor = processor
|
|
|
|
def _get_headers(self):
|
|
return self.processor.process_headers(
|
|
self.msg.headers).encode(self.output_encoding)
|
|
|
|
def _iter_body(self):
|
|
for line, lf in self.msg.iter_lines(self.CHUNK_SIZE):
|
|
if b'\0' in line:
|
|
raise BinarySuppressedError()
|
|
yield self._process_body(line) + lf
|
|
|
|
def _process_body(self, chunk):
|
|
return self.processor.process_body(
|
|
body=chunk.decode(self.msg.encoding, 'replace'),
|
|
mime=self.msg.content_type.split(';')[0]
|
|
).encode(self.output_encoding, 'replace')
|
|
|
|
|
|
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):
|
|
|
|
# Read the whole body before prettifying it,
|
|
# but bail out immediately if the body is binary.
|
|
body = bytearray()
|
|
for chunk in self.msg.iter_body(self.CHUNK_SIZE):
|
|
if b'\0' in chunk:
|
|
raise BinarySuppressedError()
|
|
body.extend(chunk)
|
|
|
|
yield self._process_body(body)
|