Pass tqdm through logger

Author: bkille

extend.py

@@ -3,6 +3,7 @@
 from Bio.SeqIO import SeqRecord
 from Bio.Align import MultipleSeqAlignment
 from glob import glob
+import logging
 import tempfile
 from pathlib import Path
 import re
@@ -15,7 +16,7 @@
 import numpy as np
 from Bio.AlignIO.MafIO import MafWriter, MafIterator
 from Bio.AlignIO.MauveIO import MauveWriter, MauveIterator
-from logger import logger
+from logger import logger, TqdmToLogger, MIN_TQDM_INTERVAL
 from tqdm import tqdm
 import time
 import spoa
@@ -77,7 +78,7 @@ def xmfa_to_covered(xmfa_file, index_to_gid, gid_to_index_to_cid):
     seqid_parser = re.compile(r'^cluster(\d+) s(\d+):p(\d+)/.*')
     idpair_to_segments = defaultdict(list)
     cluster_to_named_segments = defaultdict(list)
-    for aln in tqdm(AlignIO.parse(xmfa_file, "mauve")):
+    for aln in AlignIO.parse(xmfa_file, "mauve"):
         for seq in aln:
             # Skip reference for now...
             aln_len = seq.annotations["end"] - seq.annotations["start"]
@@ -146,7 +147,11 @@ def extend_clusters(xmfa_file, gid_to_index, gid_to_cid_to_index, idpair_to_segm
     ret_lcbs = []
     seqid_parser = re.compile(r'^cluster(\d+) s(\d+):p(\d+)/.*')
 
-    for aln_idx, aln in enumerate(tqdm(AlignIO.parse(xmfa_file, "mauve"), total=len(cluster_to_named_segments))):
+    for aln_idx, aln in enumerate(tqdm(
+            AlignIO.parse(xmfa_file, "mauve"), 
+            total=len(cluster_to_named_segments),
+            file=TqdmToLogger(logger, level=logging.INFO),
+            mininterval=MIN_TQDM_INTERVAL)):
         # validate_lcb(aln, gid_to_records, parsnp_header=True)
         seq = aln[0]
         cluster_idx, contig_idx, startpos = [int(x) for x in seqid_parser.match(seq.id).groups()]

logger.py

@@ -1,4 +1,5 @@
 import logging
+import io
 ############################################# Logging ##############################################
 BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
 #These are the sequences need to get colored ouput
@@ -14,6 +15,28 @@
 COLOR_SEQ = "\033[1;%dm"
 BOLD_SEQ = "\033[1m"
 
+MIN_TQDM_INTERVAL=30
+
+
+# Logging redirect copied from https://stackoverflow.com/questions/14897756/python-progress-bar-through-logging-module
+class TqdmToLogger(io.StringIO):
+    """
+        Output stream for TQDM which will output to logger module instead of
+        the StdOut.
+    """
+    logger = None
+    level = None
+    buf = ''
+    def __init__(self,logger,level=None):
+        super(TqdmToLogger, self).__init__()
+        self.logger = logger
+        self.level = level or logging.INFO
+    def write(self,buf):
+        self.buf = buf.strip('\r\n\t ')
+    def flush(self):
+        self.logger.log(self.level, self.buf)
+
+
 def formatter_message(message, use_color = True):
     if use_color:
         message = message.replace("$RESET", RESET_SEQ).replace("$BOLD", BOLD_SEQ)

parsnp

@@ -13,7 +13,7 @@ import shlex
 from tempfile import TemporaryDirectory
 import re
 import logging
-from logger import logger
+from logger import logger, TqdmToLogger, MIN_TQDM_INTERVAL
 import multiprocessing
 import argparse
 import signal
@@ -1650,7 +1650,11 @@ SETTINGS:
         logger.info("Running partitions...")
         good_chunks = set(chunk_labels)
         with Pool(args.threads) as pool:
-            return_codes = tqdm(pool.imap(run_parsnp_aligner, chunk_output_dirs, chunksize=1), total=len(chunk_output_dirs))
+            return_codes = tqdm(
+                pool.imap(run_parsnp_aligner, chunk_output_dirs, chunksize=1), 
+                total=len(chunk_output_dirs), 
+                file=TqdmToLogger(logger,level=logging.INFO),
+                mininterval=MIN_TQDM_INTERVAL)
             for cl, rc in zip(chunk_labels, return_codes):
                 if rc != 0:
                     logger.error(f"Partition {cl} failed...")

partition.py

@@ -7,6 +7,7 @@
 import os
 import copy
 import math
+import logging
 from multiprocessing import Pool
 from functools import partial
 from collections import namedtuple, defaultdict, Counter
@@ -18,7 +19,7 @@
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 from Bio.Align import MultipleSeqAlignment
-from logger import logger
+from logger import logger, TqdmToLogger, MIN_TQDM_INTERVAL
 from tqdm import tqdm
 
 
@@ -634,7 +635,11 @@ def trim_xmfas(
     trim_partial = partial(trim_single_xmfa, intersected_interval_dict=intersected_interval_dict)
     orig_xmfa_files = [f"{partition_output_dir}/{CHUNK_PREFIX}-{cl}-out/parsnp.xmfa" for cl in chunk_labels]
     with Pool(threads) as pool:
-        num_clusters_per_xmfa = list(tqdm(pool.imap_unordered(trim_partial, orig_xmfa_files), total=len(orig_xmfa_files)))
+        num_clusters_per_xmfa = list(tqdm(
+            pool.imap_unordered(trim_partial, orig_xmfa_files), 
+            total=len(orig_xmfa_files), 
+            file=TqdmToLogger(logger,level=logging.INFO),
+            mininterval=MIN_TQDM_INTERVAL))
         #TODO clean up
         if not all(num_clusters_per_xmfa[0][1] == nc for xmfa, nc in num_clusters_per_xmfa):
             logger.critical("One of the partitions has a different number of clusters after trimming...")
@@ -712,7 +717,9 @@ def merge_xmfas(
         with Pool(threads) as pool:
             tmp_xmfas = list(tqdm(
                     pool.imap_unordered(merge_single_LCB_star_partial, enumerate(pairs_list)), 
-                    total=num_clusters)
+                    total=num_clusters,
+                    file=TqdmToLogger(logger,level=logging.INFO),
+                    mininterval=MIN_TQDM_INTERVAL)
             )
 
         with open(xmfa_out_f, 'a') as xmfa_out_handle: