7061: Wrap User...._progress in a context manager

svn: r23091

gramps/cli/test/user_test.py

@@ -151,5 +151,32 @@ class TestUser_quiet(unittest.TestCase):
         assert len(self.user._fileout.method_calls
                 ) == 0, list(self.user._fileout.method_calls)
 
+@unittest.skipUnless(MOCKING, "Requires unittest.mock to run")
+class TestUser_progress(unittest.TestCase):
+    _progress_begin_step_end = \
+            TestUser_quiet.test_progress_can_begin_step_end.__func__
+
+    def setUp(self):
+        self.user = user.User()
+        self.user._fileout = Mock(spec=sys.stderr)
+        # Collect baseline output from the old-style interface (begin/step/end)
+        self._progress_begin_step_end()
+        self.expected_output = list(self.user._fileout.method_calls)
+        self.user._fileout.reset_mock()
+        self.assertTrue(
+                len(self.user._fileout.method_calls) == 0, 
+                list(self.user._fileout.method_calls))
+
+    def test(self):
+        with self.user.progress("Foo", "Bar", 0) as step:
+            for i in range(10):
+                step()
+
+    def tearDown(self):
+        # Output using `with' differs from one with `progress_...'
+        self.assertEqual(self.expected_output, 
+                list(self.user._fileout.method_calls))
+
+
 if __name__ == "__main__":
     unittest.main()

gramps/gen/user.py

@@ -21,6 +21,7 @@
 #
 
 import sys
+from contextlib import contextmanager
 
 """
 The User class provides basic interaction with the user.
@@ -84,6 +85,13 @@ class User():
         """
         pass
 
+    # Context-manager wrapper of the begin/step/end_progress above
+    @contextmanager
+    def progress(self, *args, **kwargs):
+        self.begin_progress(*args, **kwargs)
+        yield self.step_progress
+        self.end_progress()
+    
     def prompt(self, title, message, accept_label, reject_label):
         """
         Prompt the user with a message to select an alternative.