Merge pull request #531 from mmahut/mmahut/rich

feat: introducing Rich for polling progress bars

Author: grafuls

requirements.txt

@@ -7,4 +7,5 @@ pyyaml>=3.10
 aiohttp>=3.10
 setuptools>=39.0
 async-lru>=2.0
+rich>=13.0
 build>=0.7.0

setup.cfg

@@ -33,6 +33,7 @@ install_requires =
     aiohttp>=3.10
     setuptools>=39.0
     async-lru>=2.0
+    rich>=13.0
 package_dir =
     = src
 zip_safe = True

src/badfish/helpers/progress.py

@@ -0,0 +1,19 @@
+from contextlib import contextmanager
+
+from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn
+
+
+@contextmanager
+def polling_progress(console, total, prompt, disable=False):
+    progress = Progress(
+        TextColumn("- POLLING:"),
+        BarColumn(bar_width=20),
+        TaskProgressColumn(),
+        TextColumn("- {task.fields[prompt]}: {task.fields[state]}"),
+        console=console,
+        transient=True,
+        disable=disable or not console.is_terminal,
+    )
+    with progress:
+        task_id = progress.add_task("", total=total, prompt=prompt, state="")
+        yield progress, task_id

src/badfish/main.py

@@ -13,11 +13,14 @@
 import tempfile
 from urllib.parse import urlparse
 
+from rich.console import Console
+
 from badfish.helpers import get_now
 from badfish.helpers.parser import parse_arguments
 from badfish.helpers.logger import BadfishLogger
 from badfish.helpers.http_client import HTTPClient
 from badfish.helpers.exceptions import BadfishException
+from badfish.helpers.progress import polling_progress
 
 from logging import (
     DEBUG,
@@ -30,18 +33,49 @@
 RETRIES = 15
 
 
-async def badfish_factory(_host, _username, _password, _logger=None, _retries=RETRIES, _loop=None, _insecure=False):
+async def badfish_factory(
+    _host,
+    _username,
+    _password,
+    _logger=None,
+    _retries=RETRIES,
+    _loop=None,
+    _insecure=False,
+    _console=None,
+    _progress_disabled=False,
+):
     if not _logger:
         bfl = BadfishLogger()
         _logger = bfl.logger
 
-    badfish = Badfish(_host, _username, _password, _logger, _retries, _loop, _insecure)
+    badfish = Badfish(
+        _host,
+        _username,
+        _password,
+        _logger,
+        _retries,
+        _loop,
+        _insecure,
+        _console,
+        _progress_disabled,
+    )
     await badfish.init()
     return badfish
 
 
 class Badfish:
-    def __init__(self, _host, _username, _password, _logger, _retries, _loop=None, _insecure=False):
+    def __init__(
+        self,
+        _host,
+        _username,
+        _password,
+        _logger,
+        _retries,
+        _loop=None,
+        _insecure=False,
+        _console=None,
+        _progress_disabled=False,
+    ):
         self.host = _host
         self.username = _username
         self.password = _password
@@ -63,6 +97,8 @@ def __init__(self, _host, _username, _password, _logger, _retries, _loop=None, _
         self.session_id = None
         self.token = None
         self.vendor = None
+        self.console = _console if _console is not None else Console()
+        self._progress_disabled = _progress_disabled
 
     async def __aenter__(self):
         await self.init()
@@ -86,19 +122,6 @@ async def init(self):
             self.bios_uri = None
         self.manager_resource = await self.find_managers_resource()
 
-    @staticmethod
-    def progress_bar(value, end_value, state, prompt="Host state", bar_length=20):
-        ratio = float(value) / end_value
-        arrow = "-" * int(round(ratio * bar_length) - 1) + ">"
-        spaces = " " * (bar_length - len(arrow))
-        percent = int(round(ratio * 100))
-
-        if state.lower() == "on":
-            state = "On  "
-        ret = "\r" if percent != 100 else "\n"
-        sys.stdout.write(f"\r- POLLING: [{arrow + spaces}] {percent}% - {prompt}: {state}{ret}")
-        sys.stdout.flush()
-
     async def error_handler(self, _response, message=None):
         try:
             raw = await _response.text("utf-8", "ignore")
@@ -844,35 +867,38 @@ async def check_schedule_job_status(self, job_id):
             return False
 
     async def check_job_status(self, job_id):
-        for count in range(self.retries):
-            _url = f"{self.host_uri}{self.manager_resource}/Jobs/{job_id}"
-            self.http_client.get_request.cache_clear()
-            _response = await self.get_request(_url)
+        with polling_progress(
+            self.console, self.retries, "Status", disable=self._progress_disabled
+        ) as (progress, task_id):
+            for count in range(self.retries):
+                _url = f"{self.host_uri}{self.manager_resource}/Jobs/{job_id}"
+                self.http_client.get_request.cache_clear()
+                _response = await self.get_request(_url)
 
-            status_code = _response.status
-            raw = await _response.text("utf-8", "ignore")
-            data = json.loads(raw.strip())
-            if status_code != 200:
-                self.logger.error(f"Command failed to check job status, return code is {status_code}")
-                self.logger.debug(f"Extended Info Message: {data}")
-                return False
+                status_code = _response.status
+                raw = await _response.text("utf-8", "ignore")
+                data = json.loads(raw.strip())
+                if status_code != 200:
+                    self.logger.error(f"Command failed to check job status, return code is {status_code}")
+                    self.logger.debug(f"Extended Info Message: {data}")
+                    return False
 
-            if "Message" not in data:
-                self.logger.warning("Job status response missing Message field")
-                return False
+                if "Message" not in data:
+                    self.logger.warning("Job status response missing Message field")
+                    return False
 
-            if "Fail" in data["Message"] or "fail" in data["Message"]:
-                self.logger.debug(f"\n{job_id} job failed.")
-                return False
-            elif data["Message"] == "Job completed successfully.":
-                self.logger.info(f"JobID: {data[u'Id']}")
-                self.logger.info(f"Name: {data[u'Name']}")
-                self.logger.info(f"Message: {data[u'Message']}")
-                self.logger.info(f"PercentComplete: {str(data[u'PercentComplete'])}")
-                break
-            else:
-                self.progress_bar(count, self.retries, data["Message"], prompt="Status")
-                await asyncio.sleep(30)
+                if "Fail" in data["Message"] or "fail" in data["Message"]:
+                    self.logger.debug(f"\n{job_id} job failed.")
+                    return False
+                elif data["Message"] == "Job completed successfully.":
+                    self.logger.info(f"JobID: {data[u'Id']}")
+                    self.logger.info(f"Name: {data[u'Name']}")
+                    self.logger.info(f"Message: {data[u'Message']}")
+                    self.logger.info(f"PercentComplete: {str(data[u'PercentComplete'])}")
+                    break
+                else:
+                    progress.update(task_id, completed=count + 1, state=data["Message"])
+                    await asyncio.sleep(30)
 
     def _extract_job_id_from_response(self, response, warn_on_missing=True, context=""):
         """
@@ -1293,32 +1319,38 @@ async def polling_host_state(self, state, equals=True):
         state_str = "Not %s" % state if not equals else state
         self.logger.info("Polling for host state: %s" % state_str)
         desired_state = False
-        for count in range(self.retries):
-            current_state = await self.get_power_state()
-            if equals:
-                desired_state = current_state.lower() == state.lower()
-            else:
-                desired_state = current_state.lower() != state.lower()
-            await asyncio.sleep(5)
-            if desired_state:
-                self.progress_bar(self.retries, self.retries, current_state)
-                break
-            self.progress_bar(count, self.retries, current_state)
+        with polling_progress(
+            self.console, self.retries, "Host state", disable=self._progress_disabled
+        ) as (progress, task_id):
+            for count in range(self.retries):
+                current_state = await self.get_power_state()
+                if equals:
+                    desired_state = current_state.lower() == state.lower()
+                else:
+                    desired_state = current_state.lower() != state.lower()
+                await asyncio.sleep(5)
+                if desired_state:
+                    progress.update(task_id, completed=self.retries, state=current_state)
+                    break
+                progress.update(task_id, completed=count + 1, state=current_state)
 
         return desired_state
 
     async def poll_until_ready(self, check_func, description, sleep_interval=5, clear_cache=False):
         self.logger.info("Polling for %s" % description)
-        for count in range(self.retries):
-            if clear_cache:
-                self.http_client.get_request.cache_clear()
-            ready = await check_func()
-            if ready:
-                self.progress_bar(self.retries, self.retries, "Ready")
-                self.logger.info("%s is ready." % description)
-                return True
-            self.progress_bar(count, self.retries, "Not Ready")
-            await asyncio.sleep(sleep_interval)
+        with polling_progress(
+            self.console, self.retries, "Host state", disable=self._progress_disabled
+        ) as (progress, task_id):
+            for count in range(self.retries):
+                if clear_cache:
+                    self.http_client.get_request.cache_clear()
+                ready = await check_func()
+                if ready:
+                    progress.update(task_id, completed=self.retries, state="Ready")
+                    self.logger.info("%s is ready." % description)
+                    return True
+                progress.update(task_id, completed=count + 1, state="Not Ready")
+                await asyncio.sleep(sleep_interval)
         self.logger.warning("%s did not become ready after %d retry attempts." % (description, self.retries))
         return False
 
@@ -2710,7 +2742,7 @@ async def set_nic_attribute(self, fqdd, attribute, value):
             return True
 
 
-async def execute_badfish(_host, _args, logger, format_handler=None):
+async def execute_badfish(_host, _args, logger, format_handler=None, console=None, progress_disabled=False):
     _username = _args.get("u") or os.environ.get("BADFISH_USERNAME")
     _password = _args.get("p") or os.environ.get("BADFISH_PASSWORD")
 
@@ -2793,6 +2825,8 @@ async def execute_badfish(_host, _args, logger, format_handler=None):
             _logger=logger,
             _retries=retries,
             _insecure=insecure,
+            _console=console,
+            _progress_disabled=progress_disabled,
         )
 
         if _args["host_list"] and not _args["output"]:
@@ -2949,6 +2983,9 @@ def main(argv=None):
     output = _args["output"]
     bfl = BadfishLogger(_args["verbose"], multi_host, _args["log"], output)
 
+    console = Console()
+    progress_disabled = bool(output) or multi_host or bool(_args["log"]) or not console.is_terminal
+
     try:
         loop = asyncio.get_event_loop()
     except RuntimeError:
@@ -2975,6 +3012,8 @@ def main(argv=None):
                         _args,
                         logger,
                         bfl.queue_listener.handlers[0] if output else None,
+                        console=console,
+                        progress_disabled=progress_disabled,
                     )
                     tasks.append(fn)
         except IOError as ex:
@@ -3004,7 +3043,14 @@ def main(argv=None):
     else:
         try:
             _host, result = loop.run_until_complete(
-                execute_badfish(host, _args, bfl.logger, bfl.queue_listener.handlers[0])
+                execute_badfish(
+                    host,
+                    _args,
+                    bfl.logger,
+                    bfl.queue_listener.handlers[0],
+                    console=console,
+                    progress_disabled=progress_disabled,
+                )
             )
         except KeyboardInterrupt:
             bfl.logger.warning("Badfish terminated")

tests/test_execution.py

@@ -22,11 +22,11 @@
 from tests.test_base import TestBase
 
 
-def raise_keyb_interrupt_stub(ignore1, ignore2, ignore3, ignore4=None):
+def raise_keyb_interrupt_stub(ignore1, ignore2, ignore3, ignore4=None, **kwargs):
     raise KeyboardInterrupt
 
 
-def raise_badfish_exception_stub(ignore1, ignore2, ignore3, ignore4=None):
+def raise_badfish_exception_stub(ignore1, ignore2, ignore3, ignore4=None, **kwargs):
     raise BadfishException
 
 

tests/test_progress.py

@@ -0,0 +1,41 @@
+import io
+
+from rich.console import Console
+
+from badfish.helpers.progress import polling_progress
+
+
+def _render(console_kwargs, disable):
+    buf = io.StringIO()
+    console = Console(file=buf, **console_kwargs)
+    with polling_progress(console, 3, "Host state", disable=disable) as (progress, task_id):
+        for i in range(3):
+            progress.update(task_id, completed=i + 1, state="Off")
+    return buf.getvalue()
+
+
+def _has_progress_ui(out):
+    """Progress UI is considered rendered if any bar-specific token appears.
+    Rich may emit a trailing newline in no-op modes on some versions; we only
+    care that no actual POLLING frame was drawn."""
+    return "POLLING" in out or "\x1b[" in out or "Host state" in out
+
+
+def test_disabled_emits_no_progress_ui_even_on_terminal():
+    """Explicit disable=True must suppress all progress rendering."""
+    out = _render({"force_terminal": True, "width": 80}, disable=True)
+    assert not _has_progress_ui(out), f"unexpected progress output: {out!r}"
+
+
+def test_non_terminal_emits_no_progress_ui():
+    """Non-TTY console (pipe/redirect/capsys) must not render the progress bar."""
+    out = _render({"force_terminal": False, "width": 80}, disable=False)
+    assert not _has_progress_ui(out), f"unexpected progress output: {out!r}"
+
+
+def test_terminal_renders_progress_when_enabled():
+    """TTY + disable=False is the only combination that produces visible output."""
+    out = _render({"force_terminal": True, "width": 80}, disable=False)
+    assert "POLLING" in out
+    assert "Host state" in out
+    assert "\x1b[" in out  # ANSI escape sequences from Rich styling