CM-45715 - Add rich progress bar with spinner and elapsed time (#286)

Author: MarshalX

cycode/cli/apps/report/sbom/common.py

@@ -30,7 +30,7 @@ def _poll_report_execution_until_completed(
         report_execution = client.get_report_execution(report_execution_id)
         report_label = report_execution.error_message or report_execution.status_message
 
-        progress_bar.update_label(report_label)
+        progress_bar.update_right_side_label(report_label)
 
         if report_execution.status == consts.REPORT_STATUS_COMPLETED:
             return report_execution

cycode/cli/utils/progress_bar.py

@@ -1,16 +1,12 @@
 from abc import ABC, abstractmethod
 from enum import auto
-from typing import TYPE_CHECKING, Dict, NamedTuple, Optional
+from typing import Dict, NamedTuple, Optional
 
-import click
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn, TimeElapsedColumn
 
 from cycode.cli.utils.enum_utils import AutoCountEnum
 from cycode.logger import get_logger
 
-if TYPE_CHECKING:
-    from click._termui_impl import ProgressBar
-    from click.termui import V as ProgressBarValue
-
 # use LOGGING_LEVEL=DEBUG env var to see debug logs of this module
 logger = get_logger('Progress Bar', control_level_in_runtime=False)
 
@@ -32,6 +28,14 @@ class ProgressBarSectionInfo(NamedTuple):
 
 
 _PROGRESS_BAR_LENGTH = 100
+_PROGRESS_BAR_COLUMNS = (
+    SpinnerColumn(),
+    TextColumn('[progress.description]{task.description}'),
+    TextColumn('{task.fields[right_side_label]}'),
+    BarColumn(bar_width=None),
+    TaskProgressColumn(),
+    TimeElapsedColumn(),
+)
 
 ProgressBarSections = Dict[ProgressBarSection, ProgressBarSectionInfo]
 
@@ -91,12 +95,6 @@ class BaseProgressBar(ABC):
     def __init__(self, *args, **kwargs) -> None:
         pass
 
-    @abstractmethod
-    def __enter__(self) -> 'BaseProgressBar': ...
-
-    @abstractmethod
-    def __exit__(self, *args, **kwargs) -> None: ...
-
     @abstractmethod
     def start(self) -> None: ...
 
@@ -110,19 +108,13 @@ def set_section_length(self, section: 'ProgressBarSection', length: int = 0) ->
     def update(self, section: 'ProgressBarSection') -> None: ...
 
     @abstractmethod
-    def update_label(self, label: Optional[str] = None) -> None: ...
+    def update_right_side_label(self, label: Optional[str] = None) -> None: ...
 
 
 class DummyProgressBar(BaseProgressBar):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
-    def __enter__(self) -> 'DummyProgressBar':
-        return self
-
-    def __exit__(self, *args, **kwargs) -> None:
-        pass
-
     def start(self) -> None:
         pass
 
@@ -135,46 +127,49 @@ def set_section_length(self, section: 'ProgressBarSection', length: int = 0) ->
     def update(self, section: 'ProgressBarSection') -> None:
         pass
 
-    def update_label(self, label: Optional[str] = None) -> None:
+    def update_right_side_label(self, label: Optional[str] = None) -> None:
         pass
 
 
 class CompositeProgressBar(BaseProgressBar):
     def __init__(self, progress_bar_sections: ProgressBarSections) -> None:
         super().__init__()
 
-        self._progress_bar_sections = progress_bar_sections
-
-        self._progress_bar_context_manager = click.progressbar(
-            length=_PROGRESS_BAR_LENGTH,
-            item_show_func=self._progress_bar_item_show_func,
-            update_min_steps=0,
-        )
-        self._progress_bar: Optional['ProgressBar'] = None
         self._run = False
+        self._progress_bar_sections = progress_bar_sections
 
         self._section_lengths: Dict[ProgressBarSection, int] = {}
         self._section_values: Dict[ProgressBarSection, int] = {}
 
         self._current_section_value = 0
         self._current_section: ProgressBarSectionInfo = _get_initial_section(self._progress_bar_sections)
+        self._current_right_side_label = ''
 
-    def __enter__(self) -> 'CompositeProgressBar':
-        self._progress_bar = self._progress_bar_context_manager.__enter__()
-        self._run = True
-        return self
+        self._progress_bar = Progress(
+            *_PROGRESS_BAR_COLUMNS,
+            transient=True,
+        )
+        self._progress_bar_task_id = self._progress_bar.add_task(
+            description=self._current_section.label,
+            total=_PROGRESS_BAR_LENGTH,
+            right_side_label=self._current_right_side_label,
+        )
 
-    def __exit__(self, *args, **kwargs) -> None:
-        self._progress_bar_context_manager.__exit__(*args, **kwargs)
-        self._run = False
+    def _progress_bar_update(self, advance: int = 0) -> None:
+        self._progress_bar.update(
+            self._progress_bar_task_id,
+            advance=advance,
+            description=self._current_section.label,
+            right_side_label=self._current_right_side_label,
+        )
 
     def start(self) -> None:
         if not self._run:
-            self.__enter__()
+            self._progress_bar.start()
 
     def stop(self) -> None:
         if self._run:
-            self.__exit__(None, None, None)
+            self._progress_bar.stop()
 
     def set_section_length(self, section: 'ProgressBarSection', length: int = 0) -> None:
         logger.debug('Calling set_section_length, %s', {'section': str(section), 'length': length})
@@ -190,7 +185,7 @@ def _get_section_length(self, section: 'ProgressBarSection') -> int:
         return section_info.stop_percent - section_info.start_percent
 
     def _skip_section(self, section: 'ProgressBarSection') -> None:
-        self._progress_bar.update(self._get_section_length(section))
+        self._progress_bar_update(self._get_section_length(section))
         self._maybe_update_current_section()
 
     def _increment_section_value(self, section: 'ProgressBarSection', value: int) -> None:
@@ -205,13 +200,13 @@ def _increment_section_value(self, section: 'ProgressBarSection', value: int) ->
 
     def _rerender_progress_bar(self) -> None:
         """Used to update label right after changing the progress bar section."""
-        self._progress_bar.update(0)
+        self._progress_bar_update()
 
     def _increment_progress(self, section: 'ProgressBarSection') -> None:
         increment_value = self._get_increment_progress_value(section)
 
         self._current_section_value += increment_value
-        self._progress_bar.update(increment_value)
+        self._progress_bar_update(increment_value)
 
     def _maybe_update_current_section(self) -> None:
         if not self._current_section.section.has_next():
@@ -237,13 +232,7 @@ def _get_increment_progress_value(self, section: 'ProgressBarSection') -> int:
 
         return expected_value - self._current_section_value
 
-    def _progress_bar_item_show_func(self, _: Optional['ProgressBarValue'] = None) -> str:
-        return self._current_section.label
-
     def update(self, section: 'ProgressBarSection', value: int = 1) -> None:
-        if not self._progress_bar:
-            raise ValueError('Progress bar is not initialized. Call start() first or use "with" statement.')
-
         if section not in self._section_lengths:
             raise ValueError(f'{section} section is not initialized. Call set_section_length() first.')
         if section is not self._current_section.section:
@@ -255,12 +244,9 @@ def update(self, section: 'ProgressBarSection', value: int = 1) -> None:
         self._increment_progress(section)
         self._maybe_update_current_section()
 
-    def update_label(self, label: Optional[str] = None) -> None:
-        if not self._progress_bar:
-            raise ValueError('Progress bar is not initialized. Call start() first or use "with" statement.')
-
-        self._progress_bar.label = label or ''
-        self._progress_bar.render_progress()
+    def update_right_side_label(self, label: Optional[str] = None) -> None:
+        self._current_right_side_label = f'({label})' or ''
+        self._progress_bar_update()
 
 
 def get_progress_bar(*, hidden: bool, sections: ProgressBarSections) -> BaseProgressBar:
@@ -284,9 +270,9 @@ def get_progress_bar(*, hidden: bool, sections: ProgressBarSections) -> BaseProg
 
         for _i in range(section_capacity):
             time.sleep(0.01)
-            bar.update_label(f'{bar_section} {_i}/{section_capacity}')
+            bar.update_right_side_label(f'{bar_section} {_i}/{section_capacity}')
             bar.update(bar_section)
 
-        bar.update_label()
+        bar.update_right_side_label()
 
     bar.stop()

cycode/cyclient/__init__.py

@@ -1,5 +0,0 @@
-from cycode.cyclient.config import logger
-
-__all__ = [
-    'logger',
-]

cycode/cyclient/cycode_client_base.py

@@ -15,8 +15,9 @@
     RequestSslError,
     RequestTimeout,
 )
-from cycode.cyclient import config, logger
+from cycode.cyclient import config
 from cycode.cyclient.headers import get_cli_user_agent, get_correlation_id
+from cycode.cyclient.logger import logger
 
 
 class SystemStorageSslContext(HTTPAdapter):

cycode/cyclient/headers.py

@@ -6,7 +6,7 @@
 from cycode.cli import consts
 from cycode.cli.user_settings.configuration_manager import ConfigurationManager
 from cycode.cli.utils.sentry import add_correlation_id_to_scope
-from cycode.cyclient import logger
+from cycode.cyclient.logger import logger
 
 
 def get_cli_user_agent() -> str: