From 4c56d894ba9e2bb1c097a3a6067006843ac2944d Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Wed, 29 Dec 2021 12:41:44 +0300 Subject: [PATCH] Fix --raw with --chunked (#1254) * Fix --raw with --chunked * Better naming / annotations * More annotations --- CHANGELOG.md | 1 + httpie/client.py | 36 ++++++---- httpie/uploads.py | 158 ++++++++++++++++++++++++++---------------- tests/test_uploads.py | 13 ++++ 4 files changed, 133 insertions(+), 75 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c831fd10..32324935 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ This project adheres to [Semantic Versioning](https://semver.org/). - Broken plugins will no longer crash the whole application. ([#1204](https://github.com/httpie/httpie/issues/1204)) - Fixed auto addition of XML declaration to every formatted XML response. ([#1156](https://github.com/httpie/httpie/issues/1156)) - Fixed highlighting when `Content-Type` specifies `charset`. ([#1242](https://github.com/httpie/httpie/issues/1242)) +- Fixed an unexpected crash when `--raw` is used with `--chunked`. ([#1253](https://github.com/httpie/httpie/issues/1253)) ## [2.6.0](https://github.com/httpie/httpie/compare/2.5.0...2.6.0) (2021-10-14) diff --git a/httpie/client.py b/httpie/client.py index ba0cc773..dcf24826 100644 --- a/httpie/client.py +++ b/httpie/client.py @@ -4,7 +4,7 @@ import json import sys from contextlib import contextmanager from pathlib import Path -from typing import Callable, Iterable +from typing import Any, Dict, Callable, Iterable from urllib.parse import urlparse, urlunparse import requests @@ -273,6 +273,24 @@ def make_send_kwargs_mergeable_from_env(args: argparse.Namespace) -> dict: } +def json_dict_to_request_body(data: Dict[str, Any]) -> str: + # Propagate the top-level list if there is only one + # item in the object, with an en empty key. + if len(data) == 1: + [(key, value)] = data.items() + if key == '' and isinstance(value, list): + data = value + + if data: + data = json.dumps(data) + else: + # We need to set data to an empty string to prevent requests + # from assigning an empty list to `response.request.data`. + data = '' + + return data + + def make_request_kwargs( args: argparse.Namespace, base_headers: HTTPHeadersDict = None, @@ -287,19 +305,7 @@ def make_request_kwargs( data = args.data auto_json = data and not args.form if (args.json or auto_json) and isinstance(data, dict): - # Propagate the top-level list if there is only one - # item in the object, with an en empty key. - if len(data) == 1: - [(key, value)] = data.items() - if key == '' and isinstance(value, list): - data = value - - if data: - data = json.dumps(data) - else: - # We need to set data to an empty string to prevent requests - # from assigning an empty list to `response.request.data`. - data = '' + data = json_dict_to_request_body(data) # Finalize headers. headers = make_default_headers(args) @@ -324,7 +330,7 @@ def make_request_kwargs( 'url': args.url, 'headers': headers, 'data': prepare_request_body( - body=data, + data, body_read_callback=request_body_read_callback, chunked=args.chunked, offline=args.offline, diff --git a/httpie/uploads.py b/httpie/uploads.py index c3c4e766..0bb307e5 100644 --- a/httpie/uploads.py +++ b/httpie/uploads.py @@ -1,5 +1,6 @@ import zlib -from typing import Callable, IO, Iterable, Tuple, Union, TYPE_CHECKING +import functools +from typing import Any, Callable, IO, Iterable, Optional, Tuple, Union, TYPE_CHECKING from urllib.parse import urlencode import requests @@ -11,7 +12,12 @@ if TYPE_CHECKING: from .cli.dicts import MultipartRequestDataDict, RequestDataDict -class ChunkedUploadStream: +class ChunkedStream: + def __iter__(self) -> Iterable[Union[str, bytes]]: + raise NotImplementedError + + +class ChunkedUploadStream(ChunkedStream): def __init__(self, stream: Iterable, callback: Callable): self.callback = callback self.stream = stream @@ -22,7 +28,7 @@ class ChunkedUploadStream: yield chunk -class ChunkedMultipartUploadStream: +class ChunkedMultipartUploadStream(ChunkedStream): chunk_size = 100 * 1024 def __init__(self, encoder: 'MultipartEncoder'): @@ -36,69 +42,101 @@ class ChunkedMultipartUploadStream: yield chunk +def as_bytes(data: Union[str, bytes]) -> bytes: + if isinstance(data, str): + return data.encode() + else: + return data + + +CallbackT = Callable[[bytes], bytes] + + +def _wrap_function_with_callback( + func: Callable[..., Any], + callback: CallbackT +) -> Callable[..., Any]: + @functools.wraps(func) + def wrapped(*args, **kwargs): + chunk = func(*args, **kwargs) + callback(chunk) + return chunk + return wrapped + + +def _prepare_file_for_upload( + file: Union[IO, 'MultipartEncoder'], + callback: CallbackT, + chunked: bool = False, + content_length_header_value: Optional[int] = None, +) -> Union[bytes, IO, ChunkedStream]: + if not super_len(file): + # Zero-length -> assume stdin. + if content_length_header_value is None and not chunked: + # Read the whole stdin to determine `Content-Length`. + # + # TODO: Instead of opt-in --chunked, consider making + # `Transfer-Encoding: chunked` for STDIN opt-out via + # something like --no-chunked. + # This would be backwards-incompatible so wait until v3.0.0. + # + file = as_bytes(file.read()) + else: + file.read = _wrap_function_with_callback( + file.read, + callback + ) + + if chunked: + from requests_toolbelt import MultipartEncoder + if isinstance(file, MultipartEncoder): + return ChunkedMultipartUploadStream( + encoder=file, + ) + else: + return ChunkedUploadStream( + stream=file, + callback=callback, + ) + else: + return file + + def prepare_request_body( - body: Union[str, bytes, IO, 'MultipartEncoder', RequestDataDict], - body_read_callback: Callable[[bytes], bytes], - content_length_header_value: int = None, - chunked=False, - offline=False, -) -> Union[str, bytes, IO, 'MultipartEncoder', ChunkedUploadStream]: - - is_file_like = hasattr(body, 'read') - - if isinstance(body, RequestDataDict): - body = urlencode(body, doseq=True) + raw_body: Union[str, bytes, IO, 'MultipartEncoder', RequestDataDict], + body_read_callback: CallbackT, + offline: bool = False, + chunked: bool = False, + content_length_header_value: Optional[int] = None, +) -> Union[bytes, IO, 'MultipartEncoder', ChunkedStream]: + is_file_like = hasattr(raw_body, 'read') + if isinstance(raw_body, (bytes, str)): + body = as_bytes(raw_body) + elif isinstance(raw_body, RequestDataDict): + body = as_bytes(urlencode(raw_body, doseq=True)) + else: + body = raw_body if offline: if is_file_like: - return body.read() - return body - - if not is_file_like: - if chunked: - body = ChunkedUploadStream( - # Pass the entire body as one chunk. - stream=(chunk.encode() for chunk in [body]), - callback=body_read_callback, - ) - else: - # File-like object. - - if not super_len(body): - # Zero-length -> assume stdin. - if content_length_header_value is None and not chunked: - # - # Read the whole stdin to determine `Content-Length`. - # - # TODO: Instead of opt-in --chunked, consider making - # `Transfer-Encoding: chunked` for STDIN opt-out via - # something like --no-chunked. - # This would be backwards-incompatible so wait until v3.0.0. - # - body = body.read() + return as_bytes(raw_body.read()) else: - orig_read = body.read + return body - def new_read(*args): - chunk = orig_read(*args) - body_read_callback(chunk) - return chunk - - body.read = new_read - - if chunked: - from requests_toolbelt import MultipartEncoder - if isinstance(body, MultipartEncoder): - body = ChunkedMultipartUploadStream( - encoder=body, - ) - else: - body = ChunkedUploadStream( - stream=body, - callback=body_read_callback, - ) - - return body + if is_file_like: + return _prepare_file_for_upload( + body, + chunked=chunked, + callback=body_read_callback, + content_length_header_value=content_length_header_value + ) + elif chunked: + return ChunkedUploadStream( + stream=iter([body]), + callback=body_read_callback + ) + else: + return body def get_multipart_data_and_content_type( diff --git a/tests/test_uploads.py b/tests/test_uploads.py index 2c701477..28d427ae 100644 --- a/tests/test_uploads.py +++ b/tests/test_uploads.py @@ -1,4 +1,5 @@ import os +import json import pytest @@ -70,6 +71,18 @@ def test_chunked_stdin_multiple_chunks(httpbin_with_chunked_support): assert r.count(FILE_CONTENT) == 4 +def test_chunked_raw(httpbin_with_chunked_support): + r = http( + '--verbose', + '--chunked', + httpbin_with_chunked_support + '/post', + '--raw', + json.dumps({'a': 1, 'b': '2fafds', 'c': '🥰'}), + ) + assert HTTP_OK in r + assert 'Transfer-Encoding: chunked' in r + + class TestMultipartFormDataFileUpload: def test_non_existent_file_raises_parse_error(self, httpbin):