more flexible progress logging

src/kyupy/__init__.py

@@ -198,25 +198,35 @@ class Log:
         """Log an error message."""
         self.log('E', message)
 
+    def progress(self):
+        class Progress:
+            def __init__(self, l): self.l = l
+            def __enter__(self):
+                self.start_time = time.perf_counter()
+                self.lastlog_time = self.start_time
+                self.log_interval = 5
+                return self
+            def update(self, completion: float, msg: str = ''):
+                assert completion > 0 and completion <= 1
+                current_time = time.perf_counter()
+                if current_time > self.lastlog_time + self.log_interval:
+                    elapsed_time = current_time - self.start_time
+                    total_time = elapsed_time / completion
+                    rem_time = total_time - elapsed_time
+                    self.l.log(':', f'{msg} - {completion*100:.0f}% done {hr_time(elapsed_time)} elapsed {hr_time(rem_time)} remaining')
+                    self.log_interval = min(600, int(self.log_interval*1.5))
+                    self.lastlog_time = current_time
+            def __exit__(self, *args): pass
+        return Progress(self)
+
     def range(self, *args):
         """A generator that operates just like the ``range()`` built-in, and also occasionally logs the progress
         and compute time estimates."""
         elems = len(range(*args))
-        start_time = time.perf_counter()
-        lastlog_time = start_time
-        log_interval = 5
-        for elem, i in enumerate(range(*args)):
-            yield i
-            current_time = time.perf_counter()
-            if current_time > lastlog_time + log_interval:
-                done = (elem + 1) / elems
-                elapsed_time = current_time - start_time
-                total_time = elapsed_time / done
-                rem_time = total_time - elapsed_time
-                self.log(
-                    ':', f'{done*100:.0f}% done {hr_time(elapsed_time)} elapsed {hr_time(rem_time)} remaining')
-                log_interval = min(600, int(log_interval*1.5))
-                lastlog_time = current_time
+        with self.progress() as p:
+            for elem, i in enumerate(range(*args)):
+                yield i
+                p.update((elem + 1) / elems)
 
 
 log = Log()