Merge pull request #10 from mortacious/feature/reset_bar

Add a set() function to the ProgressBar class to allow resetting the bar to a specific value

Author: mortacious

build/lib/numba_progress/__init__.py

@@ -0,0 +1,2 @@
+from .progress import ProgressBar
+from ._version import __version__

examples/example_nested_loops.py

@@ -0,0 +1,21 @@
+from sleep import usleep
+import numba as nb
+from numba_progress import ProgressBar
+
+@nb.njit(nogil=True)
+def numba_sleeper(num_iterations, sleep_us, progress):
+    for i in range(num_iterations):
+        progress[0].update()
+        for j in range(num_iterations):
+            usleep(sleep_us)
+            progress[1].update(1)
+        # reset the second progress bar to 0
+        progress[1].set(0)
+
+
+if __name__ == "__main__":
+    num_iterations = 30
+    sleep_time_us = 25_000
+    with ProgressBar(total=num_iterations, ncols=80) as numba_progress1, ProgressBar(total=num_iterations, ncols=80) as numba_progress2:
+        # note: progressbar object must be passed as a tuple (a list will not work due to different treatment in numba)
+        numba_sleeper(num_iterations, sleep_time_us, (numba_progress1, numba_progress2))

examples/example_parallel.py

@@ -4,10 +4,11 @@
 
 
 @nb.njit(nogil=True, parallel=True)
-def numba_parallel_sleeper(num_iterations, sleep_us, progress_hook):
+def numba_parallel_sleeper(num_iterations, sleep_us, progress_hook=None):
     for i in nb.prange(num_iterations):
         usleep(sleep_us)
-        progress_hook.update(1)
+        if progress_hook is not None:
+            progress_hook.update(1)
 
 
 if __name__ == "__main__":

examples/example_sequential.py

@@ -4,10 +4,11 @@
 
 
 @nb.njit(nogil=True)
-def numba_sleeper(num_iterations, sleep_us, progress_hook):
+def numba_sleeper(num_iterations, sleep_us, progress_hook=None):
     for i in range(num_iterations):
         usleep(sleep_us)
-        progress_hook.update(1)
+        if progress_hook is not None:
+            progress_hook.update(1)
 
 
 if __name__ == "__main__":

numba_progress/_version.py

@@ -1 +1 @@
-__version__ = "1.0.0"
+__version__ = "1.1.0"

numba_progress/numba_atomic.py

@@ -21,7 +21,7 @@
 from numba.extending import lower_builtin, type_callable
 from numba.np.arrayobj import basic_indexing, make_array, normalize_indices
 
-__all__ = ["atomic_add", "atomic_sub", "atomic_max", "atomic_min"]
+__all__ = ["atomic_add", "atomic_sub", "atomic_max", "atomic_min", "atomic_xchg"]
 
 
 def atomic_rmw(context, builder, op, arrayty, val, ptr):
@@ -147,4 +147,10 @@ def atomic_min(ary, i, v):
     """
     orig = ary[i]
     ary[i] = min(ary[i], v)
+    return orig
+
+@declare_atomic_array_op("xchg", "xchg", "xchg")
+def atomic_xchg(ary, i, v):
+    orig = ary[i]
+    ary[i] = v
     return orig

numba_progress/progress.py

@@ -7,8 +7,8 @@
 from threading import Thread, Event
 
 from numba.extending import overload_method, typeof_impl, as_numba_type, models, register_model, \
-    make_attribute_wrapper, overload_attribute, unbox, NativeValue, box
-from .numba_atomic import atomic_add
+    make_attribute_wrapper, overload_attribute, unbox, NativeValue, box, lower_getattr, lower_setattr
+from .numba_atomic import atomic_add, atomic_xchg
 from numba import types
 from numba.core import cgutils
 from numba.core.boxing import unbox_array
@@ -92,18 +92,24 @@ def close(self):
         self._tqdm.close()
 
     @property
-    def value(self):
+    def n(self):
         return self.hook[0]
+    
+    def set(self, n=0):
+        atomic_xchg(self.hook, 0, n)
+        self._update_tqdm()
 
     def update(self, n=1):
         atomic_add(self.hook, 0, n)
         self._update_tqdm()
 
     def _update_tqdm(self):
-        value = self.value
-        diff = value - self._last_value
-        self._last_value = value
-        self._tqdm.update(diff)
+        value = self.hook[0]
+        #diff = value - self._last_value
+        #self._last_value = value
+        self._tqdm.n = value
+        self._tqdm.refresh()
+        #self._tqdm.update(diff)
 
     def _update_function(self):
         """Background thread for updating the progress bar.
@@ -151,11 +157,12 @@ def __init__(self, dmm, fe_type):
 make_attribute_wrapper(ProgressBarTypeImpl, 'hook', 'hook')
 
 
-@overload_attribute(ProgressBarTypeImpl, 'value')
+
+@overload_attribute(ProgressBarTypeImpl, 'n')
 def get_value(progress_bar):
-    def getter(progress_bar):
-        return progress_bar.hook[0]
-    return getter
+   def getter(progress_bar):
+       return progress_bar.hook[0]
+   return getter
 
 
 @unbox(ProgressBarTypeImpl)
@@ -186,5 +193,15 @@ def _ol_update(self, n=1):
         def _update_impl(self, n=1):
             atomic_add(self.hook, 0, n)
         return _update_impl
+    
+@overload_method(ProgressBarTypeImpl, "set", jit_options={"nogil": True})
+def _ol_set(self, n=0):
+    """
+    Numpy implementation of the update method.
+    """
+    if isinstance(self, ProgressBarTypeImpl):
+        def _set_impl(self, n=0):
+            atomic_xchg(self.hook, 0, n)
+        return _set_impl