Added support for the rich progress bar style (#115)

* Added support for the ``rich`` progress bar style (Fixes #96 )

---------

Co-authored-by: sybrenjansen <sybren.jansen@gmail.com>

.github/workflows/github-pages.yml

@@ -16,7 +16,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install setuptools wheel twine
+        pip install setuptools wheel twine rich
         pip install .[dashboard]
         pip install .[dill]
         pip install .[docs]

README.rst

@@ -35,7 +35,7 @@ Features
 - Each worker can have its own state and with convenient worker init and exit functionality this state can be easily
   manipulated (e.g., to load a memory-intensive model only once for each worker without the need of sending it through a
   queue)
-- Progress bar support using tqdm_
+- Progress bar support using tqdm_ (``rich`` and notebook widgets are supported)
 - Progress dashboard support
 - Worker insights to provide insight into your multiprocessing efficiency
 - Graceful and user-friendly exception handling

docs/changelog.rst

@@ -4,15 +4,19 @@ Changelog
 Unreleased
 ----------
 
-* Add the option to only show progress on the dashboard. (`#107`_)
-* Import escape directly from markupsafe, instead of from flask. (`#106`_)
+* Added support for the ``rich`` progress bar style (`#96`_)
+* Added the option to only show progress on the dashboard. (`#107`_)
+* Progress bars are now supported on Windows.
 * Insights now also work when using the ``forkserver`` and ``spawn`` start methods. (`#104`_)
 * When using insights on Windows the arguments of the top 5 longest tasks are now available as well.
-* Progress bars are now supported on Windows.
+* Fixed deprecated ``escape`` import from ``flask`` by importing directly from ``markupsafe``. (`#106`_)
+* Added ``py.typed`` file to prompt ``mypy`` for type checking. (`#108`_)
 
-.. _#108: https://github.com/sybrenjansen/mpire/pull/107
-.. _#107: https://github.com/sybrenjansen/mpire/issues/106
+.. _#96: https://github.com/sybrenjansen/mpire/issues/96
+.. _#107: https://github.com/sybrenjansen/mpire/pull/107
 .. _#104: https://github.com/sybrenjansen/mpire/issues/104
+.. _#106: https://github.com/sybrenjansen/mpire/issues/106
+.. _#108: https://github.com/sybrenjansen/mpire/pull/108
 
 
 2.8.1

docs/index.rst

@@ -18,7 +18,7 @@ Features
 - Each worker can have its own state and with convenient worker init and exit functionality this state can be easily
   manipulated (e.g., to load a memory-intensive model only once for each worker without the need of sending it through a
   queue)
-- Progress bar support using tqdm_
+- Progress bar support using tqdm_ (``rich`` and notebook widgets are supported)
 - Progress dashboard support
 - Worker insights to provide insight into your multiprocessing efficiency
 - Graceful and user-friendly exception handling

docs/install.rst

@@ -18,19 +18,14 @@ and is available through conda-forge:
     conda install -c conda-forge mpire
 
 
-.. note::
-
-    MPIRE is available for Python >= 3.6.
-
 Dependencies
 ------------
 
-- Python >= 3.6
+- Python >= 3.8
 
 Python packages (installed automatically when installing MPIRE):
 
 - tqdm
-- dataclasses (Python 3.6 only)
 - pygments
 - pywin32 (Windows only)
 
@@ -62,6 +57,22 @@ This will install multiprocess_, which uses ``dill`` under the hood. You can ena
 .. _multiprocess: https://github.com/uqfoundation/multiprocess
 .. _BSD license of multiprocess: https://github.com/uqfoundation/multiprocess/blob/master/LICENSE
 
+
+.. _richdep:
+
+Rich progress bars
+~~~~~~~~~~~~~~~~~~
+
+If you want to use rich_ progress bars, you have to install the dependencies for it manually:
+
+.. code-block:: bash
+
+    pip install rich
+
+
+.. _rich: https://github.com/Textualize/rich
+
+
 .. _dashboarddep:
 
 Dashboard

docs/usage/map/progress_bar.rst

@@ -17,6 +17,48 @@ This will display a basic ``tqdm`` progress bar displaying the time elapsed and
 (including a percentage value) and the speed (i.e., number of tasks completed per time unit).
 
 
+.. _progress_bar_style:
+
+Progress bar style
+------------------
+
+You can switch to a different progress bar style by changing the ``progress_bar_style`` parameter. For example, when
+you require a notebook widget use ``'notebook'`` as the style:
+
+.. code-block:: python
+
+    with WorkerPool(n_jobs=4) as pool:
+        pool.map(task, range(100), progress_bar=True, progress_bar_style='notebook')
+
+The available styles are:
+
+- ``None``: use the default style (= ``'std'`` , see below)
+- ``'std'``: use the standard ``tqdm`` progress bar
+- ``'rich'``: use the rich progress bar (requires the ``rich`` package to be installed, see :ref:`richdep`)
+- ``'notebook'``: use the Jupyter notebook widget
+- ``'dashboard'``: use only the progress bar on the dashboard
+
+When in a terminal and using the ``'notebook'`` style, the progress bar will behave weirdly. This is not recommended.
+
+.. note::
+
+    If you run into problems with getting the progress bar to work in a Jupyter notebook (with ``'notebook'`` style),
+    have a look at :ref:`troubleshooting_progress_bar`.
+
+Changing the default style
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+You can change the default style by setting the :obj:`mpire.tqdm_utils.PROGRESS_BAR_DEFAULT_STYLE` variable:
+
+.. code-block:: python
+
+    import mpire.tqdm_utils
+
+    mpire.tqdm_utils.PROGRESS_BAR_DEFAULT_STYLE = 'notebook'
+
+.. _tqdm: https://pypi.python.org/pypi/tqdm
+
+
 Progress bar options
 --------------------
 
@@ -68,45 +110,11 @@ It goes without saying that you shouldn't specify the same progress bar position
 
 .. note::
 
-    Most progress bar options are completely ignored when in a Jupyter/IPython notebook session or in the MPIRE
-    dashboard.
-
-
-.. _progress_bar_style:
-
-Progress bar style
-------------------
-
-You can switch to a notebook widget by changing the ``progress_bar_style`` parameter to ``'notebook'``:
-
-.. code-block:: python
-
-    with WorkerPool(n_jobs=4) as pool:
-        pool.map(task, range(100), progress_bar=True, progress_bar_style='notebook')
-
-The available styles are:
-
-- ``None``: use the default style (= ``'std'`` , see below)
-- ``'std'``: use the standard ``tqdm`` progress bar
-- ``'notebook'``: use the Jupyter notebook widget
-- ``'dashboard'``: use only the progressbar on the dashboard
-
-When in a terminal and using the ``'notebook'`` style, the progress bar will behave weirdly. This is not recommended.
+    When using the ``rich`` progress bar style, the ``position`` parameter cannot be used. An exception will be raised
+    when trying to do so.
 
 .. note::
 
-    If you run into problems with getting the progress bar to work in a Jupyter notebook (with ``'notebook'`` style),
-    have a look at :ref:`troubleshooting_progress_bar`.
-
-Changing the default style
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-You can change the default style by setting the :obj:`mpire.tqdm_utils.PROGRESS_BAR_DEFAULT_STYLE` variable:
-
-.. code-block:: python
-
-    import mpire.tqdm_utils
-
-    mpire.tqdm_utils.PROGRESS_BAR_DEFAULT_STYLE = 'notebook'
+    Most progress bar options are completely ignored when in a Jupyter/IPython notebook session or in the MPIRE
+    dashboard.
 
-.. _tqdm: https://pypi.python.org/pypi/tqdm

mpire/params.py

@@ -2,12 +2,10 @@
 import math
 import multiprocessing as mp
 import warnings
-from io import StringIO
+from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, Iterable, List, Optional, Sized, Tuple, Type, Union
 
-from dataclasses import dataclass, field
-from unittest.mock import patch
-from tqdm import tqdm, TqdmKeyError
+from tqdm import TqdmKeyError
 
 from mpire.context import DEFAULT_START_METHOD
 from mpire.tqdm_utils import get_tqdm
@@ -295,6 +293,11 @@ def check_progress_bar_options(progress_bar_options: Optional[Dict[str, Any]], p
                           "precedence", RuntimeWarning, stacklevel=2)
         else:
             progress_bar_options["position"] = progress_bar_position
+
+    # We currently do not support the position parameter for rich progress bars. Although this can be implemented by
+    # using a single rich progress bar for all workers and using `add_task`, but this is not trivial to implement.
+    if progress_bar_style == "rich" and "position" in progress_bar_options:
+        raise NotImplementedError("The 'position' parameter is currently not supported for rich progress bars")
 
     # Set some defaults and overwrite others
     progress_bar_options["total"] = n_tasks
@@ -305,14 +308,13 @@ def check_progress_bar_options(progress_bar_options: Optional[Dict[str, Any]], p
     progress_bar_options.setdefault("maxinterval", 0.5)
 
     # Check if the tqdm progress bar style is valid
-    get_tqdm(progress_bar_style)
+    tqdm = get_tqdm(progress_bar_style)
 
-    # Check that all progress bar options are properly formatted. We need to that here, because when an error occurs
+    # Check that all progress bar options are properly formatted. We need to do that here, because when an error occurs
     # within the progress bar handler it will deadlock (it's not technically impossible to do it there, but might as
     # well do it here)
     try:
-        with patch(progress_bar_options.get('file', 'sys.stderr'), new=StringIO()):
-            tqdm(**progress_bar_options)
+        tqdm.check_options(progress_bar_options)
     except (TqdmKeyError, TypeError) as e:
         raise e from ValueError("There's an error in progress_bar_options. Either one of the parameters doesn't exist "
                                 "or it's not properly formatted. See tqdm.tqdm() for details.")

mpire/pool.py

@@ -714,7 +714,7 @@ def imap_unordered(self, func: Callable, iterable_of_args: Union[Sized, Iterable
             iterator_of_chunked_args = chunk_tasks(iterable_of_args, n_tasks, chunk_size, n_splits)
 
         # Grab original lock in case we have a progress bar and we need to restore it
-        tqdm, _ = get_tqdm(progress_bar_style)
+        tqdm = get_tqdm(progress_bar_style)
         original_tqdm_lock = tqdm.get_lock()
         tqdm_manager_owner = False
 

mpire/progress_bar.py

@@ -3,8 +3,9 @@
 from datetime import datetime, timedelta
 from threading import Event, Thread
 from typing import Any, Dict, Optional, Type
+import warnings
 
-from tqdm import tqdm as tqdm_type
+from tqdm import TqdmExperimentalWarning, tqdm as tqdm_type
 
 from mpire.comms import WorkerComms, POISON_PILL
 from mpire.exception import remove_highlighting
@@ -106,22 +107,21 @@ def _progress_bar_handler(self) -> None:
         Keeps track of the progress made by the workers and updates the progress bar accordingly
         """
         # Obtain the progress bar tqdm class
-        tqdm, in_notebook = get_tqdm(self.progress_bar_style)
+        tqdm = get_tqdm(self.progress_bar_style)
 
         # Connect to the tqdm manager
         tqdm_manager = TqdmManager()
         tqdm_lock, tqdm_position_register = tqdm_manager.get_lock_and_position_register()
         tqdm.set_lock(tqdm_lock)
-        main_progress_bar = tqdm_position_register.register_progress_bar_position(self.progress_bar_options["position"])
-
-        # In case we're running tqdm in a notebook we need to apply a dirty hack to get progress bars working.
-        # Solution adapted from https://github.com/tqdm/tqdm/issues/485#issuecomment-473338308
-        if in_notebook and not main_progress_bar:
-            print(' ', end='', flush=True)
-
-        # Create progress bar and register the start time
-        tqdm.monitor_interval = False
-        progress_bar = tqdm(**self.progress_bar_options)
+        tqdm.set_main_progress_bar(
+            tqdm_position_register.register_progress_bar_position(self.progress_bar_options["position"])
+        )
+
+        # Create progress bar and register the start time. Ignore the experimental warning for rich progress bars
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", TqdmExperimentalWarning)
+            tqdm.monitor_interval = False
+            progress_bar = tqdm(**self.progress_bar_options)
         self.start_t = datetime.fromtimestamp(progress_bar.start_t)
 
         # Notify that the main process can continue working. We set it after the progress bar has been created, instead
@@ -152,25 +152,14 @@ def _progress_bar_handler(self) -> None:
                     elif self.worker_comms.kill_signal_received():
                         self._send_dashboard_update(progress_bar, failed=True, traceback_str='Kill signal received')
 
-                # Final update of the progress bar. When we're not in a notebook and this is the main progress bar, we
-                # add as many newlines as the highest progress bar position, such that new output is added after the
-                # progress bars.
-                progress_bar.refresh()
-                if in_notebook:
-                    progress_bar.close()
-                else:
-                    progress_bar.disable = True
-                    if main_progress_bar:
-                        progress_bar.fp.write('\n' * (tqdm_position_register.get_highest_progress_bar_position() + 1))
+                # Final update of the progress bar
+                progress_bar.final_refresh(tqdm_position_register.get_highest_progress_bar_position())
                 break
 
             # Check if the total has been updated. It could be that we didn't know the total number of tasks at the
-            # beginning, but we do now. In a notebook we also need to update the max value of the progress bar widget.
+            # beginning, but we do now.
             if self.total_updated.is_set():
-                progress_bar.total = self.total
-                if in_notebook:
-                    progress_bar.container.children[1].max = self.total
-                progress_bar.refresh()
+                progress_bar.update_total(self.total)
                 self._send_dashboard_update(progress_bar)
                 self.total_updated.clear()
 
@@ -180,12 +169,7 @@ def _progress_bar_handler(self) -> None:
 
             # Update progress bar
             progress_bar.update(tasks_completed - progress_bar.n)
-
-            # Force a refresh when we're at 100%
             if progress_bar.n == progress_bar.total:
-                if in_notebook:
-                    progress_bar.close()
-                progress_bar.refresh()
                 self.worker_comms.signal_progress_bar_complete()
                 self.worker_comms.wait_until_progress_bar_is_complete()
                 self._send_dashboard_update(progress_bar)