Experiments with rich

bio2zarr/cli.py

@@ -1,6 +1,8 @@
+import logging
+
+import rich.logging as rl
 import click
 import tabulate
-import coloredlogs
 
 from . import vcf
 from . import vcf_utils
@@ -22,19 +24,18 @@
 version = click.version_option(version=provenance.__version__)
 
 
-# Note: logging hasn't been implemented in the code at all, this is just
-# a first pass to try out some ways of doing things to see what works.
 def setup_logging(verbosity):
     level = "WARNING"
     if verbosity == 1:
         level = "INFO"
     elif verbosity >= 2:
         level = "DEBUG"
-    # NOTE: I'm not that excited about coloredlogs, just trying it out
-    # as it is installed by cyvcf2 anyway. We will have some complicated
-    # stuff doing on with threads and processes, to logs might not work
-    # so well anyway.
-    coloredlogs.install(level=level)
+    logging.basicConfig(
+        level=level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[rl.RichHandler()],
+    )
 
 
 @click.command
@@ -158,7 +159,9 @@ def vcf2zarr():
 @verbose
 @chunk_length
 @chunk_width
-def convert_plink(in_path, out_path, verbose, worker_processes, chunk_length, chunk_width):
+def convert_plink(
+    in_path, out_path, verbose, worker_processes, chunk_length, chunk_width
+):
     """
     In development; DO NOT USE!
     """

bio2zarr/core.py

@@ -8,7 +8,8 @@
 
 import zarr
 import numpy as np
-import tqdm
+# import tqdm
+import rich.progress as rp
 import numcodecs
 
 
@@ -173,14 +174,26 @@ def __init__(self, worker_processes=1, progress_config=None):
         if progress_config is None:
             progress_config = ProgressConfig()
         self.progress_config = progress_config
-        self.progress_bar = tqdm.tqdm(
-            total=progress_config.total,
-            desc=f"{progress_config.title:>7}",
-            unit_scale=True,
-            unit=progress_config.units,
-            smoothing=0.1,
-            disable=not progress_config.show,
-        )
+        # self.progress_bar = tqdm.tqdm(
+        #     total=progress_config.total,
+        #     desc=f"{progress_config.title:>7}",
+        #     unit_scale=True,
+        #     unit=progress_config.units,
+        #     smoothing=0.1,
+        #     disable=not progress_config.show,
+        # )
+        self.progress_bar = rp.Progress(
+            rp.TimeElapsedColumn(),
+            rp.TextColumn("[progress.description]{task.description}"),
+            rp.BarColumn(),
+            rp.TaskProgressColumn(),
+            rp.MofNCompleteColumn(),
+            rp.TextColumn(f"{progress_config.units} ETA:"),
+            rp.TimeRemainingColumn())
+        self.progress_task = self.progress_bar.add_task(
+            f"{progress_config.title:>7}",
+            total=progress_config.total)
+        self.progress_bar.start()
         self.completed = False
         self.completed_lock = threading.Lock()
         self.progress_thread = threading.Thread(
@@ -191,9 +204,10 @@ def __init__(self, worker_processes=1, progress_config=None):
 
     def _update_progress(self):
         current = get_progress()
-        inc = current - self.progress_bar.n
-        # print("UPDATE PROGRESS: current = ", current, self.progress_config.total, inc)
-        self.progress_bar.update(inc)
+        # inc = current - self.progress_bar.n
+        # print("UPDATE PROGRESS: current = ", current)
+        # self.progress_bar.update(inc)
+        self.progress_bar.update(self.progress_task, completed=current, refresh=True)
 
     def _update_progress_worker(self):
         completed = False
@@ -228,5 +242,5 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         # setting a total of zero or something.
         self.progress_thread.join()
         self._update_progress()
-        self.progress_bar.close()
+        self.progress_bar.stop()
         return False