Merge pull request #85 from box/stream_uploads

Jeff-Meadows · Jeff-Meadows · commit 42a03fa84ab0 · 2015-11-05T11:43:42.000-08:00
Stream uploads from disk.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -9,8 +9,9 @@ Upcoming
 - CPython 3.5 support.
 - Support for cryptography>=1.0 on PyPy 2.6.
 - Travis CI testing for CPython 3.5 and PyPy 2.6.0.
+- Stream uploads of files from disk.
 
-1.2.1 (2015-07-22)
+1.2.2 (2015-07-22)
 ++++++++++++++++++
 
 - The SDK now supports setting a password when creating a shared link.
diff --git a/boxsdk/session/box_session.py b/boxsdk/session/box_session.py
@@ -3,6 +3,7 @@
 from __future__ import unicode_literals
 
 from boxsdk.exception import BoxAPIException
+from boxsdk.util.multipart_stream import MultipartStream
 from boxsdk.util.shared_link import get_shared_link_header
 
 
@@ -323,18 +324,25 @@ def _make_request(
 
         # Reset stream positions to what they were when the request was made so the same data is sent even if this
         # is a retried attempt.
+        request_kwargs = kwargs
         files, file_stream_positions = kwargs.get('files'), kwargs.pop('file_stream_positions')
         if files and file_stream_positions:
+            request_kwargs = kwargs.copy()
             for name, position in file_stream_positions.items():
                 files[name][1].seek(position)
+            data = request_kwargs.pop('data', {})
+            multipart_stream = MultipartStream(data, files)
+            request_kwargs['data'] = multipart_stream
+            del request_kwargs['files']
+            headers['Content-Type'] = multipart_stream.content_type
 
         # send the request
         network_response = self._network_layer.request(
             method,
             url,
             access_token=access_token_will_be_used,
             headers=headers,
-            **kwargs
+            **request_kwargs
         )
 
         network_response = self._retry_request_if_necessary(
diff --git a/boxsdk/util/multipart_stream.py b/boxsdk/util/multipart_stream.py
@@ -0,0 +1,22 @@
+# coding: utf-8
+
+from __future__ import unicode_literals
+
+from requests_toolbelt.multipart.encoder import MultipartEncoder
+
+from boxsdk.util.ordered_dict import OrderedDict
+
+
+class MultipartStream(MultipartEncoder):
+    """
+    Subclass of the requests_toolbelt's :class:`MultipartEncoder` that ensures that data
+    is encoded before files. This allows a server to process information in the data before
+    receiving the file bytes.
+    """
+    def __init__(self, data, files):
+        fields = OrderedDict()
+        for k in data:
+            fields[k] = data[k]
+        for k in files:
+            fields[k] = files[k]
+        super(MultipartStream, self).__init__(fields)
diff --git a/docs/source/boxsdk.util.rst b/docs/source/boxsdk.util.rst
@@ -20,6 +20,14 @@ boxsdk.util.lru_cache module
     :undoc-members:
     :show-inheritance:
 
+boxsdk.util.multipart_stream module
+-----------------------------------
+
+.. automodule:: boxsdk.util.multipart_stream
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 boxsdk.util.ordered_dict module
 -------------------------------
 
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,6 @@
 cryptography>=0.9.2
 pyjwt>=1.3.0
 requests>=2.4.3
+requests-toolbelt>=0.4.0
 six >= 1.4.0
 .
diff --git a/test/conftest.py b/test/conftest.py
@@ -128,8 +128,8 @@ def auth_code():
 
 @pytest.fixture(params=[
     b'Hello',
-    'Goodbye',
-    '42',
+    b'Goodbye',
+    b'42',
 ])
 def test_file_content(request):
     return request.param
diff --git a/test/functional/conftest.py b/test/functional/conftest.py
@@ -1,7 +1,7 @@
 # coding: utf-8
 
 from __future__ import unicode_literals
-from mock import mock_open, patch
+from mock import patch
 import pytest
 import re
 import requests
@@ -11,6 +11,7 @@
 from boxsdk.config import API
 from boxsdk.client import Client
 from test.functional.mock_box.box import Box
+from test.util.streamable_mock_open import streamable_mock_open
 
 
 @pytest.fixture()
@@ -95,7 +96,7 @@ def user_login():
 @pytest.fixture()
 def uploaded_file(box_client, test_file_path, test_file_content, file_name):
     # pylint:disable=redefined-outer-name
-    with patch('boxsdk.object.folder.open', mock_open(read_data=test_file_content), create=True):
+    with patch('boxsdk.object.folder.open', streamable_mock_open(read_data=test_file_content), create=True):
         return box_client.folder('0').upload(test_file_path, file_name)
 
 
diff --git a/test/functional/test_delete.py b/test/functional/test_delete.py
@@ -1,14 +1,15 @@
 # coding: utf-8
 
 from __future__ import unicode_literals
-from mock import patch, mock_open
+from mock import patch
 import pytest
 from boxsdk.client import Client
 from boxsdk.exception import BoxAPIException
+from test.util.streamable_mock_open import streamable_mock_open
 
 
 def test_upload_then_delete(box_client, test_file_path, test_file_content, file_name):
-    with patch('boxsdk.object.folder.open', mock_open(read_data=test_file_content), create=True):
+    with patch('boxsdk.object.folder.open', streamable_mock_open(read_data=test_file_content), create=True):
         file_object = box_client.folder('0').upload(test_file_path, file_name)
     assert file_object.delete()
     assert len(box_client.folder('0').get_items(1)) == 0
diff --git a/test/functional/test_file_upload_update_download.py b/test/functional/test_file_upload_update_download.py
@@ -1,12 +1,13 @@
 # coding: utf-8
 
 from __future__ import unicode_literals
-from mock import mock_open, patch
+from mock import patch
 import six
+from test.util.streamable_mock_open import streamable_mock_open
 
 
 def test_upload_then_update(box_client, test_file_path, test_file_content, update_file_content, file_name):
-    with patch('boxsdk.object.folder.open', mock_open(read_data=test_file_content), create=True):
+    with patch('boxsdk.object.folder.open', streamable_mock_open(read_data=test_file_content), create=True):
         file_object = box_client.folder('0').upload(test_file_path, file_name)
     assert file_object.name == file_name
     file_object_with_info = file_object.get()
@@ -20,13 +21,15 @@ def test_upload_then_update(box_client, test_file_path, test_file_content, updat
     assert len(folder_items) == 1
     assert folder_items[0].object_id == file_object.object_id
     assert folder_items[0].name == file_object.name
-    with patch('boxsdk.object.file.open', mock_open(read_data=update_file_content), create=True):
+    with patch('boxsdk.object.file.open', streamable_mock_open(read_data=update_file_content), create=True):
         updated_file_object = file_object.update_contents(test_file_path)
     assert updated_file_object.name == file_name
     file_object_with_info = updated_file_object.get()
     assert file_object_with_info.id == updated_file_object.object_id
     assert file_object_with_info.name == file_name
     file_content = updated_file_object.content()
+    expected_file_content = update_file_content.encode('utf-8') if isinstance(update_file_content, six.text_type)\
+        else update_file_content
     assert file_content == expected_file_content
     folder_items = box_client.folder('0').get_items(100)
     assert len(folder_items) == 1
@@ -35,7 +38,7 @@ def test_upload_then_update(box_client, test_file_path, test_file_content, updat
 
 
 def test_upload_then_download(box_client, test_file_path, test_file_content, file_name):
-    with patch('boxsdk.object.folder.open', mock_open(read_data=test_file_content), create=True):
+    with patch('boxsdk.object.folder.open', streamable_mock_open(read_data=test_file_content), create=True):
         file_object = box_client.folder('0').upload(test_file_path, file_name)
     writeable_stream = six.BytesIO()
     file_object.download_to(writeable_stream)
diff --git a/test/unit/session/test_box_session.py b/test/unit/session/test_box_session.py
@@ -119,8 +119,8 @@ def test_box_session_seeks_file_after_retry(box_session, server_error_response,
     assert box_response.status_code == 200
     assert box_response.json() == generic_successful_response.json()
     assert box_response.ok == generic_successful_response.ok
-    mock_file_1.tell.assert_called_once_with()
-    mock_file_2.tell.assert_called_once_with()
+    mock_file_1.tell.assert_called_with()
+    mock_file_2.tell.assert_called_with()
     mock_file_1.seek.assert_called_with(0)
     assert mock_file_1.seek.call_count == 2
     assert mock_file_1.seek.has_calls(call(0) * 2)
diff --git a/test/unit/util/test_multipart_stream.py b/test/unit/util/test_multipart_stream.py
@@ -0,0 +1,30 @@
+# coding: utf-8
+
+from __future__ import unicode_literals, absolute_import
+
+import pytest
+
+from boxsdk.util.multipart_stream import MultipartStream
+
+
+@pytest.fixture(params=({}, {'data_1': b'data_1_value', 'data_2': b'data_2_value'}))
+def multipart_stream_data(request):
+    return request.param
+
+
+@pytest.fixture(params=({}, {'file_1': b'file_1_value', 'file_2': b'file_2_value'}))
+def multipart_stream_files(request):
+    return request.param
+
+
+def test_multipart_stream_orders_data_before_files(multipart_stream_data, multipart_stream_files):
+    # pylint:disable=redefined-outer-name
+    if not multipart_stream_data and not multipart_stream_files:
+        pytest.xfail('Encoder does not support empty fields.')
+    stream = MultipartStream(multipart_stream_data, multipart_stream_files)
+    encoded_stream = stream.to_string()
+    data_indices = [encoded_stream.find(value) for value in multipart_stream_data.values()]
+    file_indices = [encoded_stream.find(value) for value in multipart_stream_files.values()]
+    assert -1 not in data_indices
+    assert -1 not in file_indices
+    assert all((all((data_index < f for f in file_indices)) for data_index in data_indices))
diff --git a/test/util/__init__.py b/test/util/__init__.py
@@ -0,0 +1,3 @@
+# coding: utf-8
+
+from __future__ import unicode_literals, absolute_import
diff --git a/test/util/streamable_mock_open.py b/test/util/streamable_mock_open.py
@@ -0,0 +1,30 @@
+# coding: utf-8
+
+from __future__ import unicode_literals, absolute_import
+
+
+from mock import mock_open
+
+
+def streamable_mock_open(mock=None, read_data=b''):
+    mock = mock_open(mock, read_data)
+    handle = mock.return_value
+    handle.position = 0
+
+    def tell():
+        return handle.position
+
+    def read(size=-1):
+        if size == -1:
+            handle.position = len(read_data)
+            return read_data
+        else:
+            data = read_data[handle.position:handle.position + size]
+            handle.position += size
+            return data
+
+    handle.tell.side_effect = tell
+    handle.len = len(read_data)
+    handle.read.side_effect = read
+    del handle.getvalue
+    return mock

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,6 @@`
`1`	`1`	`cryptography>=0.9.2`
`2`	`2`	`pyjwt>=1.3.0`
`3`	`3`	`requests>=2.4.3`
	`4`	`+requests-toolbelt>=0.4.0`
`4`	`5`	`six >= 1.4.0`
`5`	`6`	`.`