From 6cd934d1b863188bf4acc69371c7728b137b0fc6 Mon Sep 17 00:00:00 2001 From: Jakub Roztocil Date: Sat, 15 Aug 2020 17:50:00 +0200 Subject: [PATCH] Add support for multipart upload streaming Close #684, #201 --- CHANGELOG.rst | 2 ++ README.rst | 4 ++++ extras/brew-deps.py | 1 + httpie/cli/requestitems.py | 6 ++---- httpie/client.py | 10 ++++++++-- httpie/output/streams.py | 27 ++++++++++++++++++++++++--- httpie/uploads.py | 20 ++++++++++++++++++++ httpie/utils.py | 2 ++ setup.py | 1 + tests/test_uploads.py | 25 ++++++++++++++++++++++--- 10 files changed, 86 insertions(+), 12 deletions(-) create mode 100644 httpie/uploads.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6ceb2923..b2c279bd 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ This project adheres to `Semantic Versioning `_. `2.3.0-dev`_ (unreleased) ------------------------- +* Added support for multipart upload streaming (`#684`_). * Added support for combining cookies specified on the CLI and in a session file (`#932`_). * Added out of the box SOCKS support with no extra installation (`#904`_). * Added ``--quiet, -q`` flag to enforce silent behaviour. @@ -451,6 +452,7 @@ This project adheres to `Semantic Versioning `_. .. _#128: https://github.com/jakubroztocil/httpie/issues/128 .. _#488: https://github.com/jakubroztocil/httpie/issues/488 .. _#668: https://github.com/jakubroztocil/httpie/issues/668 +.. _#684: https://github.com/jakubroztocil/httpie/issues/684 .. _#718: https://github.com/jakubroztocil/httpie/issues/718 .. _#719: https://github.com/jakubroztocil/httpie/issues/719 .. _#840: https://github.com/jakubroztocil/httpie/issues/840 diff --git a/README.rst b/README.rst index ec274279..a32c81fd 100644 --- a/README.rst +++ b/README.rst @@ -694,6 +694,10 @@ override the inferred content type: $ http -f POST httpbin.org/post name='John Smith' cv@'~/files/data.bin;type=application/pdf' +Larger multipart uploads (i.e., ``--form`` requests with at least one ``file@path``) +are always streamed to avoid memory issues. Additionally, the display of the +request body on the terminal is suppressed. + HTTP headers ============ diff --git a/extras/brew-deps.py b/extras/brew-deps.py index eb541ebd..3b666c26 100755 --- a/extras/brew-deps.py +++ b/extras/brew-deps.py @@ -16,6 +16,7 @@ PACKAGES = [ 'httpie', 'Pygments', 'requests', + 'requests-toolbelt', 'certifi', 'urllib3', 'idna', diff --git a/httpie/cli/requestitems.py b/httpie/cli/requestitems.py index a812d651..46d50903 100644 --- a/httpie/cli/requestitems.py +++ b/httpie/cli/requestitems.py @@ -99,15 +99,13 @@ def process_file_upload_arg(arg: KeyValueArg) -> Tuple[str, IO, str]: parts = arg.value.split(SEPARATOR_FILE_UPLOAD_TYPE) filename = parts[0] mime_type = parts[1] if len(parts) > 1 else None - try: - with open(os.path.expanduser(filename), 'rb') as f: - contents = f.read() + f = open(os.path.expanduser(filename), 'rb') except IOError as e: raise ParseError('"%s": %s' % (arg.orig, e)) return ( os.path.basename(filename), - BytesIO(contents), + f, mime_type or get_content_type(filename), ) diff --git a/httpie/client.py b/httpie/client.py index ff596bbd..d696eed6 100644 --- a/httpie/client.py +++ b/httpie/client.py @@ -11,14 +11,15 @@ from urllib.parse import urlparse, urlunparse import requests # noinspection PyPackageRequirements import urllib3 - from httpie import __version__ from httpie.cli.dicts import RequestHeadersDict from httpie.plugins.registry import plugin_manager from httpie.sessions import get_httpie_session from httpie.ssl import AVAILABLE_SSL_VERSION_ARG_MAPPING, HTTPieHTTPSAdapter +from httpie.uploads import get_multipart_data from httpie.utils import get_expired_cookies, repr_dict + urllib3.disable_warnings() FORM_CONTENT_TYPE = 'application/x-www-form-urlencoded; charset=utf-8' @@ -256,6 +257,7 @@ def make_request_kwargs( Translate our `args` into `requests.Request` keyword arguments. """ + files = args.files # Serialize JSON data, if needed. data = args.data auto_json = data and not args.form @@ -274,6 +276,10 @@ def make_request_kwargs( headers.update(args.headers) headers = finalize_headers(headers) + if args.form and files: + data, headers['Content-Type'] = get_multipart_data(data, files) + files = None + kwargs = { 'method': args.method.lower(), 'url': args.url, @@ -281,7 +287,7 @@ def make_request_kwargs( 'data': data, 'auth': args.auth, 'params': args.params, - 'files': args.files, + 'files': files, } return kwargs diff --git a/httpie/output/streams.py b/httpie/output/streams.py index b1f810b6..51dc58cc 100644 --- a/httpie/output/streams.py +++ b/httpie/output/streams.py @@ -4,6 +4,7 @@ from typing import Callable, Iterable, Union from httpie.context import Environment from httpie.models import HTTPMessage from httpie.output.processing import Conversion, Formatting +from requests_toolbelt import MultipartEncoder BINARY_SUPPRESSED_NOTICE = ( @@ -14,11 +15,29 @@ BINARY_SUPPRESSED_NOTICE = ( ) -class BinarySuppressedError(Exception): +LARGE_UPLOAD_SUPPRESSED_NOTICE = ( + b'\n' + b'+--------------------------------------------------------+\n' + b'| NOTE: large form upload data not shown in the terminal |\n' + b'+--------------------------------------------------------+' +) + + +class DataSuppressedError(Exception): + message = None + + +class BinarySuppressedError(DataSuppressedError): + """An error indicating that the body is binary and won't be written, + e.g., for terminal output).""" + message = BINARY_SUPPRESSED_NOTICE + + +class LargeUploadSuppressedError(DataSuppressedError): """An error indicating that the body is binary and won't be written, e.g., for terminal output).""" - message = BINARY_SUPPRESSED_NOTICE + message = LARGE_UPLOAD_SUPPRESSED_NOTICE class BaseStream: @@ -63,7 +82,7 @@ class BaseStream: yield chunk if self.on_body_chunk_downloaded: self.on_body_chunk_downloaded(chunk) - except BinarySuppressedError as e: + except DataSuppressedError as e: if self.with_headers: yield b'\n' yield e.message @@ -184,6 +203,8 @@ class BufferedPrettyStream(PrettyStream): body = bytearray() for chunk in self.msg.iter_body(self.CHUNK_SIZE): + if isinstance(chunk, MultipartEncoder): + raise LargeUploadSuppressedError() if not converter and b'\0' in chunk: converter = self.conversion.get_converter(self.mime) if not converter: diff --git a/httpie/uploads.py b/httpie/uploads.py new file mode 100644 index 00000000..8ae10b36 --- /dev/null +++ b/httpie/uploads.py @@ -0,0 +1,20 @@ +from typing import Tuple, Union + +from httpie.cli.dicts import RequestDataDict, RequestFilesDict +from requests_toolbelt import MultipartEncoder + + +# Multipart uploads smaller than this size gets buffered (otherwise streamed). +# NOTE: Unbuffered upload requests cannot be displayed on the terminal. +UPLOAD_BUFFER = 1024 * 100 + + +def get_multipart_data( + data: RequestDataDict, + files: RequestFilesDict +) -> Tuple[Union[MultipartEncoder, bytes], str]: + fields = list(data.items()) + list(files.items()) + encoder = MultipartEncoder(fields=fields) + content_type = encoder.content_type + data = encoder.to_string() if encoder.len < UPLOAD_BUFFER else encoder + return data, content_type diff --git a/httpie/utils.py b/httpie/utils.py index 41bdbdd4..e33adee1 100644 --- a/httpie/utils.py +++ b/httpie/utils.py @@ -111,3 +111,5 @@ def get_expired_cookies( for cookie in cookies if cookie.get('expires', float('Inf')) <= now ] + + diff --git a/setup.py b/setup.py index 0c738ef7..f2957ab9 100644 --- a/setup.py +++ b/setup.py @@ -38,6 +38,7 @@ tests_require = [ install_requires = [ 'requests[socks]>=2.22.0', 'Pygments>=2.5.2', + 'requests-toolbelt>=0.9.1', ] install_requires_win_only = [ 'colorama>=0.2.4', diff --git a/tests/test_uploads.py b/tests/test_uploads.py index d2bab006..e59892a9 100644 --- a/tests/test_uploads.py +++ b/tests/test_uploads.py @@ -1,8 +1,10 @@ import os +from unittest import mock import pytest from httpie.cli.exceptions import ParseError +from httpie.output.streams import LARGE_UPLOAD_SUPPRESSED_NOTICE from httpie.status import ExitStatus from utils import MockEnvironment, http, HTTP_OK from fixtures import FILE_PATH_ARG, FILE_PATH, FILE_CONTENT @@ -39,15 +41,32 @@ class TestMultipartFormDataFileUpload: assert r.count('Content-Type: text/plain') == 2 def test_upload_custom_content_type(self, httpbin): - r = http('--form', '--verbose', 'POST', httpbin.url + '/post', - f'test-file@{FILE_PATH_ARG};type=image/vnd.microsoft.icon') + r = http( + '--form', + '--verbose', + httpbin.url + '/post', + f'test-file@{FILE_PATH_ARG};type=image/vnd.microsoft.icon' + ) assert HTTP_OK in r # Content type is stripped from the filename assert 'Content-Disposition: form-data; name="test-file";' \ f' filename="{os.path.basename(FILE_PATH)}"' in r - assert FILE_CONTENT in r + assert r.count(FILE_CONTENT) == 2 assert 'Content-Type: image/vnd.microsoft.icon' in r + @mock.patch('httpie.uploads.UPLOAD_BUFFER', 0) + def test_large_upload_display_suppressed(self, httpbin): + r = http( + '--form', + '--verbose', + httpbin.url + '/post', + f'test-file@{FILE_PATH_ARG}', + 'foo=bar', + ) + assert HTTP_OK in r + assert r.count(FILE_CONTENT) == 1 + assert LARGE_UPLOAD_SUPPRESSED_NOTICE.decode() in r + class TestRequestBodyFromFilePath: """