Fix --raw with --chunked (#1254)

* Fix --raw with --chunked

* Better naming / annotations

* More annotations
This commit is contained in:
Batuhan Taskaya 2021-12-29 12:41:44 +03:00 committed by GitHub
parent 0e10e23dca
commit 4c56d894ba
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 133 additions and 75 deletions

View File

@ -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)

View File

@ -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,

View File

@ -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(

View File

@ -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):