Include tag in images.get after pulling if provided separately

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

docker/models/images.py

@@ -238,7 +238,7 @@ class ImageCollection(Collection):
         """
         return self.client.api.load_image(data)
 
-    def pull(self, name, **kwargs):
+    def pull(self, name, tag=None, **kwargs):
         """
         Pull an image of the given name and return it. Similar to the
         ``docker pull`` command.
@@ -267,8 +267,8 @@ class ImageCollection(Collection):
 
             >>> image = client.images.pull('busybox')
         """
-        self.client.api.pull(name, **kwargs)
-        return self.get(name)
+        self.client.api.pull(name, tag=tag, **kwargs)
+        return self.get('{0}:{1}'.format(name, tag) if tag else name)
 
     def push(self, repository, tag=None, **kwargs):
         return self.client.api.push(repository, tag=tag, **kwargs)

tests/integration/models_images_test.py

@@ -30,10 +30,12 @@ class ImageCollectionTest(BaseIntegrationTest):
 
     def test_build_with_multiple_success(self):
         client = docker.from_env(version=TEST_API_VERSION)
-        image = client.images.build(tag='some-tag', fileobj=io.BytesIO(
+        image = client.images.build(
+            tag='some-tag', fileobj=io.BytesIO(
                 "FROM alpine\n"
                 "CMD echo hello world".encode('ascii')
-            ))
+            )
+        )
         self.tmp_imgs.append(image.id)
         assert client.containers.run(image) == b"hello world\n"
 
@@ -53,6 +55,11 @@ class ImageCollectionTest(BaseIntegrationTest):
         image = client.images.pull('alpine:latest')
         assert 'alpine:latest' in image.attrs['RepoTags']
 
+    def test_pull_with_tag(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+        image = client.images.pull('alpine', tag='3.3')
+        assert 'alpine:3.3' in image.attrs['RepoTags']
+
 
 class ImageTest(BaseIntegrationTest):
 

tests/unit/models_containers_test.py

@@ -227,7 +227,7 @@ class ContainerCollectionTest(unittest.TestCase):
         container = client.containers.run('alpine', 'sleep 300', detach=True)
 
         assert container.id == FAKE_CONTAINER_ID
-        client.api.pull.assert_called_with('alpine')
+        client.api.pull.assert_called_with('alpine', tag=None)
 
     def test_run_with_error(self):
         client = make_fake_client()

tests/unit/models_images_test.py

@@ -42,7 +42,7 @@ class ImageCollectionTest(unittest.TestCase):
     def test_pull(self):
         client = make_fake_client()
         image = client.images.pull('test_image')
-        client.api.pull.assert_called_with('test_image')
+        client.api.pull.assert_called_with('test_image', tag=None)
         client.api.inspect_image.assert_called_with('test_image')
         assert isinstance(image, Image)
         assert image.id == FAKE_IMAGE_ID