Replace tqdm and coloredlogs with rich (GitHub issue #272)

Author: marcelzwiers

bidscoin/bcoin.py

@@ -7,7 +7,6 @@
 
 import os
 import types
-import coloredlogs
 import logging
 import shutil
 import sys
@@ -21,8 +20,10 @@
 from importlib.util import spec_from_file_location, module_from_spec
 from pathlib import Path
 from typing import Union
-from tqdm import tqdm
-from tqdm.contrib.logging import logging_redirect_tqdm
+from rich.progress import track, Progress, TextColumn, BarColumn, DownloadColumn, TransferSpeedColumn, TimeRemainingColumn, TimeElapsedColumn
+from rich.logging import RichHandler
+from rich.theme import Theme
+from rich.console import Console
 from importlib.util import find_spec
 if find_spec('bidscoin') is None:
     sys.path.append(str(Path(__file__).parents[1]))
@@ -31,22 +32,6 @@
 LOGGER = logging.getLogger(__name__)
 
 
-class TqdmUpTo(tqdm):
-
-    def update_to(self, b=1, bsize=1, tsize=None):
-        """
-        Adds a tqdm progress bar to urllib.request.urlretrieve()
-        https://gist.github.com/leimao/37ff6e990b3226c2c9670a2cd1e4a6f5
-
-        :param b:       Number of blocks transferred so far [default: 1].
-        :param bsize:   Size of each block (in tqdm units) [default: 1].
-        :param tsize:   Total size (in tqdm units). If [default: None] remains unchanged.
-        """
-        if tsize is not None:
-            self.total = tsize
-        self.update(b * bsize - self.n)  # will also set self.n = b * bsize
-
-
 def drmaa_nativespec(specs: str, session) -> str:
     """
     Converts (CLI default) native Torque walltime and memory specifications to the DRMAA implementation (currently only Slurm is supported)
@@ -75,42 +60,42 @@ def drmaa_nativespec(specs: str, session) -> str:
 
 def synchronize(pbatch, jobids: list, event: str, wait: int=15):
     """
-    Shows tqdm progress bars for queued and running DRMAA jobs. Waits until all jobs have finished +
-    some extra wait time to give NAS systems the opportunity to fully synchronize
+    Shows Rich progress bars for queued and running DRMAA jobs. Waits until all jobs have finished +
+    some extra wait time to give NAS systems the opportunity to fully synchronize.
 
     :param pbatch: The DRMAA session
     :param jobids: The job ids
     :param event:  The event that is passed to trackusage()
     :param wait:   The extra wait time for the NAS
     """
 
-    if jobids:
-        match = re.search(r"(slurm|pbs|torque|sge|lsf|condor|uge)", pbatch.drmaaImplementation.lower())
-        trackusage(f"{event}_{match.group(1) if match else 'drmaa'}")
-    else:
+    if not jobids:
         return
 
-    with logging_redirect_tqdm():
+    match = re.search(r"(slurm|pbs|torque|sge|lsf|condor|uge)", pbatch.drmaaImplementation.lower())
+    trackusage(f"{event}_{match.group(1) if match else 'drmaa'}")
+
+    with Progress(TextColumn('{task.description}'), BarColumn(), TextColumn('{task.completed}/{task.total}'), TimeElapsedColumn(), transient=True) as progress:
+
+        qtask = progress.add_task('[white]Queued  ', total=len(jobids))
+        rtask = progress.add_task('[green]Running ', total=len(jobids))
 
-        qbar = tqdm(total=len(jobids), desc='Queued ', unit='job', leave=False, bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]')
-        rbar = tqdm(total=len(jobids), desc='Running', unit='job', leave=False, bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]', colour='green')
         done = 0
         while done < len(jobids):
             jobs   = [pbatch.jobStatus(jobid) for jobid in jobids]
             done   = sum(status in ('done', 'failed', 'undetermined') for status in jobs)
-            qbar.n = sum(status == 'queued_active'                    for status in jobs)
-            rbar.n = sum(status == 'running'                          for status in jobs)
-            qbar.refresh(), rbar.refresh()
+            qcount = sum(status == 'queued_active'                    for status in jobs)
+            rcount = sum(status == 'running'                          for status in jobs)
+            progress.update(qtask, completed=qcount)
+            progress.update(rtask, completed=rcount)
             time.sleep(2)
-        qbar.close(), rbar.close()
 
-        failedjobs = [jobid for jobid in jobids if pbatch.jobStatus(jobid)=='failed']
-        if failedjobs:
+        if failedjobs := [jobid for jobid in jobids if pbatch.jobStatus(jobid) == 'failed']:
             LOGGER.error(f"{len(failedjobs)} HPC jobs failed to run:\n{failedjobs}\nThis may well be due to an underspecified `--cluster` input option (e.g. not enough memory)")
 
-        # Give NAS systems some time to fully synchronize
-        for t in tqdm(range(wait*100), desc='synchronizing', leave=False, bar_format='{l_bar}{bar}| [{elapsed}]'):
-            time.sleep(.01)
+        # Synchronization wait bar
+        for t in track(range(wait*100), description='[cyan]Synchronizing', transient=True):
+            time.sleep(0.01)
 
 
 def setup_logging(logfile: Path=Path()):
@@ -129,11 +114,11 @@ def setup_logging(logfile: Path=Path()):
 
     # Set the default formats
     if DEBUG:
-        fmt  = '%(asctime)s - %(name)s - %(levelname)s | %(message)s'
-        cfmt = '%(levelname)s - %(name)s | %(message)s'
+        fmt  = '%(asctime)s - %(name)s | %(message)s'
+        cfmt = '%(name)s | %(message)s'
     else:
         fmt  = '%(asctime)s - %(levelname)s | %(message)s'
-        cfmt = '%(levelname)s | %(message)s'
+        cfmt = '%(message)s'
     datefmt  = '%Y-%m-%d %H:%M:%S'
 
     # Register custom log levels
@@ -161,13 +146,16 @@ def success(self, message, *args, **kws):
     logger = logging.getLogger()
     logger.setLevel('BCDEBUG' if DEBUG else 'VERBOSE')
 
-    # Add the console streamhandler and bring some color to those boring logs! :-)
-    coloredlogs.install(level='BCDEBUG' if DEBUG else 'VERBOSE' if not logfile.name else 'INFO', fmt=cfmt, datefmt=datefmt)   # NB: Using tqdm sets the streamhandler level to 0, see: https://github.com/tqdm/tqdm/pull/1235
-    coloredlogs.DEFAULT_LEVEL_STYLES['verbose']['color'] = 245  # = Gray
+    # Add the Rich console handler and bring some color to those boring logs! :-)
+    console        = Console(theme=Theme({'logging.level.verbose': 'grey50', 'logging.level.success': 'green bold', 'logging.level.bcdebug': 'bright_yellow'}))
+    consolehandler = RichHandler(console=console, show_time=False, show_level=True, show_path=DEBUG, rich_tracebacks=True, markup=True, level='BCDEBUG' if DEBUG else 'VERBOSE' if not logfile.name else 'INFO')
+    consolehandler.set_name('console')
+    consolehandler.setFormatter(logging.Formatter(fmt=cfmt, datefmt=datefmt))
+    logger.addHandler(consolehandler)
 
     if logfile.name:
 
-        # Add the log filehandler
+        # Add the verbose filehandler
         logfile.parent.mkdir(parents=True, exist_ok=True)      # Create the log dir if it does not exist
         formatter  = logging.Formatter(fmt=fmt, datefmt=datefmt)
         loghandler = logging.FileHandler(logfile)
@@ -184,7 +172,7 @@ def success(self, message, *args, **kws):
         logger.addHandler(errorhandler)
 
     if DEBUG:
-        LOGGER.info('\t<<<<<<<<<< Running BIDScoin in DEBUG mode >>>>>>>>>>')
+        LOGGER.info('\t[bold bright_yellow]<<<<<<<<<< Running BIDScoin in DEBUG mode >>>>>>>>>>')
         settracking('show')
 
 
@@ -550,8 +538,14 @@ def pulltutorialdata(tutorialfolder: str) -> None:
 
     # Download the data
     LOGGER.info(f"Downloading the tutorial dataset...")
-    with TqdmUpTo(unit='B', unit_scale=True, unit_divisor=1024, miniters=1, desc=tutorialtargz.name) as t:
-        urllib.request.urlretrieve(tutorialurl, tutorialtargz, reporthook=t.update_to)  # NB: In case of ssl certificate issues use: with urllib.request.urlopen(tutorialurl, context=ssl.SSLContext()) as data, open(tutorialtargz, 'wb') as targz_fid: shutil.copyfileobj(data, targz_fid)
+    with Progress(TextColumn('[bold blue]{task.fields[filename]}'), BarColumn(), DownloadColumn(), TransferSpeedColumn(), TimeRemainingColumn()) as progress:
+        task = progress.add_task('[cyan]Download', filename=tutorialtargz.name, total=None)
+        def reporthook(blocknum: int, blocksize: int, totalsize: int):
+            if totalsize > 0 and progress.tasks[task].total is None:
+                progress.update(task, total=totalsize)
+            progress.update(task, completed=blocknum * blocksize)
+
+        urllib.request.urlretrieve(tutorialurl, tutorialtargz, reporthook=reporthook)  # NB: In case of ssl certificate issues use: with urllib.request.urlopen(tutorialurl, context=ssl.SSLContext()) as data, open(tutorialtargz, 'wb') as targz_fid: shutil.copyfileobj(data, targz_fid)
 
     # Unzip the data in the target folder
     LOGGER.info(f"Unpacking the downloaded data in: {tutorialfolder}")

bidscoin/bidsapps/deface.py

@@ -14,8 +14,7 @@
 import pandas as pd
 import pydeface.utils as pdu
 import tempfile
-from tqdm import tqdm
-from tqdm.contrib.logging import logging_redirect_tqdm
+from rich.progress import track
 from pathlib import Path
 from importlib.util import find_spec
 if find_spec('bidscoin') is None:
@@ -76,77 +75,76 @@ def deface(bidsfolder: str, pattern: str, participant: list, force: bool, output
             jt.joinFiles           = True
 
         # Loop over bids subject/session-directories
-        with logging_redirect_tqdm():
-            for n, subject in enumerate(tqdm(subjects, unit='subject', colour='green', leave=False), 1):
-
-                subid    = subject.name
-                sessions = lsdirs(subject, 'ses-*')
-                if not sessions:
-                    sessions = [subject]
-                for session in sessions:
-
-                    LOGGER.info('--------------------------------------')
-                    LOGGER.info(f"Processing ({n}/{len(subjects)}): {session}")
-
-                    # Search for images that need to be defaced
-                    sesid = session.name if session.name.startswith('ses-') else ''
-                    for match in sorted([match for match in session.glob(pattern) if '.nii' in match.suffixes]):
-
-                        # Construct the output filename and relative path name (used in BIDS)
-                        match_rel = match.relative_to(session).as_posix()
-                        if not output:
-                            outputfile     = match
-                            outputfile_rel = match_rel
-                        elif output == 'derivatives':
-                            srcent, suffix = match.with_suffix('').stem.rsplit('_', 1)  # Name without suffix, suffix
-                            ext = ''.join(match.suffixes)                               # Account for e.g. '.nii.gz'
-                            outputfile     = bidsdir/'derivatives'/'deface'/subid/sesid/match.parent.name/f"{srcent}_space-orig_{suffix}{ext}"
-                            outputfile_rel = outputfile.relative_to(bidsdir).as_posix()
-                        else:
-                            outputfile     = session/output/match.name
-                            outputfile_rel = outputfile.relative_to(session).as_posix()
-                        outputfile.parent.mkdir(parents=True, exist_ok=True)
-
-                        # Check the json "Defaced" field to see if it has already been defaced
-                        outputjson = outputfile.with_suffix('').with_suffix('.json')
-                        if not force and outputjson.is_file():
-                            with outputjson.open('r') as sidecar:
-                                metadata = json.load(sidecar)
-                            if metadata.get('Defaced'):
-                                LOGGER.info(f"Skipping already defaced image: {match_rel} -> {outputfile_rel}")
-                                continue
-
-                        # Deface the image
-                        LOGGER.info(f"Defacing: {match_rel} -> {outputfile_rel}")
-                        if cluster:
-                            jt.args       = [str(match), '--outfile', str(outputfile), '--force'] + [item for pair in [[f"--{key}",val] for key,val in args.items()] for item in pair]
-                            jt.jobName    = f"deface_{subid}_{sesid}"
-                            jt.outputPath = f"{os.getenv('HOSTNAME')}:{Path.cwd() if DEBUG else tempfile.gettempdir()}/{jt.jobName}.out"
-                            jobids.append(pbatch.runJob(jt))
-                            LOGGER.info(f"Your deface job has been submitted with ID: {jobids[-1]}")
-                        else:
-                            pdu.deface_image(str(match), str(outputfile), force=True, forcecleanup=True, **args)
-
-                        # Add a json sidecar-file with the "Defaced" field
-                        inputjson = match.with_suffix('').with_suffix('.json')
-                        if inputjson.is_file():
-                            with inputjson.open('r') as sidecar:
-                                metadata = json.load(sidecar)
-                        else:
-                            metadata = {}
-                        metadata['Defaced'] = True
-                        with outputjson.open('w') as sidecar:
-                            json.dump(metadata, sidecar, indent=4)
-
-                        # Update the scans.tsv file
-                        scans_tsv  = session/f"{subid}{'_'+sesid if sesid else ''}_scans.tsv"
-                        bidsignore = (bidsdir/'.bidsignore').read_text().splitlines() if (bidsdir/'.bidsignore').is_file() else ['extra_data/']
-                        if output and not bids.check_ignore(output, bidsignore) and scans_tsv.is_file():
-                            LOGGER.info(f"Adding {outputfile_rel} to {scans_tsv}")
-                            scans_table                     = pd.read_csv(scans_tsv, sep='\t', index_col='filename')
-                            scans_table.loc[outputfile_rel] = scans_table.loc[match_rel]
-                            scans_table.sort_values(by=['acq_time','filename'], inplace=True)
-                            scans_table.to_csv(scans_tsv, sep='\t', encoding='utf-8')
+        for n, subject in enumerate(track(subjects, description='[green]Subjects', transient=True), 1):
+
+            subid    = subject.name
+            sessions = lsdirs(subject, 'ses-*')
+            if not sessions:
+                sessions = [subject]
+            for session in sessions:
+
+                LOGGER.info('--------------------------------------')
+                LOGGER.info(f"Processing ({n}/{len(subjects)}): {session}")
+
+                # Search for images that need to be defaced
+                sesid = session.name if session.name.startswith('ses-') else ''
+                for match in sorted([match for match in session.glob(pattern) if '.nii' in match.suffixes]):
+
+                    # Construct the output filename and relative path name (used in BIDS)
+                    match_rel = match.relative_to(session).as_posix()
+                    if not output:
+                        outputfile     = match
+                        outputfile_rel = match_rel
+                    elif output == 'derivatives':
+                        srcent, suffix = match.with_suffix('').stem.rsplit('_', 1)  # Name without suffix, suffix
+                        ext = ''.join(match.suffixes)                               # Account for e.g. '.nii.gz'
+                        outputfile     = bidsdir/'derivatives'/'deface'/subid/sesid/match.parent.name/f"{srcent}_space-orig_{suffix}{ext}"
+                        outputfile_rel = outputfile.relative_to(bidsdir).as_posix()
+                    else:
+                        outputfile     = session/output/match.name
+                        outputfile_rel = outputfile.relative_to(session).as_posix()
+                    outputfile.parent.mkdir(parents=True, exist_ok=True)
+
+                    # Check the json "Defaced" field to see if it has already been defaced
+                    outputjson = outputfile.with_suffix('').with_suffix('.json')
+                    if not force and outputjson.is_file():
+                        with outputjson.open('r') as sidecar:
+                            metadata = json.load(sidecar)
+                        if metadata.get('Defaced'):
+                            LOGGER.info(f"Skipping already defaced image: {match_rel} -> {outputfile_rel}")
+                            continue
+
+                    # Deface the image
+                    LOGGER.info(f"Defacing: {match_rel} -> {outputfile_rel}")
+                    if cluster:
+                        jt.args       = [str(match), '--outfile', str(outputfile), '--force'] + [item for pair in [[f"--{key}",val] for key,val in args.items()] for item in pair]
+                        jt.jobName    = f"deface_{subid}_{sesid}"
+                        jt.outputPath = f"{os.getenv('HOSTNAME')}:{Path.cwd() if DEBUG else tempfile.gettempdir()}/{jt.jobName}.out"
+                        jobids.append(pbatch.runJob(jt))
+                        LOGGER.info(f"Your deface job has been submitted with ID: {jobids[-1]}")
+                    else:
+                        pdu.deface_image(str(match), str(outputfile), force=True, forcecleanup=True, **args)
+
+                    # Add a json sidecar-file with the "Defaced" field
+                    inputjson = match.with_suffix('').with_suffix('.json')
+                    if inputjson.is_file():
+                        with inputjson.open('r') as sidecar:
+                            metadata = json.load(sidecar)
+                    else:
+                        metadata = {}
+                    metadata['Defaced'] = True
+                    with outputjson.open('w') as sidecar:
+                        json.dump(metadata, sidecar, indent=4)
+
+                    # Update the scans.tsv file
+                    scans_tsv  = session/f"{subid}{'_'+sesid if sesid else ''}_scans.tsv"
+                    bidsignore = (bidsdir/'.bidsignore').read_text().splitlines() if (bidsdir/'.bidsignore').is_file() else ['extra_data/']
+                    if output and not bids.check_ignore(output, bidsignore) and scans_tsv.is_file():
+                        LOGGER.info(f"Adding {outputfile_rel} to {scans_tsv}")
+                        scans_table                     = pd.read_csv(scans_tsv, sep='\t', index_col='filename')
+                        scans_table.loc[outputfile_rel] = scans_table.loc[match_rel]
+                        scans_table.sort_values(by=['acq_time','filename'], inplace=True)
+                        scans_table.to_csv(scans_tsv, sep='\t', encoding='utf-8')
 
         if cluster and jobids:
             LOGGER.info('')