Update save / export methods to return data generators1774-export-methods

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

5 files changed, 53 insertions, 20 deletions

diff --git a/docker/api/container.py b/docker/api/container.py
index b08032c..49230c7 100644
--- a/docker/api/container.py
+++ b/docker/api/container.py
@@ -698,7 +698,7 @@ class ContainerApiMixin(object):
             container (str): The container to export
 
         Returns:
-            (str): The filesystem tar archive
+            (generator): The archived filesystem data stream
 
         Raises:
             :py:class:`docker.errors.APIError`
@@ -707,8 +707,7 @@ class ContainerApiMixin(object):
         res = self._get(
             self._url("/containers/{0}/export", container), stream=True
         )
-        self._raise_for_status(res)
-        return res.raw
+        return self._stream_raw_result(res)
 
     @utils.check_resource('container')
     @utils.minimum_version('1.20')
@@ -737,7 +736,7 @@ class ContainerApiMixin(object):
         self._raise_for_status(res)
         encoded_stat = res.headers.get('x-docker-container-path-stat')
         return (
-            res.raw,
+            self._stream_raw_result(res),
             utils.decode_json_header(encoded_stat) if encoded_stat else None
         )
 
diff --git a/docker/api/image.py b/docker/api/image.py
index 065fae3..b3dcd3a 100644
--- a/docker/api/image.py
+++ b/docker/api/image.py
@@ -21,8 +21,7 @@ class ImageApiMixin(object):
             image (str): Image name to get
 
         Returns:
-            (urllib3.response.HTTPResponse object): The response from the
-            daemon.
+            (generator): A stream of raw archive data.
 
         Raises:
             :py:class:`docker.errors.APIError`
@@ -30,14 +29,14 @@ class ImageApiMixin(object):
 
         Example:
 
-            >>> image = cli.get_image("fedora:latest")
-            >>> f = open('/tmp/fedora-latest.tar', 'w')
-            >>> f.write(image.data)
+            >>> image = cli.get_image("busybox:latest")
+            >>> f = open('/tmp/busybox-latest.tar', 'w')
+            >>> for chunk in image:
+            >>>   f.write(chunk)
             >>> f.close()
         """
         res = self._get(self._url("/images/{0}/get", image), stream=True)
-        self._raise_for_status(res)
-        return res.raw
+        return self._stream_raw_result(res)
 
     @utils.check_resource('image')
     def history(self, image):
diff --git a/docker/models/images.py b/docker/models/images.py
index dcdeac9..8229cfc 100644
--- a/docker/models/images.py
+++ b/docker/models/images.py
@@ -61,8 +61,7 @@ class Image(Model):
         Get a tarball of an image. Similar to the ``docker save`` command.
 
         Returns:
-            (urllib3.response.HTTPResponse object): The response from the
-            daemon.
+            (generator): A stream of raw archive data.
 
         Raises:
             :py:class:`docker.errors.APIError`
@@ -70,11 +69,10 @@ class Image(Model):
 
         Example:
 
-            >>> image = cli.images.get("fedora:latest")
-            >>> resp = image.save()
-            >>> f = open('/tmp/fedora-latest.tar', 'w')
-            >>> for chunk in resp.stream():
-            >>>     f.write(chunk)
+            >>> image = cli.get_image("busybox:latest")
+            >>> f = open('/tmp/busybox-latest.tar', 'w')
+            >>> for chunk in image:
+            >>>   f.write(chunk)
             >>> f.close()
         """
         return self.client.api.get_image(self.id)
diff --git a/tests/integration/api_image_test.py b/tests/integration/api_image_test.py
index 178c34e..ae93190 100644
--- a/tests/integration/api_image_test.py
+++ b/tests/integration/api_image_test.py
@@ -329,7 +329,7 @@ class PruneImagesTest(BaseAPIIntegrationTest):
         img_id = self.client.inspect_image('hello-world')['Id']
         result = self.client.prune_images()
         assert img_id not in [
-            img.get('Deleted') for img in result['ImagesDeleted']
+            img.get('Deleted') for img in result.get('ImagesDeleted') or []
         ]
         result = self.client.prune_images({'dangling': False})
         assert result['SpaceReclaimed'] > 0
@@ -339,3 +339,25 @@ class PruneImagesTest(BaseAPIIntegrationTest):
         assert img_id in [
             img.get('Deleted') for img in result['ImagesDeleted']
         ]
+
+
+class SaveLoadImagesTest(BaseAPIIntegrationTest):
+    @requires_api_version('1.23')
+    def test_get_image_load_image(self):
+        with tempfile.TemporaryFile() as f:
+            stream = self.client.get_image(BUSYBOX)
+            for chunk in stream:
+                f.write(chunk)
+
+            f.seek(0)
+            result = self.client.load_image(f.read())
+
+        success = False
+        result_line = 'Loaded image: {}\n'.format(BUSYBOX)
+        for data in result:
+            print(data)
+            if 'stream' in data:
+                if data['stream'] == result_line:
+                    success = True
+                    break
+        assert success is True
diff --git a/tests/integration/models_images_test.py b/tests/integration/models_images_test.py
index 8840e15..2a28e12 100644
--- a/tests/integration/models_images_test.py
+++ b/tests/integration/models_images_test.py
@@ -1,9 +1,10 @@
 import io
+import tempfile
 
 import docker
 import pytest
 
-from .base import BaseIntegrationTest, TEST_API_VERSION
+from .base import BaseIntegrationTest, BUSYBOX, TEST_API_VERSION
 
 
 class ImageCollectionTest(BaseIntegrationTest):
@@ -76,6 +77,20 @@ class ImageCollectionTest(BaseIntegrationTest):
         with pytest.raises(docker.errors.ImageLoadError):
             client.images.load('abc')
 
+    def test_save_and_load(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        image = client.images.get(BUSYBOX)
+        with tempfile.TemporaryFile() as f:
+            stream = image.save()
+            for chunk in stream:
+                f.write(chunk)
+
+            f.seek(0)
+            result = client.images.load(f.read())
+
+        assert len(result) == 1
+        assert result[0].id == image.id
+
 
 class ImageTest(BaseIntegrationTest):