Merge pull request #2186 from docker/adw1n-i2116

 Fix pulling images with `stream=True`

Author: shin-

docker/api/image.py

@@ -334,7 +334,8 @@ def pull(self, repository, tag=None, stream=False, auth_config=None,
         Args:
             repository (str): The repository to pull
             tag (str): The tag to pull
-            stream (bool): Stream the output as a generator
+            stream (bool): Stream the output as a generator. Make sure to
+                consume the generator, otherwise pull might get cancelled.
             auth_config (dict): Override the credentials that
                 :py:meth:`~docker.api.daemon.DaemonApiMixin.login` has set for
                 this request. ``auth_config`` should contain the ``username``

docker/models/images.py

@@ -1,5 +1,6 @@
 import itertools
 import re
+import warnings
 
 import six
 
@@ -425,7 +426,21 @@ def pull(self, repository, tag=None, **kwargs):
         if not tag:
             repository, tag = parse_repository_tag(repository)
 
-        self.client.api.pull(repository, tag=tag, **kwargs)
+        if 'stream' in kwargs:
+            warnings.warn(
+                '`stream` is not a valid parameter for this method'
+                ' and will be overridden'
+            )
+            del kwargs['stream']
+
+        pull_log = self.client.api.pull(
+            repository, tag=tag, stream=True, **kwargs
+        )
+        for _ in pull_log:
+            # We don't do anything with the logs, but we need
+            # to keep the connection alive and wait for the image
+            # to be pulled.
+            pass
         if tag:
             return self.get('{0}{2}{1}'.format(
                 repository, tag, '@' if tag.startswith('sha256:') else ':'

tests/unit/models_containers_test.py

@@ -232,7 +232,9 @@ def test_run_pull(self):
         container = client.containers.run('alpine', 'sleep 300', detach=True)
 
         assert container.id == FAKE_CONTAINER_ID
-        client.api.pull.assert_called_with('alpine', platform=None, tag=None)
+        client.api.pull.assert_called_with(
+            'alpine', platform=None, tag=None, stream=True
+        )
 
     def test_run_with_error(self):
         client = make_fake_client()

tests/unit/models_images_test.py

@@ -1,6 +1,8 @@
+import unittest
+import warnings
+
 from docker.constants import DEFAULT_DATA_CHUNK_SIZE
 from docker.models.images import Image
-import unittest
 
 from .fake_api import FAKE_IMAGE_ID
 from .fake_api_client import make_fake_client
@@ -43,15 +45,19 @@ def test_load(self):
     def test_pull(self):
         client = make_fake_client()
         image = client.images.pull('test_image:latest')
-        client.api.pull.assert_called_with('test_image', tag='latest')
+        client.api.pull.assert_called_with(
+            'test_image', tag='latest', stream=True
+        )
         client.api.inspect_image.assert_called_with('test_image:latest')
         assert isinstance(image, Image)
         assert image.id == FAKE_IMAGE_ID
 
     def test_pull_multiple(self):
         client = make_fake_client()
         images = client.images.pull('test_image')
-        client.api.pull.assert_called_with('test_image', tag=None)
+        client.api.pull.assert_called_with(
+            'test_image', tag=None, stream=True
+        )
         client.api.images.assert_called_with(
             all=False, name='test_image', filters=None
         )
@@ -61,6 +67,16 @@ def test_pull_multiple(self):
         assert isinstance(image, Image)
         assert image.id == FAKE_IMAGE_ID
 
+    def test_pull_with_stream_param(self):
+        client = make_fake_client()
+        with warnings.catch_warnings(record=True) as w:
+            client.images.pull('test_image', stream=True)
+
+        assert len(w) == 1
+        assert str(w[0].message).startswith(
+            '`stream` is not a valid parameter'
+        )
+
     def test_push(self):
         client = make_fake_client()
         client.images.push('foobar', insecure_registry=True)