Rev 3031: Start on bzrlib.progress.TaskStack. in http://people.ubuntu.com/~robertc/baz2.0/nested-pb

[Robert Collins]

At http://people.ubuntu.com/~robertc/baz2.0/nested-pb

------------------------------------------------------------
revno: 3031
revision-id:robertc at robertcollins.net-20071120223910-o578f5kb0g6ppoqv
parent: robertc at robertcollins.net-20071120221812-0pywfzrk21xxxrto
committer: Robert Collins <robertc at robertcollins.net>
branch nick: pb.simplify
timestamp: Wed 2007-11-21 09:39:10 +1100
message:
  Start on bzrlib.progress.TaskStack.
modified:
  bzrlib/progress.py             progress.py-20050610070202-df9faaab791964c0
  bzrlib/tests/test_progress.py  test_progress.py-20060308160359-978c397bc79b7fda
=== modified file 'bzrlib/progress.py'
--- a/bzrlib/progress.py	2007-11-20 21:29:15 +0000
+++ b/bzrlib/progress.py	2007-11-20 22:39:10 +0000
@@ -32,6 +32,7 @@
 
  * CountedTask: Work with an unknown total.
  * KnownLengthTask: Work when we know the total amount to do.
+ * TaskStack: Work when we want to track subordinate tasks.
  * TaskDisplay: Display of a Task.
   * SilentTaskDisplay: For headless operations.
   * DotsTaskDisplay: For non-TTY file stream display.
@@ -40,7 +41,7 @@
 """
 
 # TODO: Optionally show elapsed time instead/as well as ETA; nicer
-# when the rate is unpredictable
+# when the rate is unpredictable.
 
 import sys
 import time
@@ -178,6 +179,24 @@
         return result
 
 
+class TaskStack(KnownLengthTask):
+    """A task which tracks sub-tasks via a stack of simpler tasks."""
+
+    def __init__(self, message, total):
+        KnownLengthTask.__init__(self, message, total)
+        self.tasks = []
+
+    def push_task(self, task):
+        """Add task as a sub-task of this stack."""
+        self.tasks.append(task)
+
+    def pop_task(self, task):
+        """Pop task from the stack, task must be the top of the stack."""
+        if not self.tasks or self.tasks[-1] is not task:
+            raise errors.MissingProgressBarFinish()
+        del self.tasks[-1]
+
+
 class ProgressBarStack(object):
     """A stack of progress bars."""
 

=== modified file 'bzrlib/tests/test_progress.py'
--- a/bzrlib/tests/test_progress.py	2007-11-20 22:18:12 +0000
+++ b/bzrlib/tests/test_progress.py	2007-11-20 22:39:10 +0000
@@ -30,6 +30,7 @@
         ProgressBar,
         ProgressBarStack,
         SilentTaskDisplay,
+        TaskStack,
         TTYProgressBar,
         TTYTaskDisplay,
         )
@@ -428,6 +429,58 @@
         self.assertEqual(1.0, task.percent_complete())
 
 
+class TestTaskStack(TestCase):
+
+    def test_construct_with_message(self):
+        message = 'transferring'
+        stack = TaskStack(message, 1)
+        self.assertEqual(message, stack.get_message())
+        self.assertEqual(1, stack.tick_size)
+        self.assertEqual(0, stack.current)
+        self.assertEqual(1, stack.total)
+        self.assertEqual([], stack.tasks)
+
+    def test_core_push_pop_task(self):
+        message = 'transferring'
+        stack = TaskStack(message, 1)
+        # push and pop one task.
+        task = CountedTask("foo")
+        stack.push_task(task)
+        self.assertEqual([task], stack.tasks)
+        stack.pop_task(task)
+        self.assertEqual([], stack.tasks)
+        # push and pop two likewise.
+        task_2 = KnownLengthTask("bar", 2)
+        stack.push_task(task)
+        stack.push_task(task_2)
+        self.assertEqual([task, task_2], stack.tasks)
+        stack.pop_task(task_2)
+        self.assertEqual([task], stack.tasks)
+        stack.pop_task(task)
+        self.assertEqual([], stack.tasks)
+
+    def test_push_pop_sanity_checks(self):
+        message = 'transferring'
+        stack = TaskStack(message, 1)
+        task = CountedTask("foo")
+        task_2 = KnownLengthTask("bar", 2)
+        # pop when there is nothing to pop
+        self.assertRaises(errors.MissingProgressBarFinish,
+            stack.pop_task, task)
+        self.assertRaises(errors.MissingProgressBarFinish,
+            stack.pop_task, task)
+        # pop a task that is not the top
+        stack.push_task(task)
+        stack.push_task(task_2)
+        self.assertRaises(errors.MissingProgressBarFinish,
+            stack.pop_task, task)
+        # and we can recover and clean up state properly still.
+        self.assertEqual([task, task_2], stack.tasks)
+        stack.pop_task(task_2)
+        stack.pop_task(task)
+        self.assertEqual([], stack.tasks)
+
+
 class TestSilentTaskDisplay(TestCase):
 
     def test_construct(self):