Added filter enhancement option to remove keys from the response body

This commit is contained in:
lavanyagarg112 2024-09-26 11:15:42 +08:00
parent f4cf43ecdd
commit c99c9f2c84
7 changed files with 208 additions and 5 deletions

View File

@ -18,7 +18,7 @@ from .argtypes import (
from .constants import ( from .constants import (
HTTP_GET, HTTP_POST, BASE_OUTPUT_OPTIONS, OUTPUT_OPTIONS, OUTPUT_OPTIONS_DEFAULT, HTTP_GET, HTTP_POST, BASE_OUTPUT_OPTIONS, OUTPUT_OPTIONS, OUTPUT_OPTIONS_DEFAULT,
OUTPUT_OPTIONS_DEFAULT_OFFLINE, OUTPUT_OPTIONS_DEFAULT_STDOUT_REDIRECTED, 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_CREDENTIALS,
SEPARATOR_GROUP_ALL_ITEMS, SEPARATOR_GROUP_DATA_ITEMS, URL_SCHEME_RE, SEPARATOR_GROUP_ALL_ITEMS, SEPARATOR_GROUP_DATA_ITEMS, URL_SCHEME_RE,
) )
@ -172,6 +172,7 @@ class HTTPieArgumentParser(BaseHTTPieArgumentParser):
self._setup_standard_streams() self._setup_standard_streams()
self._process_output_options() self._process_output_options()
self._process_pretty_options() self._process_pretty_options()
self._process_filter_options()
self._process_format_options() self._process_format_options()
self._guess_method() self._guess_method()
self._parse_items() self._parse_items()
@ -539,6 +540,16 @@ class HTTPieArgumentParser(BaseHTTPieArgumentParser):
# noinspection PyTypeChecker # noinspection PyTypeChecker
self.args.prettify = PRETTY_MAP[self.args.prettify] 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): def _process_download_options(self):
if self.args.offline: if self.args.offline:
self.args.download = False self.args.download = False

View File

@ -105,6 +105,13 @@ PRETTY_MAP = {
} }
PRETTY_STDOUT_TTY_ONLY = PrettyOptions.STDOUT_TTY_ONLY 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 = [ DEFAULT_FORMAT_OPTIONS = [
'headers.sort:true', 'headers.sort:true',

View File

@ -13,6 +13,7 @@ from httpie.cli.constants import (BASE_OUTPUT_OPTIONS, DEFAULT_FORMAT_OPTIONS,
OUT_RESP_HEAD, OUT_RESP_META, OUTPUT_OPTIONS, OUT_RESP_HEAD, OUT_RESP_META, OUTPUT_OPTIONS,
OUTPUT_OPTIONS_DEFAULT, PRETTY_MAP, OUTPUT_OPTIONS_DEFAULT, PRETTY_MAP,
PRETTY_STDOUT_TTY_ONLY, PRETTY_STDOUT_TTY_ONLY,
FILTER_STDOUT_TTY_ONLY,
SEPARATOR_GROUP_ALL_ITEMS, SEPARATOR_PROXY, SEPARATOR_GROUP_ALL_ITEMS, SEPARATOR_PROXY,
SORTED_FORMAT_OPTIONS_STRING, SORTED_FORMAT_OPTIONS_STRING,
UNSORTED_FORMAT_OPTIONS_STRING, RequestType) UNSORTED_FORMAT_OPTIONS_STRING, RequestType)
@ -303,6 +304,20 @@ output_processing.add_argument(
""", """,
) )
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( output_processing.add_argument(
'--style', '--style',
'-s', '-s',

View File

@ -2,7 +2,7 @@ import argparse
from typing import Any, Dict, Union, List, NamedTuple, Optional from typing import Any, Dict, Union, List, NamedTuple, Optional
from httpie.context import Environment 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.cli.argtypes import PARSED_DEFAULT_FORMAT_OPTIONS
from httpie.output.formatters.colors import AUTO_STYLE from httpie.output.formatters.colors import AUTO_STYLE
@ -18,6 +18,7 @@ class ProcessingOptions(NamedTuple):
stream: bool = False stream: bool = False
style: str = AUTO_STYLE style: str = AUTO_STYLE
prettify: Union[List[str], PrettyOptions] = PRETTY_STDOUT_TTY_ONLY prettify: Union[List[str], PrettyOptions] = PRETTY_STDOUT_TTY_ONLY
filtery: Union[List[str], FilterOptions] = FILTER_STDOUT_TTY_ONLY
response_mime: Optional[str] = None response_mime: Optional[str] = None
response_charset: Optional[str] = None response_charset: Optional[str] = None
@ -30,6 +31,12 @@ class ProcessingOptions(NamedTuple):
return PRETTY_MAP['all' if env.stdout_isatty else 'none'] return PRETTY_MAP['all' if env.stdout_isatty else 'none']
else: else:
return self.prettify 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 @classmethod
def from_raw_args(cls, options: argparse.Namespace) -> 'ProcessingOptions': def from_raw_args(cls, options: argparse.Namespace) -> 'ProcessingOptions':

View File

@ -1,6 +1,6 @@
from abc import ABCMeta, abstractmethod from abc import ABCMeta, abstractmethod
from itertools import chain 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 .processing import Conversion, Formatting
from ..context import Environment from ..context import Environment
@ -223,6 +223,98 @@ class PrettyStream(EncodedStream):
chunk = self.decode_chunk(chunk) chunk = self.decode_chunk(chunk)
chunk = self.formatting.format_body(content=chunk, mime=self.mime) chunk = self.formatting.format_body(content=chunk, mime=self.mime)
return smart_encode(chunk, self.output_encoding) 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): class BufferedPrettyStream(PrettyStream):
@ -252,3 +344,55 @@ class BufferedPrettyStream(PrettyStream):
self.mime, body = converter.convert(body) self.mime, body = converter.convert(body)
yield self.process_body(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)

View File

@ -15,7 +15,15 @@ from ..models import (
from .models import ProcessingOptions from .models import ProcessingOptions
from .processing import Conversion, Formatting from .processing import Conversion, Formatting
from .streams import ( from .streams import (
BaseStream, BufferedPrettyStream, EncodedStream, PrettyStream, RawStream, BaseStream,
BufferedPrettyStream,
EncodedStream,
PrettyStream,
RawStream,
FilterStream,
PrettyFilterStream,
PrettyBufferedFilterStream,
BufferedFilterStream,
) )
from ..utils import parse_content_type_header from ..utils import parse_content_type_header
@ -161,6 +169,7 @@ def get_stream_type_and_kwargs(
""" """
is_stream = processing_options.stream is_stream = processing_options.stream
prettify_groups = processing_options.get_prettify(env) prettify_groups = processing_options.get_prettify(env)
filtery_groups = processing_options.get_filtery(env)
if not is_stream and message_type is HTTPResponse: if not is_stream and message_type is HTTPResponse:
# If this is a response, then check the headers for determining # If this is a response, then check the headers for determining
# auto-streaming. # 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 return stream_class, stream_kwargs

View File

@ -389,7 +389,7 @@ def http(
and '--traceback' not in args_with_config_defaults): and '--traceback' not in args_with_config_defaults):
add_to_args.append('--traceback') add_to_args.append('--traceback')
if not any('--timeout' in arg for arg in args_with_config_defaults): 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] complete_args = [program_name, *add_to_args, *args]
# print(' '.join(complete_args)) # print(' '.join(complete_args))