From f0739cc2e90f160b4f532980858e890c1a8896b2 Mon Sep 17 00:00:00 2001
From: Urvashi Mohnani <umohnani@redhat.com>
Date: Thu, 9 Nov 2023 11:05:20 -0500
Subject: [PATCH] Add progress_bar option to image pull

Add a new arg option called progress_bar that when set to
true will display a progress bar with the image pull progress.
This uses the compat endpoint as that is the data we are able
to parse to displat the progress bar.

Signed-off-by: Urvashi Mohnani <umohnani@redhat.com>
---
 podman/domain/images_manager.py | 52 ++++++++++++++++++++++++++++++++-
 test-requirements.txt           |  1 +
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/podman/domain/images_manager.py b/podman/domain/images_manager.py
index 503dcbf4..fb7e773b 100644
--- a/podman/domain/images_manager.py
+++ b/podman/domain/images_manager.py
@@ -5,8 +5,8 @@
 import logging
 import urllib.parse
 from typing import Any, Dict, Generator, Iterator, List, Mapping, Optional, Union
-
 import requests
+from rich.progress import Progress
 
 from podman import api
 from podman.api import Literal
@@ -259,6 +259,8 @@ def pull(
                 config for this request. auth_config should contain the username and password
                 keys to be valid.
             platform (str) – Platform in the format os[/arch[/variant]]
+            progress_bar (bool) - Display a progress bar with the image pull progress (uses
+                the compat endpoint). Default: False
             tls_verify (bool) - Require TLS verification. Default: True.
             stream (bool) - When True, the pull progress will be published as received.
                 Default: False.
@@ -307,9 +309,26 @@ def pull(
                 params["Variant"] = tokens[2]
 
         stream = kwargs.get("stream", False)
+        # if the user wants a progress bar, we need to use the compat endpoint
+        # so set that to true as well as stream so we can parse that output for the
+        # progress bar
+        progress_bar = kwargs.get("progress_bar", False)
+        if progress_bar:
+            params["compatMode"] = True
+            stream = True
+
         response = self.client.post("/images/pull", params=params, stream=stream, headers=headers)
         response.raise_for_status(not_found=ImageNotFound)
 
+        if progress_bar:
+            tasks = {}
+            print("Pulling", params["reference"])
+            with Progress() as progress:
+                for line in response.iter_lines():
+                    decoded_line = json.loads(line.decode('utf-8'))
+                    self.__show_progress_bar(decoded_line, progress, tasks)
+            return None
+
         if stream:
             return response.iter_lines()
 
@@ -325,6 +344,37 @@ def pull(
                 return self.get(obj["id"])
         return self.resource()
 
+    def __show_progress_bar(self, line, progress, tasks):
+        completed = False
+        if line['status'] == 'Download complete':
+            description = f'[green][Download complete  {line["id"]}]'
+            completed = True
+        elif line['status'] == 'Downloading':
+            description = f'[red][Downloading {line["id"]}]'
+        else:
+            # skip other statuses
+            return
+
+        task_id = line["id"]
+        if task_id not in tasks.keys():
+            if completed:
+                # some layers are really small that they download immediately without showing
+                # anything as Downloading in the stream.
+                # For that case, show a completed progress bar
+                tasks[task_id] = progress.add_task(description, total=100, completed=100)
+            else:
+                tasks[task_id] = progress.add_task(
+                    description, total=line['progressDetail']['total']
+                )
+        else:
+            if completed:
+                # due to the stream, the Download complete output can happen before the Downloading
+                # bar outputs the 100%. So when we detect that the download is in fact complete,
+                # update the progress bar to show 100%
+                progress.update(tasks[task_id], description=description, total=100, completed=100)
+            else:
+                progress.update(tasks[task_id], completed=line['progressDetail']['current'])
+
     def remove(
         self,
         image: Union[Image, str],
diff --git a/test-requirements.txt b/test-requirements.txt
index b208312e..4b51b081 100644
--- a/test-requirements.txt
+++ b/test-requirements.txt
@@ -7,3 +7,4 @@ pylint
 pytest
 requests-mock >= 1.11.0
 tox
+rich >= 12.5.1