diff --git a/README.rst b/README.rst index 3f3be76..6c5adfe 100644 --- a/README.rst +++ b/README.rst @@ -34,6 +34,17 @@ To use them, just call ``next`` to advance and ``finish`` to finish: bar.next() bar.finish() +or use any bar of this class as a context manager: + +.. code-block:: python + + from progress.bar import Bar + + with Bar('Processing', max=20) as bar: + for i in range(20): + # Do some work + bar.next() + The result will be a bar like the following: :: Processing |############# | 42/100 diff --git a/progress/__init__.py b/progress/__init__.py index a41f65d..827b0cd 100644 --- a/progress/__init__.py +++ b/progress/__init__.py @@ -73,12 +73,17 @@ class Infinite(object): self.update() def iter(self, it): - try: + with self: for x in it: yield x self.next() - finally: - self.finish() + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.finish() class Progress(Infinite): @@ -119,9 +124,7 @@ class Progress(Infinite): except TypeError: pass - try: + with self: for x in it: yield x self.next() - finally: - self.finish() diff --git a/test_progress.py b/test_progress.py index 0f68b01..5064190 100755 --- a/test_progress.py +++ b/test_progress.py @@ -27,9 +27,10 @@ for bar_cls in (Bar, ChargingBar, FillingSquaresBar, FillingCirclesBar): for bar_cls in (IncrementalBar, PixelBar, ShadyBar): suffix = '%(percent)d%% [%(elapsed_td)s / %(eta)d / %(eta_td)s]' - bar = bar_cls(bar_cls.__name__, suffix=suffix) - for i in bar.iter(range(200)): - sleep() + with bar_cls(bar_cls.__name__, suffix=suffix, max=200) as bar: + for i in range(200): + bar.next() + sleep() for spin in (Spinner, PieSpinner, MoonSpinner, LineSpinner, PixelSpinner): for i in spin(spin.__name__ + ' ').iter(range(100)):