From b6b16a0ff63bf407d4713db154d4a6f57b4fc980 Mon Sep 17 00:00:00 2001 From: Ben Cumming Date: Mon, 3 Jun 2024 17:41:47 +0200 Subject: [PATCH] implement pull status bar --- install | 1 + lib/oras.py | 34 +++++++++++++++++++++++----------- lib/progress.py | 35 +++++++++++++++++++++++++++++++++++ uenv-image | 30 +++++++++++++++++++++++++++++- 4 files changed, 88 insertions(+), 12 deletions(-) create mode 100644 lib/progress.py diff --git a/install b/install index 773d72a..8830703 100755 --- a/install +++ b/install @@ -172,6 +172,7 @@ run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/datastore run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/jfrog.py run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/names.py run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/oras.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/progress.py run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/record.py run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/terminal.py diff --git a/lib/oras.py b/lib/oras.py index 04fac70..66167aa 100644 --- a/lib/oras.py +++ b/lib/oras.py @@ -25,21 +25,33 @@ def find_oras() -> str: return oras_file -def run_command(args): +def run_command(args, detach=False): try: command = [find_oras()] + args terminal.info(f"calling oras: {' '.join(command)}") - result = subprocess.run( - command, - stdout=subprocess.PIPE, # Capture standard output - stderr=subprocess.PIPE, # Capture standard error - check=True, # Raise exception if command fails - encoding='utf-8' # Decode output from bytes to string - ) - - # Print standard output - terminal.info("Output:\n{result.stdout}") + if detach: + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, # Capture standard output + stderr=subprocess.PIPE, # Capture standard error + encoding='utf-8' # Decode output from bytes to string + ) + return process + + else: + result = subprocess.run( + command, + stdout=subprocess.PIPE, # Capture standard output + stderr=subprocess.PIPE, # Capture standard error + check=True, # Raise exception if command fails + encoding='utf-8' # Decode output from bytes to string + ) + + # Print standard output + terminal.info("Output:\n{result.stdout}") + + return None except subprocess.CalledProcessError as e: # Print error message along with captured standard error diff --git a/lib/progress.py b/lib/progress.py new file mode 100644 index 0000000..f3913a7 --- /dev/null +++ b/lib/progress.py @@ -0,0 +1,35 @@ +import math +import sys +from terminal import colorize + +def progress_bar(progress: float, width=25, msg="", stream=sys.stdout): + progress = min(1., max(0., progress)) + + fancy_bar = True + filled_slots = math.floor(width * progress) + lsep = colorize("[", "yellow") + rsep = colorize("]", "yellow") + if filled_slots==width: + if fancy_bar: + bar = colorize("≡"*width, "green") + else: + bar = colorize("="*width, "green") + else: + part_slot = (width * progress) - filled_slots + empty_slots = width - (filled_slots+1) + if fancy_bar: + part_width = math.floor(part_slot * 3) + part_char = [" ", "-", "="][part_width] + bar = colorize("≡"*filled_slots + part_char, "green") + " "*empty_slots + else: + part_width = math.floor(part_slot * 2) + part_char = [" ", "-"][part_width] + bar = colorize("="*filled_slots + part_char, "green") + " "*empty_slots + + bar = lsep + bar + rsep + + pc_str = f"{int(100*progress):3d}%" + + # Use '\r' to return to the start of the line + stream.write('\r ' + bar + " " + pc_str + " " + msg) + stream.flush() diff --git a/uenv-image b/uenv-image index 96e51e1..08dc46e 100755 --- a/uenv-image +++ b/uenv-image @@ -12,6 +12,7 @@ import pathlib import re import sys import textwrap +import time prefix = pathlib.Path(__file__).parent.resolve() libpath = prefix / 'lib' @@ -22,6 +23,7 @@ import datastore import jfrog import names import oras +import progress import record import terminal from terminal import colorize @@ -444,7 +446,33 @@ if __name__ == "__main__": if cache.database.get_record(t.sha256).is_empty: terminal.stdout(f"downloading image {t.sha256} {image_size_string(t.size)}") # download the image using oras - oras.run_command(["pull", "-o", image_path, source_address]) + try: + # run the oras command in a separate process so that this process can + # draw a progress bar. + proc = oras.run_command(["pull", "-o", image_path, source_address], detach=True) + while proc.poll() is None: + total_mb = t.size/(1024*1024) + time.sleep(0.25) + if os.path.exists(image_path): + current_size = os.path.getsize(image_path + "/store.squashfs") + current_mb = current_size / (1024*1024) + p = current_mb/total_mb + msg = f"{int(current_mb)}/{int(total_mb)} MB" + progress.progress_bar(p, width=50, msg=msg) + except KeyboardInterrupt: + proc.terminate() + terminal.stdout("") + terminal.error(f"image pull cancelled by user.") + except Exception as e: + proc.terminate() + terminal.stdout("") + terminal.error(f"image pull failed: {str(e)}") + sys.stdout.write("\n") + sys.stdout.flush() + if proc.returncode == 0: + terminal.info(f"oras download completed successfully\n{proc.stdout}") + else: + terminal.error(f"oras download failed {proc.stderr}") else: terminal.stdout(f"image {t.sha256} is already available locally") # update all the tags associated with the image.