Add chunk_size parameter to data downloading methods (export, get_archive, save)1352-data_stream_control

Signed-off-by: Joffrey F <joffrey@docker.com>

3 files changed, 20 insertions, 9 deletions

diff --git a/docker/api/client.py b/docker/api/client.py
index e69d143..bddab61 100644
--- a/docker/api/client.py
+++ b/docker/api/client.py
@@ -350,10 +350,10 @@ class APIClient(
                 break
             yield data
 
-    def _stream_raw_result(self, response):
-        ''' Stream result for TTY-enabled container '''
+    def _stream_raw_result(self, response, chunk_size=1, decode=True):
+        ''' Stream result for TTY-enabled container and raw binary data'''
         self._raise_for_status(response)
-        for out in response.iter_content(chunk_size=1, decode_unicode=True):
+        for out in response.iter_content(chunk_size, decode):
             yield out
 
     def _read_from_socket(self, response, stream, tty=False):
diff --git a/docker/api/container.py b/docker/api/container.py
index 962d8cb..e986cf2 100644
--- a/docker/api/container.py
+++ b/docker/api/container.py
@@ -3,6 +3,7 @@ from datetime import datetime
 
 from .. import errors
 from .. import utils
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 from ..types import (
     ContainerConfig, EndpointConfig, HostConfig, NetworkingConfig
 )
@@ -643,12 +644,15 @@ class ContainerApiMixin(object):
         )
 
     @utils.check_resource('container')
-    def export(self, container):
+    def export(self, container, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Export the contents of a filesystem as a tar archive.
 
         Args:
             container (str): The container to export
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (generator): The archived filesystem data stream
@@ -660,10 +664,10 @@ class ContainerApiMixin(object):
         res = self._get(
             self._url("/containers/{0}/export", container), stream=True
         )
-        return self._stream_raw_result(res)
+        return self._stream_raw_result(res, chunk_size, False)
 
     @utils.check_resource('container')
-    def get_archive(self, container, path):
+    def get_archive(self, container, path, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Retrieve a file or folder from a container in the form of a tar
         archive.
@@ -671,6 +675,9 @@ class ContainerApiMixin(object):
         Args:
             container (str): The container where the file is located
             path (str): Path to the file or folder to retrieve
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (tuple): First element is a raw tar data stream. Second element is
@@ -688,7 +695,7 @@ class ContainerApiMixin(object):
         self._raise_for_status(res)
         encoded_stat = res.headers.get('x-docker-container-path-stat')
         return (
-            self._stream_raw_result(res),
+            self._stream_raw_result(res, chunk_size, False),
             utils.decode_json_header(encoded_stat) if encoded_stat else None
         )
 
diff --git a/docker/api/image.py b/docker/api/image.py
index fa832a3..3ebca32 100644
--- a/docker/api/image.py
+++ b/docker/api/image.py
@@ -4,6 +4,7 @@ import os
 import six
 
 from .. import auth, errors, utils
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 
 log = logging.getLogger(__name__)
 
@@ -11,12 +12,15 @@ log = logging.getLogger(__name__)
 class ImageApiMixin(object):
 
     @utils.check_resource('image')
-    def get_image(self, image):
+    def get_image(self, image, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Get a tarball of an image. Similar to the ``docker save`` command.
 
         Args:
             image (str): Image name to get
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (generator): A stream of raw archive data.
@@ -34,7 +38,7 @@ class ImageApiMixin(object):
             >>> f.close()
         """
         res = self._get(self._url("/images/{0}/get", image), stream=True)
-        return self._stream_raw_result(res)
+        return self._stream_raw_result(res, chunk_size, False)
 
     @utils.check_resource('image')
     def history(self, image):