Add named parameter to image.save to identify which repository name to use in the resulting tarball

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

2 files changed, 45 insertions, 2 deletions

diff --git a/docker/models/images.py b/docker/models/images.py
index 7d9ab70..28b1fd3 100644
--- a/docker/models/images.py
+++ b/docker/models/images.py
@@ -59,7 +59,7 @@ class Image(Model):
         """
         return self.client.api.history(self.id)
 
-    def save(self, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
+    def save(self, chunk_size=DEFAULT_DATA_CHUNK_SIZE, named=False):
         """
         Get a tarball of an image. Similar to the ``docker save`` command.
 
@@ -67,6 +67,12 @@ class Image(Model):
             chunk_size (int): The generator will return up to that much data
                 per iteration, but may return less. If ``None``, data will be
                 streamed as it is received. Default: 2 MB
+            named (str or bool): If ``False`` (default), the tarball will not
+                retain repository and tag information for this image. If set
+                to ``True``, the first tag in the :py:attr:`~tags` list will
+                be used to identify the image. Alternatively, any element of
+                the :py:attr:`~tags` list can be used as an argument to use
+                that specific tag as the saved identifier.
 
         Returns:
             (generator): A stream of raw archive data.
@@ -83,7 +89,17 @@ class Image(Model):
             >>>   f.write(chunk)
             >>> f.close()
         """
-        return self.client.api.get_image(self.id, chunk_size)
+        img = self.id
+        if named:
+            img = self.tags[0] if self.tags else img
+            if isinstance(named, six.string_types):
+                if named not in self.tags:
+                    raise InvalidArgument(
+                        "{} is not a valid tag for this image".format(named)
+                    )
+                img = named
+
+        return self.client.api.get_image(img, chunk_size)
 
     def tag(self, repository, tag=None, **kwargs):
         """
diff --git a/tests/integration/models_images_test.py b/tests/integration/models_images_test.py
index ae735ba..31fab10 100644
--- a/tests/integration/models_images_test.py
+++ b/tests/integration/models_images_test.py
@@ -5,6 +5,7 @@ import docker
 import pytest
 
 from .base import BaseIntegrationTest, BUSYBOX, TEST_API_VERSION
+from ..helpers import random_name
 
 
 class ImageCollectionTest(BaseIntegrationTest):
@@ -108,6 +109,32 @@ class ImageCollectionTest(BaseIntegrationTest):
         assert len(result) == 1
         assert result[0].id == image.id
 
+    def test_save_and_load_repo_name(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        image = client.images.get(BUSYBOX)
+        additional_tag = random_name()
+        image.tag(additional_tag)
+        self.tmp_imgs.append(additional_tag)
+        image.reload()
+        with tempfile.TemporaryFile() as f:
+            stream = image.save(named='{}:latest'.format(additional_tag))
+            for chunk in stream:
+                f.write(chunk)
+
+            f.seek(0)
+            client.images.remove(additional_tag, force=True)
+            result = client.images.load(f.read())
+
+        assert len(result) == 1
+        assert result[0].id == image.id
+        assert '{}:latest'.format(additional_tag) in result[0].tags
+
+    def test_save_name_error(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        image = client.images.get(BUSYBOX)
+        with pytest.raises(docker.errors.InvalidArgument):
+            image.save(named='sakuya/izayoi')
+
 
 class ImageTest(BaseIntegrationTest):