From 086cfd5599773076ad0d7ef91f82d20030aa068e Mon Sep 17 00:00:00 2001
From: Tobias Gruetzmacher <tobias-git@23.gs>
Date: Sun, 31 Jul 2016 21:58:51 +0200
Subject: [PATCH] Allow usage as a context manager.

---
 README.rst           | 11 +++++++++++
 progress/__init__.py | 15 +++++++++------
 test_progress.py     |  7 ++++---
 3 files changed, 24 insertions(+), 9 deletions(-)

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)):