Add TqdmHandler

yzhangcs · yzhangcs · commit 673839ddf929 · 2021-04-27T00:49:42.000+08:00
diff --git a/supar/utils/logging.py b/supar/utils/logging.py
@@ -11,6 +11,22 @@ def get_logger(name):
     return logging.getLogger(name)
 
 
+class TqdmHandler(logging.StreamHandler):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            tqdm.write(msg)
+            self.flush()
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except Exception:
+            self.handleError(record)
+
+
 def init_logger(logger,
                 path=None,
                 mode='w',
@@ -19,9 +35,9 @@ def init_logger(logger,
                 verbose=True):
     level = level or logging.WARNING
     if not handlers:
-        handlers = [logging.StreamHandler()]
+        handlers = [TqdmHandler()]
         if path:
-            os.makedirs(os.path.dirname(path), exist_ok=True)
+            os.makedirs(os.path.dirname(path) or './', exist_ok=True)
             handlers.append(logging.FileHandler(path, mode))
     logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                         datefmt='%Y-%m-%d %H:%M:%S',
@@ -33,13 +49,15 @@ def init_logger(logger,
 def progress_bar(iterator,
                  ncols=None,
                  bar_format='{l_bar}{bar:18}| {n_fmt}/{total_fmt} {elapsed}<{remaining}, {rate_fmt}{postfix}',
-                 leave=True):
+                 leave=False,
+                 **kwargs):
     return tqdm(iterator,
                 ncols=ncols,
                 bar_format=bar_format,
                 ascii=True,
                 disable=(not (logger.level == logging.INFO and is_master())),
-                leave=leave)
+                leave=leave,
+                **kwargs)
 
 
 logger = get_logger('supar')
