Introduce progress bar class for shared logic

Author: FlorentinD

graphdatascience/query_runner/progress/progress_bar.py

@@ -0,0 +1,55 @@
+from typing import Any, Optional
+
+from tqdm.auto import tqdm
+
+from graphdatascience.query_runner.progress.progress_provider import TaskWithProgress
+
+
+class TqdmProgressBar:
+    def __init__(self, task_name: str, relative_progress: Optional[float], bar_options: dict[str, Any] = {}):
+        root_task_name = task_name
+        if relative_progress is None:  # Qualitative progress report
+            self._tqdm_bar = tqdm(
+                total=None,
+                unit="",
+                desc=root_task_name,
+                bar_format="{desc} [elapsed: {elapsed} {postfix}]",
+                **bar_options,
+            )
+        else:
+            self._tqdm_bar = tqdm(
+                total=100,
+                unit="%",
+                desc=root_task_name,
+                **bar_options,
+            )
+
+    def update(
+        self,
+        status: str,
+        progress: Optional[float],
+        sub_tasks_description: Optional[str] = None,
+    ) -> None:
+        postfix = f"status: {status}, task: {sub_tasks_description}" if sub_tasks_description else f"status: {status}"
+        self._tqdm_bar.set_postfix_str(postfix, refresh=False)
+        if progress is not None:
+            new_progress = progress - self._tqdm_bar.n
+            self._tqdm_bar.update(new_progress)
+        else:
+            self._tqdm_bar.refresh()
+
+    def finish(self, success: bool) -> None:
+        if not success:
+            self._tqdm_bar.set_postfix_str("status: FAILED", refresh=True)
+            return
+
+        if self._tqdm_bar.total is not None:
+            self._tqdm_bar.update(self._tqdm_bar.total - self._tqdm_bar.n)
+        self._tqdm_bar.set_postfix_str("status: FINISHED", refresh=True)
+
+    @staticmethod
+    def _relative_progress(task: TaskWithProgress) -> Optional[float]:
+        try:
+            return float(task.progress_percent.removesuffix("%"))
+        except ValueError:
+            return None

graphdatascience/query_runner/progress/progress_provider.py

@@ -10,6 +10,12 @@ class TaskWithProgress:
     status: str
     sub_tasks_description: Optional[str] = None
 
+    def relative_progress(self) -> Optional[float]:
+        try:
+            return float(self.progress_percent.removesuffix("%"))
+        except ValueError:
+            return None
+
 
 class ProgressProvider(ABC):
     @abstractmethod

graphdatascience/query_runner/progress/query_progress_logger.py

@@ -1,9 +1,10 @@
 import warnings
 from concurrent.futures import Future, ThreadPoolExecutor, wait
-from typing import Any, Callable, NoReturn, Optional
+from typing import Any, Callable, Optional
 
 from pandas import DataFrame
-from tqdm.auto import tqdm
+
+from graphdatascience.query_runner.progress.progress_bar import TqdmProgressBar
 
 from ...server_version.server_version import ServerVersion
 from .progress_provider import ProgressProvider, TaskWithProgress
@@ -28,6 +29,8 @@ def __init__(
         self._polling_interval = polling_interval
         self._progress_bar_options = progress_bar_options
 
+        self._progress_bar_options.setdefault("maxinterval", self._polling_interval)
+
     def run_with_progress_logging(
         self, runnable: DataFrameProducer, job_id: str, database: Optional[str] = None
     ) -> DataFrame:
@@ -58,7 +61,7 @@ def _select_progress_provider(self, job_id: str) -> ProgressProvider:
     def _log(
         self, future: Future[Any], job_id: str, progress_provider: ProgressProvider, database: Optional[str] = None
     ) -> None:
-        pbar: Optional[tqdm[NoReturn]] = None
+        pbar: Optional[TqdmProgressBar] = None
         warn_if_failure = True
 
         while wait([future], timeout=self._polling_interval).not_done:
@@ -83,53 +86,19 @@ def _log(
         if pbar is not None:
             self._finish_pbar(future, pbar)
 
-    def _init_pbar(self, task_with_progress: TaskWithProgress) -> tqdm:  # type: ignore
-        root_task_name = task_with_progress.task_name
-        parsed_progress = QueryProgressLogger._relative_progress(task_with_progress)
-        if parsed_progress is None:  # Qualitative progress report
-            return tqdm(
-                total=None,
-                unit="",
-                desc=root_task_name,
-                maxinterval=self._polling_interval,
-                bar_format="{desc} [elapsed: {elapsed} {postfix}]",
-                **self._progress_bar_options,
-            )
-        else:
-            return tqdm(
-                total=100,
-                unit="%",
-                desc=root_task_name,
-                maxinterval=self._polling_interval,
-                **self._progress_bar_options,
-            )
-
-    def _update_pbar(self, pbar: tqdm, task_with_progress: TaskWithProgress) -> None:  # type: ignore
-        parsed_progress = QueryProgressLogger._relative_progress(task_with_progress)
-        postfix = (
-            f"status: {task_with_progress.status}, task: {task_with_progress.sub_tasks_description}"
-            if task_with_progress.sub_tasks_description
-            else f"status: {task_with_progress.status}"
+    def _update_pbar(self, pbar: TqdmProgressBar, task: TaskWithProgress) -> None:
+        pbar.update(
+            task.status,
+            task.relative_progress(),
+            task.sub_tasks_description,
         )
-        pbar.set_postfix_str(postfix, refresh=False)
-        if parsed_progress is not None:
-            new_progress = parsed_progress - pbar.n
-            pbar.update(new_progress)
-        else:
-            pbar.refresh()
-
-    def _finish_pbar(self, future: Future[Any], pbar: tqdm) -> None:  # type: ignore
-        if future.exception():
-            pbar.set_postfix_str("status: FAILED", refresh=True)
-            return
-
-        if pbar.total is not None:
-            pbar.update(pbar.total - pbar.n)
-        pbar.set_postfix_str("status: FINISHED", refresh=True)
-
-    @staticmethod
-    def _relative_progress(task: TaskWithProgress) -> Optional[float]:
-        try:
-            return float(task.progress_percent.removesuffix("%"))
-        except ValueError:
-            return None
+
+    def _init_pbar(self, task: TaskWithProgress) -> TqdmProgressBar:
+        return TqdmProgressBar(
+            task.task_name,
+            task.relative_progress(),
+            bar_options=self._progress_bar_options,
+        )
+
+    def _finish_pbar(self, future: Future[Any], pbar: TqdmProgressBar) -> None:
+        pbar.finish(future.exception() is None)