From c99c9f2c849b52c4645acbc5e1e1355ba1a0a425 Mon Sep 17 00:00:00 2001 From: lavanyagarg112 Date: Thu, 26 Sep 2024 11:15:42 +0800 Subject: [PATCH 1/3] Added filter enhancement option to remove keys from the response body --- httpie/cli/argparser.py | 13 +++- httpie/cli/constants.py | 7 ++ httpie/cli/definition.py | 15 ++++ httpie/output/models.py | 9 ++- httpie/output/streams.py | 146 ++++++++++++++++++++++++++++++++++++++- httpie/output/writer.py | 21 +++++- tests/utils/__init__.py | 2 +- 7 files changed, 208 insertions(+), 5 deletions(-) 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..16491bcb67 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,98 @@ 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: + if word in chunk_dict: + del chunk_dict[word] + 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 + 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 +344,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 diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index 0a9af608a5..7f10893570 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -389,7 +389,7 @@ def http( and '--traceback' not in args_with_config_defaults): add_to_args.append('--traceback') if not any('--timeout' in arg for arg in args_with_config_defaults): - add_to_args.append('--timeout=3') + add_to_args.append('--timeout=10') complete_args = [program_name, *add_to_args, *args] # print(' '.join(complete_args)) From 8377dedf4d82f06533a5c0ac05cbdbec1369da7d Mon Sep 17 00:00:00 2001 From: lavanyagarg112 Date: Thu, 26 Sep 2024 11:20:25 +0800 Subject: [PATCH 2/3] Revert back the timeout change --- tests/utils/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index 7f10893570..0a9af608a5 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -389,7 +389,7 @@ def http( and '--traceback' not in args_with_config_defaults): add_to_args.append('--traceback') if not any('--timeout' in arg for arg in args_with_config_defaults): - add_to_args.append('--timeout=10') + add_to_args.append('--timeout=3') complete_args = [program_name, *add_to_args, *args] # print(' '.join(complete_args)) From 17f875041d566f21a1eb03d609fd880cdf53e416 Mon Sep 17 00:00:00 2001 From: lavanyagarg112 Date: Thu, 26 Sep 2024 12:16:01 +0800 Subject: [PATCH 3/3] modify process_body function to correctly access the nested keys --- httpie/output/streams.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/httpie/output/streams.py b/httpie/output/streams.py index 16491bcb67..a20b6b774e 100644 --- a/httpie/output/streams.py +++ b/httpie/output/streams.py @@ -273,8 +273,18 @@ def process_body(self, chunk: Union[str, bytes]) -> bytes: chunk = self.decode_chunk(chunk) chunk_dict = eval(chunk) for word in self.filters: - if word in chunk_dict: - del chunk_dict[word] + 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) @@ -308,9 +318,10 @@ def process_body(self, chunk: Union[str, bytes]) -> bytes: temp_dict = temp_dict[subword] else: break - subword = splitwords[-1] - if subword in temp_dict: - del temp_dict[subword] + 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)