Fixes #11 - Compatibility with newer tqdm versions (>= 4.62.2)

sybrenjansen · Aug 31, 2021 · dec07f9 · dec07f9
1 parent 71fe424
commit dec07f9
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 9 deletions.
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -1,18 +1,28 @@
 Changelog
 =========
 
+Not yet released
+----------------
+
+* Fixed compatibility with newer tqdm versions (``>= 4.62.2``) (`#11`_)
+
+.. _#11: https://github.com/Slimmer-AI/mpire/issues/11
+
 2.2.0
 -----
 
 *(2021-08-30)*
 
-* Added support for Windows. Support has a few caveats:
+* Added support for Windows (`#6`_, `#7`_). Support has a few caveats:
 
   * When using worker insights the arguments of the top 5 longest tasks are not available
   * Progress bar is not supported when using threading as start method
   * When using ``dill`` and an exception occurs, or when the exception occurs in an exit function, it can print
     additional ``OSError`` messages in the terminal, but these can be safely ignored.
 
+.. _#6: https://github.com/Slimmer-AI/mpire/issues/6
+.. _#7: https://github.com/Slimmer-AI/mpire/issues/7
+
 2.1.1
 -----
 

diff --git a/mpire/progress_bar.py b/mpire/progress_bar.py
@@ -167,21 +167,23 @@ def _progress_bar_handler(self, tqdm_connection_details: TqdmConnectionDetails,
                 # and send it to the dashboard, if available.)
                 if self.worker_comms.exception_thrown() or self.worker_comms.kill_signal_received():
                     progress_bar.set_description('Exception occurred, terminating ... ')
-                    if in_notebook:
-                        progress_bar.sp(bar_style='danger')
                     if self.worker_comms.exception_thrown():
                         _, traceback_str = self.worker_comms.get_exception()
                         self._send_update(progress_bar, failed=True, traceback_str=traceback_str)
                         self.worker_comms.task_done_exception()
                     elif self.worker_comms.kill_signal_received():
                         self._send_update(progress_bar, failed=True, traceback_str='Kill signal received')
 
-                # Final update of the progress bar. When 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.
+                # 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()
-                progress_bar.disable = True
-                if main_progress_bar and not in_notebook:
-                    progress_bar.fp.write('\n' * (tqdm_position_register.get_highest_progress_bar_position() + 1))
+                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))
                 if from_queue:
                     self.worker_comms.task_done_progress_bar()
                 break
@@ -193,7 +195,7 @@ def _progress_bar_handler(self, tqdm_connection_details: TqdmConnectionDetails,
             # Force a refresh when we're at 100%
             if progress_bar.n == progress_bar.total:
                 if in_notebook:
-                    progress_bar.sp(bar_style='success')
+                    progress_bar.close()
                 progress_bar.refresh()
                 self.worker_comms.set_progress_bar_complete()
                 self.worker_comms.wait_until_progress_bar_is_complete()