make a Stopwatch usable with the "with" statement

This simplifies its use. Previously:

    w.start('foo')
    ...
    w.stop('foo')

Now:

    with w('foo'):
        ...

With "with", it is immediately clear from the code,
and its indentation, what is being measured. With
the old code, you have to hunt for the start and stop
method calls.

1 files changed, 15 insertions, 0 deletions

diff --git a/morphlib/stopwatch.py b/morphlib/stopwatch.py
index 492738d5..7446f897 100644
--- a/morphlib/stopwatch.py
+++ b/morphlib/stopwatch.py
@@ -23,6 +23,7 @@ class Stopwatch(object):
 
     def __init__(self):
         self.ticks = {}
+        self.context_stack = []
 
     def tick(self, reference_object, name):
         if not reference_object in self.ticks:
@@ -56,3 +57,17 @@ class Stopwatch(object):
         return (delta.days * 24 * 3600 +
                 delta.seconds +
                 operator.truediv(delta.microseconds, 10**6))
+
+    def __call__(self, reference_object):
+        self.context_stack.append(reference_object)
+        return self
+                
+    def __enter__(self):
+        self.start(self.context_stack[-1])
+        return self
+        
+    def __exit__(self, *args):
+        self.stop(self.context_stack[-1])
+        self.context_stack.pop()
+        return False # cause any exception to be re-raised
+