Adds tqdm handler to logging_config.py for progress bar integration

1 files changed, 14 insertions, 0 deletions

diff --git a/modules/logging_config.py b/modules/logging_config.py
index 7db23d4b..ce831b5c 100644
--- a/modules/logging_config.py
+++ b/modules/logging_config.py
@@ -1,6 +1,19 @@
 import os

 import logging

 

+from tqdm.auto import tqdm

+

+class TqdmLoggingHandler(logging.Handler):

+    def __init__(self, level=logging.INFO):

+        super().__init__(level)

+

+    def emit(self, record):

+        try:

+            msg = self.format(record)

+            tqdm.write(msg)

+            self.flush()

+        except Exception:

+            self.handleError(record)

 

 def setup_logging(loglevel):

     if loglevel is None:

@@ -12,5 +25,6 @@ def setup_logging(loglevel):
             level=log_level,

             format='%(asctime)s %(levelname)s [%(name)s] %(message)s',

             datefmt='%Y-%m-%d %H:%M:%S',

+            handlers=[TqdmLoggingHandler()]

         )