Rewrite Client.import_image() and add less-magical versions

Rather than have one function that can do 5 things and attempts to guess
what the user wanted, it's good to have a function for each possible
method of importing an image. Error cases are handled a lot better this
way.

The unit test test_import_image_from_file() was changed to be
test_import_image_from_bytes(), because if we try to import a temporary
file in the test, the 'data' parameter of the 'fake_request' object is
an instance of a funny internal type like <_io.BufferedReader
name='/tmp/tmpc9chux'> and it's very hard to match such a thing with
`mock.assert_called_with()`.

2 files changed, 80 insertions, 27 deletions

diff --git a/docker/client.py b/docker/client.py
index d6c96a2..02a731c 100644
--- a/docker/client.py
+++ b/docker/client.py
@@ -646,33 +646,86 @@ class Client(requests.Session):
         return res
 
     def import_image(self, src=None, repository=None, tag=None, image=None):
+        if src:
+            if isinstance(src, six.string_types):
+                try:
+                    result = self.import_image_from_file(
+                        src, repository=repository, tag=tag)
+                except IOError:
+                    result = self.import_image_from_url(
+                        src, repository=repository, tag=tag)
+            else:
+                result = self.import_image_from_data(
+                    src, repository=repository, tag=tag)
+        elif image:
+            result = self.import_image_from_image(
+                image, repository=repository, tag=tag)
+        else:
+            raise Exception("Must specify a src or image")
+
+        return result
+
+    def import_image_from_data(self, data, repository=None, tag=None):
         u = self._url("/images/create")
         params = {
+            'fromSrc': '-',
             'repo': repository,
             'tag': tag
         }
+        headers = {
+            'Content-Type': 'application/tar',
+        }
+        return self._result(
+            self._post(u, data=data, params=params, headers=headers))
 
-        if src:
-            try:
-                # XXX: this is ways not optimal but the only way
-                # for now to import tarballs through the API
-                fic = open(src)
-                data = fic.read()
-                fic.close()
-                src = "-"
-            except IOError:
-                # file does not exists or not a file (URL)
-                data = None
-            if isinstance(src, six.string_types):
-                params['fromSrc'] = src
-                return self._result(self._post(u, data=data, params=params))
-            return self._result(self._post(u, data=src, params=params))
+    def import_image_from_file(self, filename, repository=None, tag=None):
+        u = self._url("/images/create")
+        params = {
+            'fromSrc': '-',
+            'repo': repository,
+            'tag': tag
+        }
+        headers = {
+            'Content-Type': 'application/tar',
+        }
+        with open(filename, 'rb') as f:
+            return self._result(
+                self._post(u, data=f, params=params, headers=headers,
+                           timeout=None))
+
+    def import_image_from_stream(self, stream, repository=None, tag=None):
+        u = self._url("/images/create")
+        params = {
+            'fromSrc': '-',
+            'repo': repository,
+            'tag': tag
+        }
+        headers = {
+            'Content-Type': 'application/tar',
+            'Transfer-Encoding': 'chunked',
+        }
+        return self._result(
+            self._post(u, data=stream, params=params, headers=headers))
 
-        if image:
-            params['fromImage'] = image
-            return self._result(self._post(u, data=None, params=params))
+    def import_image_from_url(self, url, repository=None, tag=None):
+        u = self._url("/images/create")
+        params = {
+            'fromSrc': url,
+            'repo': repository,
+            'tag': tag
+        }
+        return self._result(
+            self._post(u, data=None, params=params))
 
-        raise Exception("Must specify a src or image")
+    def import_image_from_image(self, image, repository=None, tag=None):
+        u = self._url("/images/create")
+        params = {
+            'fromImage': image,
+            'repo': repository,
+            'tag': tag
+        }
+        return self._result(
+            self._post(u, data=None, params=params))
 
     def info(self):
         return self._result(self._get(self._url("/info")),
diff --git a/tests/test.py b/tests/test.py
index 4a17a8a..0437889 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -1700,12 +1700,11 @@ class DockerClientTest(Cleanup, unittest.TestCase):
             timeout=docker.client.DEFAULT_TIMEOUT_SECONDS
         )
 
-    def test_import_image_from_file(self):
-        buf = tempfile.NamedTemporaryFile(delete=False)
+    def test_import_image_from_bytes(self):
+        stream = (i for i in range(0, 100))
         try:
-            # pretent the buffer is a file
             self.client.import_image(
-                buf.name,
+                stream,
                 repository=fake_api.FAKE_REPO_NAME,
                 tag=fake_api.FAKE_TAG_NAME
             )
@@ -1717,13 +1716,14 @@ class DockerClientTest(Cleanup, unittest.TestCase):
             params={
                 'repo': fake_api.FAKE_REPO_NAME,
                 'tag': fake_api.FAKE_TAG_NAME,
-                'fromSrc': '-'
+                'fromSrc': '-',
             },
-            data='',
+            headers={
+                'Content-Type': 'application/tar',
+            },
+            data=stream,
             timeout=docker.client.DEFAULT_TIMEOUT_SECONDS
         )
-        buf.close()
-        os.remove(buf.name)
 
     def test_import_image_from_image(self):
         try: