Skip to content

Commit

Permalink
Added filter enhancement option to remove keys from the response body
Browse files Browse the repository at this point in the history
  • Loading branch information
lavanyagarg112 committed Sep 26, 2024
1 parent f4cf43e commit c99c9f2
Show file tree
Hide file tree
Showing 7 changed files with 208 additions and 5 deletions.
13 changes: 12 additions & 1 deletion httpie/cli/argparser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions httpie/cli/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
15 changes: 15 additions & 0 deletions httpie/cli/definition.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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',
Expand Down
9 changes: 8 additions & 1 deletion httpie/output/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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':
Expand Down
146 changes: 145 additions & 1 deletion httpie/output/streams.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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)
21 changes: 20 additions & 1 deletion httpie/output/writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion tests/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down

0 comments on commit c99c9f2

Please sign in to comment.