Merge pull request #706 from jojje/issue_651

Update progress after iteration (fix #651)

2 files changed, 60 insertions, 19 deletions

diff --git a/click/_termui_impl.py b/click/_termui_impl.py
index 44b4ae4..5cb942d 100644
--- a/click/_termui_impl.py
+++ b/click/_termui_impl.py
@@ -103,7 +103,7 @@ class ProgressBar(object):
         if not self.entered:
             raise RuntimeError('You need to use progress bars in a with block.')
         self.render_progress()
-        return self
+        return self.generator()
 
     def render_finish(self):
         if self.is_hidden:
@@ -131,13 +131,13 @@ class ProgressBar(object):
 
     def format_eta(self):
         if self.eta_known:
-            t = self.eta + 1
+            t = int(self.eta)
             seconds = t % 60
-            t /= 60
+            t //= 60
             minutes = t % 60
-            t /= 60
+            t //= 60
             hours = t % 24
-            t /= 24
+            t //= 24
             if t > 0:
                 days = t
                 return '%dd %02d:%02d:%02d' % (days, hours, minutes, seconds)
@@ -251,23 +251,25 @@ class ProgressBar(object):
         self.current_item = None
         self.finished = True
 
-    def next(self):
+    def generator(self):
+        """
+        Returns a generator which yields the items added to the bar during
+        construction, and updates the progress bar *after* the yielded block
+        returns.
+        """
+        if not self.entered:
+            raise RuntimeError('You need to use progress bars in a with block.')
+
         if self.is_hidden:
-            return next(self.iter)
-        try:
-            rv = next(self.iter)
-            self.current_item = rv
-        except StopIteration:
+            for rv in self.iter:
+                yield rv
+        else:
+            for rv in self.iter:
+                self.current_item = rv
+                yield rv
+                self.update(1)
             self.finish()
             self.render_progress()
-            raise StopIteration()
-        else:
-            self.update(1)
-            return rv
-
-    if not PY2:
-        __next__ = next
-        del next
 
 
 def pager(generator, color=None):
diff --git a/tests/test_termui.py b/tests/test_termui.py
index 1d5118c..23f064f 100644
--- a/tests/test_termui.py
+++ b/tests/test_termui.py
@@ -1,4 +1,5 @@
 import click
+import time
 
 
 def test_progressbar_strip_regression(runner, monkeypatch):
@@ -82,3 +83,41 @@ def test_secho(runner):
         click.secho(None, nl=False)
         bytes = outstreams[0].getvalue()
         assert bytes == b''
+
+
+def test_progressbar_yields_all_items(runner):
+    with click.progressbar(range(3)) as progress:
+        assert len(list(progress)) == 3
+
+
+def test_progressbar_update(runner, monkeypatch):
+    class FakeClock(object):
+        def __init__(self):
+            self.now = time.time()
+
+        def advance_time(self, seconds=1):
+            self.now += seconds
+
+        def time(self):
+            return self.now
+
+    fake_clock = FakeClock()
+
+    @click.command()
+    def cli():
+        with click.progressbar(range(4)) as progress:
+            for _ in progress:
+                fake_clock.advance_time()
+                print("")
+
+    monkeypatch.setattr(time, 'time', fake_clock.time)
+    monkeypatch.setattr(click._termui_impl, 'isatty', lambda _: True)
+    output = runner.invoke(cli, []).output
+
+    lines = [line for line in output.split('\n') if '[' in line]
+
+    assert '  0%' in lines[0]
+    assert ' 25%  00:00:03' in lines[1]
+    assert ' 50%  00:00:02' in lines[2]
+    assert ' 75%  00:00:01' in lines[3]
+    assert '100%          ' in lines[4]