diff --git a/httpie/cli/argparser.py b/httpie/cli/argparser.py index 9bf09b3b73..abc1e841d2 100644 --- a/httpie/cli/argparser.py +++ b/httpie/cli/argparser.py @@ -18,7 +18,7 @@ from .constants import ( HTTP_GET, HTTP_POST, BASE_OUTPUT_OPTIONS, OUTPUT_OPTIONS, OUTPUT_OPTIONS_DEFAULT, OUTPUT_OPTIONS_DEFAULT_OFFLINE, OUTPUT_OPTIONS_DEFAULT_STDOUT_REDIRECTED, - OUT_RESP_BODY, PRETTY_MAP, PRETTY_STDOUT_TTY_ONLY, RequestType, + OUT_RESP_BODY, PRETTY_MAP, PRETTY_STDOUT_TTY_ONLY, FILTER_STDOUT_TTY_ONLY, RequestType, SEPARATOR_CREDENTIALS, SEPARATOR_GROUP_ALL_ITEMS, SEPARATOR_GROUP_DATA_ITEMS, URL_SCHEME_RE, ) @@ -172,6 +172,7 @@ def parse_args( self._setup_standard_streams() self._process_output_options() self._process_pretty_options() + self._process_filter_options() self._process_format_options() self._guess_method() self._parse_items() @@ -539,6 +540,16 @@ def _process_pretty_options(self): # noinspection PyTypeChecker self.args.prettify = PRETTY_MAP[self.args.prettify] + def _process_filter_options(self): + if self.args.filtery == FILTER_STDOUT_TTY_ONLY: + pass + elif (self.args.filtery and self.env.is_windows + and self.args.output_file): + self.error('Only terminal output can be colorized on Windows.') + else: + # noinspection PyTypeChecker + pass + def _process_download_options(self): if self.args.offline: self.args.download = False diff --git a/httpie/cli/constants.py b/httpie/cli/constants.py index 09ca19e4af..ab1cf80b98 100644 --- a/httpie/cli/constants.py +++ b/httpie/cli/constants.py @@ -105,6 +105,13 @@ class PrettyOptions(enum.Enum): } PRETTY_STDOUT_TTY_ONLY = PrettyOptions.STDOUT_TTY_ONLY +# FILTER + +class FilterOptions(enum.Enum): + STDOUT_TTY_ONLY = enum.auto() + +FILTER_STDOUT_TTY_ONLY = FilterOptions.STDOUT_TTY_ONLY + DEFAULT_FORMAT_OPTIONS = [ 'headers.sort:true', diff --git a/httpie/cli/definition.py b/httpie/cli/definition.py index 843b29c9cf..473f3233e3 100644 --- a/httpie/cli/definition.py +++ b/httpie/cli/definition.py @@ -13,6 +13,7 @@ OUT_RESP_HEAD, OUT_RESP_META, OUTPUT_OPTIONS, OUTPUT_OPTIONS_DEFAULT, PRETTY_MAP, PRETTY_STDOUT_TTY_ONLY, + FILTER_STDOUT_TTY_ONLY, SEPARATOR_GROUP_ALL_ITEMS, SEPARATOR_PROXY, SORTED_FORMAT_OPTIONS_STRING, UNSORTED_FORMAT_OPTIONS_STRING, RequestType) @@ -303,6 +304,20 @@ def format_style_help(available_styles, *, isolation_mode: bool = False): """, ) + +output_processing.add_argument( + '--filter-keys', + dest='filtery', + default=FILTER_STDOUT_TTY_ONLY, + # choices=sorted(PRETTY_MAP.keys()), + short_help='Control the processing of console outputs.', + help=""" + Controls output processing. Filters the comma separated + key values in the output json + + """, +) + output_processing.add_argument( '--style', '-s', diff --git a/httpie/output/models.py b/httpie/output/models.py index e65199b07b..9818746903 100644 --- a/httpie/output/models.py +++ b/httpie/output/models.py @@ -2,7 +2,7 @@ from typing import Any, Dict, Union, List, NamedTuple, Optional from httpie.context import Environment -from httpie.cli.constants import PrettyOptions, PRETTY_MAP, PRETTY_STDOUT_TTY_ONLY +from httpie.cli.constants import PrettyOptions, PRETTY_MAP, PRETTY_STDOUT_TTY_ONLY, FilterOptions, FILTER_STDOUT_TTY_ONLY from httpie.cli.argtypes import PARSED_DEFAULT_FORMAT_OPTIONS from httpie.output.formatters.colors import AUTO_STYLE @@ -18,6 +18,7 @@ class ProcessingOptions(NamedTuple): stream: bool = False style: str = AUTO_STYLE prettify: Union[List[str], PrettyOptions] = PRETTY_STDOUT_TTY_ONLY + filtery: Union[List[str], FilterOptions] = FILTER_STDOUT_TTY_ONLY response_mime: Optional[str] = None response_charset: Optional[str] = None @@ -30,6 +31,12 @@ def get_prettify(self, env: Environment) -> List[str]: return PRETTY_MAP['all' if env.stdout_isatty else 'none'] else: return self.prettify + + def get_filtery(self, env: Environment) -> List[str]: + if self.filtery is FILTER_STDOUT_TTY_ONLY: + return [] + else: + return self.filtery.split(",") @classmethod def from_raw_args(cls, options: argparse.Namespace) -> 'ProcessingOptions': diff --git a/httpie/output/streams.py b/httpie/output/streams.py index 811093808a..a20b6b774e 100644 --- a/httpie/output/streams.py +++ b/httpie/output/streams.py @@ -1,6 +1,6 @@ from abc import ABCMeta, abstractmethod from itertools import chain -from typing import Callable, Iterable, Optional, Union +from typing import Callable, Iterable, List, Optional, Union from .processing import Conversion, Formatting from ..context import Environment @@ -223,6 +223,109 @@ def process_body(self, chunk: Union[str, bytes]) -> bytes: chunk = self.decode_chunk(chunk) chunk = self.formatting.format_body(content=chunk, mime=self.mime) return smart_encode(chunk, self.output_encoding) + +class FilterStream(EncodedStream): + + CHUNK_SIZE = 1 + + def __init__( + self, + conversion: Conversion, + filter: List['str'], + **kwargs, + ): + super().__init__(**kwargs) + self.filters = filter + self.conversion = conversion + + def get_headers(self) -> bytes: + return self.msg.headers.encode(self.output_encoding) + + + def get_metadata(self) -> bytes: + return 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) + 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_dict = eval(chunk) + for word in self.filters: + temp_dict = chunk_dict + splitwords = word.split(".") + for i in range(len(splitwords)-1): + subword = splitwords[i] + if subword in temp_dict: + temp_dict = temp_dict[subword] + else: + break + else: + subword = splitwords[-1] + if subword in temp_dict: + del temp_dict[subword] + chunk = f'{chunk_dict}' + return smart_encode(chunk, self.output_encoding) + + +class PrettyFilterStream(PrettyStream): + + CHUNK_SIZE = 1 + + def __init__( + self, + conversion: Conversion, + formatting: Formatting, + filter: List['str'], + **kwargs, + ): + super().__init__(conversion=conversion, formatting=formatting, **kwargs) + self.filters = filter + + 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_dict = eval(chunk) + for word in self.filters: + temp_dict = chunk_dict + splitwords = word.split(".") + for i in range(len(splitwords)-1): + subword = splitwords[i] + if subword in temp_dict: + temp_dict = temp_dict[subword] + else: + break + else: + subword = splitwords[-1] + if subword in temp_dict: + del temp_dict[subword] + chunk = (f'{chunk_dict}').replace(" ", "").replace("'", '"') + chunk = self.formatting.format_body(content=chunk, mime=self.mime) + return smart_encode(chunk, self.output_encoding) + class BufferedPrettyStream(PrettyStream): @@ -252,3 +355,55 @@ def iter_body(self) -> Iterable[bytes]: self.mime, body = converter.convert(body) yield self.process_body(body) + + +class PrettyBufferedFilterStream(PrettyFilterStream): + + 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) + + +class BufferedFilterStream(FilterStream): + """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) \ No newline at end of file diff --git a/httpie/output/writer.py b/httpie/output/writer.py index 4a2949bce2..786161e795 100644 --- a/httpie/output/writer.py +++ b/httpie/output/writer.py @@ -15,7 +15,15 @@ from .models import ProcessingOptions from .processing import Conversion, Formatting from .streams import ( - BaseStream, BufferedPrettyStream, EncodedStream, PrettyStream, RawStream, + BaseStream, + BufferedPrettyStream, + EncodedStream, + PrettyStream, + RawStream, + FilterStream, + PrettyFilterStream, + PrettyBufferedFilterStream, + BufferedFilterStream, ) from ..utils import parse_content_type_header @@ -161,6 +169,7 @@ def get_stream_type_and_kwargs( """ is_stream = processing_options.stream prettify_groups = processing_options.get_prettify(env) + filtery_groups = processing_options.get_filtery(env) if not is_stream and message_type is HTTPResponse: # If this is a response, then check the headers for determining # auto-streaming. @@ -201,4 +210,14 @@ def get_stream_type_and_kwargs( ) }) + if filtery_groups: + stream_class = FilterStream if is_stream else BufferedFilterStream + stream_kwargs.update({ + 'conversion': Conversion(), + 'filter': filtery_groups + }) + + if prettify_groups and filtery_groups: + stream_class = PrettyFilterStream if is_stream else PrettyBufferedFilterStream + return stream_class, stream_kwargs