Refactor async funcs

Author: pavinjosdev

zypperoni

@@ -29,7 +29,7 @@ from shlex import quote
 import xml.etree.ElementTree as ET
 
 # Constants
-ZYPPERONI_VERSION = "1.0.6"
+ZYPPERONI_VERSION = "1.0.7"
 ZYPPER_PID_FILE = "/run/zypp.pid"
 ZYPPER_ENV = "ZYPP_CURL2=1 ZYPP_PCK_PRELOAD=1 ZYPP_SINGLE_RPMTRANS=1"
 
@@ -142,50 +142,32 @@ def shell_exec(command):
     return output.strip(), res.returncode
 
 # Async function to perform zypper shell commands
-async def zypper_task(lock, UUID, task_type, task_item, total_items, item_counter):
+async def zypper_task(lock, UUID, args, task_item, total_items, item_counter):
     try:
         async with lock:
             uuid = UUID.pop()
         log_messages = {}
         commands = ""
         temp_dir = f"{ZYPPERONI_TMP_DIR}/{uuid}/rootfs"
-        if task_type == "ref":
-            log_messages.update({"start": f"Refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
-            log_messages.update({"success": f"Successfully refreshed repo {task_item!r}"})
-            log_messages.update({"error": f"Error refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
-            log_messages.update({"exception": f"Received SIGINT while refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
+        if args.command_name in ["refresh", "ref"]:
+            log_messages.update({"start": f"{'Force ' if args.force else ''}Refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
+            log_messages.update({"success": f"Successfully {'force ' if args.force else ''}refreshed repo {task_item!r}"})
+            log_messages.update({"error": f"Error {'force ' if args.force else ''}refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
+            log_messages.update({"exception": f"Received SIGINT while {'force ' if args.force else ''}refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
             if not os.path.isdir(temp_dir):
                 commands = refresh_mount_commands + refresh_shell_commands
                 commands = commands.format(
                     uuid=uuid,
-                    refresh_type="refresh",
+                    refresh_type="refresh --force" if args.force else "refresh",
                     repo_alias=task_item,
                 )
             else:
                 commands = refresh_shell_commands.format(
                     uuid=uuid,
-                    refresh_type="refresh",
+                    refresh_type="refresh --force" if args.force else "refresh",
                     repo_alias=task_item,
                 )
-        elif task_type == "force-ref":
-            log_messages.update({"start": f"Force refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
-            log_messages.update({"success": f"Successfully force refreshed repo {task_item!r}"})
-            log_messages.update({"error": f"Error force refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
-            log_messages.update({"exception": f"Received SIGINT while force refreshing repo [{item_counter}/{total_items}] {task_item!r}"})
-            if not os.path.isdir(temp_dir):
-                commands = refresh_mount_commands + refresh_shell_commands
-                commands = commands.format(
-                    uuid=uuid,
-                    refresh_type="refresh --force",
-                    repo_alias=task_item,
-                )
-            else:
-                commands = refresh_shell_commands.format(
-                    uuid=uuid,
-                    refresh_type="refresh --force",
-                    repo_alias=task_item,
-                )
-        elif task_type in ["dup", "dup-download", "in", "in-download", "inr", "inr-download"]:
+        elif args.command_name in ["dist-upgrade", "dup", "install", "in", "install-new-recommends", "inr"]:
             log_messages.update({"start": f"Downloading package [{item_counter}/{total_items}] {task_item!r}"})
             log_messages.update({"success": f"Successfully downloaded package {task_item!r}"})
             log_messages.update({"error": f"Error downloading package [{item_counter}/{total_items}] {task_item!r}"})
@@ -222,14 +204,14 @@ async def zypper_task(lock, UUID, task_type, task_item, total_items, item_counte
         logging.debug(log_messages.get("exception"))
 
 # Async function to perform multiple tasks concurrently
-async def main_task(num_jobs, task_type, task_items, no_confirm=None):
+async def main_task(task_items, args):
     EXCEPTION_OCCUR = False
     # init array of temp dir UUIDs corresponding to max num of jobs
-    UUID = [f"{uuid4()!s}" for _ in range(num_jobs)]
+    UUID = [f"{uuid4()!s}" for _ in range(args.jobs)]
     UUID_UNCHANGED = UUID.copy()
     total_items = len(task_items)
     item_counter = 0
-    if task_type == "in":
+    if args.command_name in ["install", "in"]:
         install_pkgs = task_items.copy()
     try:
         # start processing tasks
@@ -241,13 +223,11 @@ async def main_task(num_jobs, task_type, task_items, no_confirm=None):
             log_messages = {}
             task_item = task_items.pop(0)
             item_counter += 1
-            if task_type == "ref":
-                log_messages.update({"exception": "Received SIGINT while processing tasks to refresh repo"})
-            elif task_type == "force-ref":
-                log_messages.update({"exception": "Received SIGINT while processing tasks to force refresh repo"})
-            elif task_type in ["dup", "dup-download", "in", "in-download", "inr", "inr-download"]:
+            if args.command_name in ["refresh", "ref"]:
+                log_messages.update({"exception": f"Received SIGINT while processing tasks to {'force ' if args.force else ''}refresh repo"})
+            elif args.command_name in ["dist-upgrade", "dup", "install", "in", "install-new-recommends", "inr"]:
                 log_messages.update({"exception": "Received SIGINT while processing tasks to download packages"})
-            asyncio.create_task(zypper_task(lock, UUID, task_type, task_item, total_items, item_counter))
+            asyncio.create_task(zypper_task(lock, UUID, args, task_item, total_items, item_counter))
             await asyncio.sleep(0.1)
         # finished processing all tasks
         tasks = asyncio.all_tasks()
@@ -272,19 +252,19 @@ async def main_task(num_jobs, task_type, task_items, no_confirm=None):
         # release zypper exclusive lock
         release_zypp_lock()
         # perform additional zypper commands (if any) on no exception
-        if not EXCEPTION_OCCUR:
+        if not EXCEPTION_OCCUR and \
+        args.command_name in ["dist-upgrade", "dup", "install", "in", "install-new-recommends", "inr"] and \
+        not args.download_only:
             msg = "Zypperoni has finished its tasks. Handing you over to zypper..."
-            if task_type == "dup":
-                logging.info(color("info", msg))
-                command = f"env {ZYPPER_ENV} zypper {'--non-interactive' if no_confirm else ''} --no-cd dist-upgrade"
+            logging.info(color("info", msg))
+            if args.command_name in ["dist-upgrade", "dup"]:
+                command = f"env {ZYPPER_ENV} zypper {'--non-interactive' if args.no_confirm else ''} --no-cd dist-upgrade"
                 os.system(command)
-            elif task_type == "in":
-                logging.info(color("info", msg))
-                command = f"env {ZYPPER_ENV} zypper {'--non-interactive' if no_confirm else ''} --no-cd install {' '.join(install_pkgs)}"
+            elif args.command_name in ["install", "in"]:
+                command = f"env {ZYPPER_ENV} zypper {'--non-interactive' if args.no_confirm else ''} --no-cd install {' '.join(install_pkgs)}"
                 os.system(command)
-            elif task_type == "inr":
-                logging.info(color("info", msg))
-                command = f"env {ZYPPER_ENV} zypper {'--non-interactive' if no_confirm else ''} --no-cd install-new-recommends"
+            elif args.command_name in ["install-new-recommends", "inr"]:
+                command = f"env {ZYPPER_ENV} zypper {'--non-interactive' if args.no_confirm else ''} --no-cd install-new-recommends"
                 os.system(command)
 
 ################################
@@ -466,7 +446,7 @@ if args.command_name in ["refresh", "ref"]:
         zypperoni_cleanup()
         sys.exit()
     try:
-        asyncio.run(main_task(args.jobs, "force-ref" if args.force else "ref", REPO_ALIAS))
+        asyncio.run(main_task(REPO_ALIAS, args))
     except asyncio.exceptions.CancelledError:
         logging.debug("Received SIGINT for asyncio runner")
     except:
@@ -542,7 +522,7 @@ elif args.command_name in ["dist-upgrade", "dup"]:
         zypperoni_cleanup()
         sys.exit()
     try:
-        asyncio.run(main_task(args.jobs, "dup-download" if args.download_only else "dup", DUP_PKG, args.no_confirm))
+        asyncio.run(main_task(DUP_PKG, args))
     except asyncio.exceptions.CancelledError:
         logging.debug("Received SIGINT for asyncio runner")
     except:
@@ -618,7 +598,7 @@ elif args.command_name in ["install", "in"]:
         zypperoni_cleanup()
         sys.exit()
     try:
-        asyncio.run(main_task(args.jobs, "in-download" if args.download_only else "in", IN_PKG, args.no_confirm))
+        asyncio.run(main_task(IN_PKG, args))
     except asyncio.exceptions.CancelledError:
         logging.debug("Received SIGINT for asyncio runner")
     except:
@@ -685,7 +665,7 @@ elif args.command_name in ["install-new-recommends", "inr"]:
         zypperoni_cleanup()
         sys.exit()
     try:
-        asyncio.run(main_task(args.jobs, "inr-download" if args.download_only else "inr", INR_PKG, args.no_confirm))
+        asyncio.run(main_task(INR_PKG, args))
     except asyncio.exceptions.CancelledError:
         logging.debug("Received SIGINT for asyncio runner")
     except: