From 90ddac7b8b27690dbfefe2f2e01c009b59b1de42 Mon Sep 17 00:00:00 2001
From: "STEFANINI Axel (INTERN)" <42176370+axel7083@users.noreply.github.com>
Date: Tue, 4 Jul 2023 15:57:32 +0200
Subject: [PATCH] Adding build progress callback to client.images.build

---
 docker/models/images.py                 | 17 ++++++++++++++++-
 tests/integration/models_images_test.py | 18 ++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/docker/models/images.py b/docker/models/images.py
index e3ec39d28..283db7cfd 100644
--- a/docker/models/images.py
+++ b/docker/models/images.py
@@ -219,7 +219,7 @@ def reload(self):
 class ImageCollection(Collection):
     model = Image
 
-    def build(self, **kwargs):
+    def build(self, progress_callback=None, **kwargs):
         """
         Build an image and return it. Similar to the ``docker build``
         command. Either ``path`` or ``fileobj`` must be set.
@@ -281,6 +281,9 @@ def build(self, **kwargs):
                 configuration file (``~/.docker/config.json`` by default)
                 contains a proxy configuration, the corresponding environment
                 variables will be set in the container being built.
+            progress_callback Callable[[int, int, str], None]: Will get called
+                with the following arguments: current_step, max_step, details
+                Default: `None`.
 
         Returns:
             (tuple): The first item is the :py:class:`Image` object for the
@@ -305,12 +308,24 @@ def build(self, **kwargs):
             if 'error' in chunk:
                 raise BuildError(chunk['error'], result_stream)
             if 'stream' in chunk:
+
                 match = re.search(
                     r'(^Successfully built |sha256:)([0-9a-f]+)$',
                     chunk['stream']
                 )
                 if match:
                     image_id = match.group(2)
+                else:
+                    match = re.match(
+                        r'Step (\d+)/(\d+) : (.+)',
+                        chunk['stream']
+                    )
+                    if match and callable(progress_callback):
+                        progress_callback(
+                            int(match.group(1)),
+                            int(match.group(2)),
+                            match.group(3)
+                        )
             last_event = chunk
         if image_id:
             return (self.get(image_id), result_stream)
diff --git a/tests/integration/models_images_test.py b/tests/integration/models_images_test.py
index 94aa20100..c1ceec6c5 100644
--- a/tests/integration/models_images_test.py
+++ b/tests/integration/models_images_test.py
@@ -19,6 +19,24 @@ def test_build(self):
         self.tmp_imgs.append(image.id)
         assert client.containers.run(image) == b"hello world\n"
 
+    def test_build_with_progress_callback(self):
+        client = docker.from_env(version=TEST_API_VERSION)
+
+        def progress_callback(current_step, max_step, details):
+            assert max_step == 3
+            if current_step == 1:
+                assert details == "FROM alpine"
+            elif current_step == 2:
+                assert details == "RUN echo 1"
+            elif current_step == 3:
+                assert details == "RUN echo 2"
+
+        image, _ = client.images.build(fileobj=io.BytesIO(
+            b"FROM alpine\n"
+            b"RUN echo 1\n"
+            b"RUN echo 2"
+        ), progress_callback=progress_callback)
+
     # @pytest.mark.xfail(reason='Engine 1.13 responds with status 500')
     def test_build_with_error(self):
         client = docker.from_env(version=TEST_API_VERSION)