Add the context manager flushing()

alphatwirl · May 19, 2024 · 265a5b7 · 265a5b7
1 parent e390719
commit 265a5b7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/atpbar/__init__.py b/atpbar/__init__.py
@@ -3,9 +3,10 @@
     "find_reporter",
     "register_reporter",
     "flush",
+    "flushing",
     "disable",
     "__version__",
 ]
 from .__about__ import __version__
-from .funcs import disable, find_reporter, flush, register_reporter
+from .funcs import disable, find_reporter, flush, flushing, register_reporter
 from .main import atpbar
diff --git a/atpbar/funcs.py b/atpbar/funcs.py
@@ -60,6 +60,15 @@ def flush() -> None:
     _machine.flush()
 
 
+@contextlib.contextmanager
+def flushing() -> Iterator[None]:
+    '''Flushes progress bars on exit'''
+    try:
+        yield
+    finally:
+        flush()
+
+
 def disable() -> None:
     """disables progress bars