pull, push, build are streamable

docker/client.py

@@ -30,7 +30,7 @@ if not six.PY3:
 
 class APIError(requests.exceptions.HTTPError):
     def __init__(self, message, response, explanation=None):
-        super(APIError, self).__init__(message, response=response)
+        super(APIError, self).__init__(message, response)
 
         self.explanation = explanation
 
@@ -81,7 +81,18 @@ class Client(requests.Session):
         try:
             response.raise_for_status()
         except requests.exceptions.HTTPError as e:
-            raise APIError(e, response=response, explanation=explanation)
+            raise APIError(e, response, explanation=explanation)
+
+    def _stream_result(self, response):
+        self._raise_for_status(response)
+        for line in response.iter_lines(chunk_size=1):
+            # filter out keep-alive new lines
+            if line:
+                yield line + '\n'
+
+    def _stream_result_socket(self, response):
+        self._raise_for_status(response)
+        return response.raw._fp.fp._sock
 
     def _result(self, response, json=False):
         self._raise_for_status(response)
@@ -192,7 +203,7 @@ class Client(requests.Session):
         }
 
     def build(self, path=None, tag=None, quiet=False, fileobj=None,
-              nocache=False, rm=False):
+              nocache=False, rm=False, stream=False):
         remote = context = headers = None
         if path is None and fileobj is None:
             raise Exception("Either path or fileobj needs to be provided.")
@@ -214,15 +225,19 @@ class Client(requests.Session):
         }
         if context is not None:
             headers = {'Content-Type': 'application/tar'}
-        res = self._result(self.post(u, context, params=params,
+        response = self.post(
-                           headers=headers, stream=True))
+            u, context, params=params, headers=headers, stream=stream)
         if context is not None:
             context.close()
-        srch = r'Successfully built ([0-9a-f]+)'
+        if stream:
-        match = re.search(srch, res)
+            return self._stream_result(response)
-        if not match:
+        else:
-            return None, res
+            output = self._result(response)
-        return match.group(1), res
+            srch = r'Successfully built ([0-9a-f]+)'
+            match = re.search(srch, output)
+            if not match:
+                return None, output
+            return match.group(1), output
 
     def commit(self, container, repository=None, tag=None, message=None,
                author=None, conf=None):
@@ -410,7 +425,7 @@ class Client(requests.Session):
 
         return f_port
 
-    def pull(self, repository, tag=None):
+    def pull(self, repository, tag=None, stream=False):
         registry, repo_name = auth.resolve_repository_name(repository)
         if repo_name.count(":") == 1:
             repository, tag = repository.rsplit(":", 1)
@@ -430,9 +445,30 @@ class Client(requests.Session):
             if authcfg:
                 headers['X-Registry-Auth'] = auth.encode_header(authcfg)
         u = self._url("/images/create")
-        return self._result(self.post(u, params=params, headers=headers))
+        response = self.post(u, params=params, headers=headers, stream=stream)
 
-    def push(self, repository):
+        if stream:
+            return self.stream_helper(response)
+        else:
+            return self._result(response)
+
+    def stream_helper(self, response):
+        socket = self._stream_result_socket(response)
+        while True:
+            chunk = socket.recv(4096)
+            if chunk:
+                parts = chunk.strip().split('\r\n')
+                for i in range(len(parts)):
+                    if i % 2 != 0:
+                        yield parts[i] + '\n'
+                    else:
+                        size = int(parts[i], 16)
+                if size <= 0:
+                    break
+            else:
+                break
+
+    def push(self, repository, stream=False):
         registry, repo_name = auth.resolve_repository_name(repository)
         u = self._url("/images/{0}/push".format(repository))
         headers = {}
@@ -444,8 +480,17 @@ class Client(requests.Session):
             # for this specific registry as we can have an anon push
             if authcfg:
                 headers['X-Registry-Auth'] = auth.encode_header(authcfg)
-            return self._result(self._post_json(u, None, headers=headers))
+            if stream:
-        return self._result(self._post_json(u, authcfg))
+                return self.stream_helper(
+                    self._post_json(u, None, headers=headers, stream=True))
+            else:
+                return self._result(
+                    self._post_json(u, None, headers=headers, stream=False))
+        if stream:
+            return self.stream_helper(
+                self._post_json(u, authcfg, stream=True))
+        else:
+            return self._result(self._post_json(u, authcfg, stream=False))
 
     def remove_container(self, container, v=False):
         if isinstance(container, dict):

tests/integration_test.py

@@ -459,6 +459,28 @@ class TestPull(BaseTestCase):
         self.tmp_imgs.append('376968a23351')
 
 
+class TestPullStream(BaseTestCase):
+    def runTest(self):
+        try:
+            self.client.remove_image('joffrey/test001')
+            self.client.remove_image('376968a23351')
+        except docker.APIError:
+            pass
+        info = self.client.info()
+        self.assertIn('Images', info)
+        img_count = info['Images']
+        stream = self.client.pull('joffrey/test001', stream=True)
+        res = u''
+        for chunk in stream:
+            res += chunk
+        self.assertEqual(type(res), six.text_type)
+        self.assertEqual(img_count + 2, self.client.info()['Images'])
+        img_info = self.client.inspect_image('joffrey/test001')
+        self.assertIn('id', img_info)
+        self.tmp_imgs.append('joffrey/test001')
+        self.tmp_imgs.append('376968a23351')
+
+
 class TestCommit(BaseTestCase):
     def runTest(self):
         container = self.client.create_container('busybox', ['touch', '/test'])
@@ -529,6 +551,23 @@ class TestBuild(BaseTestCase):
         self.tmp_imgs.append(img)
 
 
+class TestBuildStream(BaseTestCase):
+    def runTest(self):
+        script = io.BytesIO('\n'.join([
+            'FROM busybox',
+            'MAINTAINER docker-py',
+            'RUN mkdir -p /tmp/test',
+            'EXPOSE 8080',
+            'ADD https://dl.dropboxusercontent.com/u/20637798/silence.tar.gz'
+            ' /tmp/silence.tar.gz'
+        ]).encode('ascii'))
+        stream = self.client.build(fileobj=script, stream=True)
+        logs = ''
+        for chunk in stream:
+            logs += chunk
+        self.assertNotEqual(logs, '')
+
+
 class TestBuildFromStringIO(BaseTestCase):
     def runTest(self):
         if six.PY3:

tests/test.py

@@ -467,7 +467,21 @@ class DockerClientTest(unittest.TestCase):
         fake_request.assert_called_with(
             'unix://var/run/docker.sock/v1.4/images/create',
             headers={},
-            params={'tag': None, 'fromImage': 'joffrey/test001'}
+            params={'tag': None, 'fromImage': 'joffrey/test001'},
+            stream=False
+        )
+
+    def test_pull_stream(self):
+        try:
+            self.client.pull('joffrey/test001', stream=True)
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
+        fake_request.assert_called_with(
+            'unix://var/run/docker.sock/v1.4/images/create',
+            headers={},
+            params={'tag': None, 'fromImage': 'joffrey/test001'},
+            stream=True
         )
 
     def test_commit(self):
@@ -517,6 +531,20 @@ class DockerClientTest(unittest.TestCase):
         except Exception as e:
             self.fail('Command should not raise exception: {0}'.format(e))
 
+    def test_build_container_stream(self):
+        script = io.BytesIO('\n'.join([
+            'FROM busybox',
+            'MAINTAINER docker-py',
+            'RUN mkdir -p /tmp/test',
+            'EXPOSE 8080',
+            'ADD https://dl.dropboxusercontent.com/u/20637798/silence.tar.gz'
+            ' /tmp/silence.tar.gz'
+        ]).encode('ascii'))
+        try:
+            self.client.build(fileobj=script, stream=True)
+        except Exception as e:
+            self.fail('Command should not raise exception: {0}'.format(e))
+
     #######################
     ## PY SPECIFIC TESTS ##
     #######################