From 5b6c7876425d43a8c1cab9ddd316ed90fb19f773 Mon Sep 17 00:00:00 2001
From: Martin Kristiansen <martin.kristiansen@bodylabs.com>
Date: Fri, 10 Jun 2016 11:31:42 -0400
Subject: [PATCH] TimeoutTimer class (#2)

* Adding TimeoutTimer class
* Unit tests for TimeoutTimer
---
 harrison/__init__.py   |  2 +-
 harrison/test_timer.py | 63 ++++++++++++++++++++++++++++++++++++++++++
 harrison/timer.py      | 39 ++++++++++++++++++++++++++
 3 files changed, 103 insertions(+), 1 deletion(-)
 create mode 100644 harrison/test_timer.py

diff --git a/harrison/__init__.py b/harrison/__init__.py
index 172c537..7df8ac1 100644
--- a/harrison/__init__.py
+++ b/harrison/__init__.py
@@ -1,4 +1,4 @@
 from harrison.profile import profile
 from harrison.timer import Timer
 
-__version__ = '1.0.0'
+__version__ = '1.1.0'
diff --git a/harrison/test_timer.py b/harrison/test_timer.py
new file mode 100644
index 0000000..a2886c0
--- /dev/null
+++ b/harrison/test_timer.py
@@ -0,0 +1,63 @@
+import unittest
+import signal
+from harrison.timer import TimeoutTimer
+from harrison.timer import TimeoutError
+
+class TestTimeoutTimer(unittest.TestCase):
+    TEST_KWARGS_LIST = [
+        {'timeout': 5},
+        {'timeout': 5, 'desc': 'Description'},
+        {'timeout': 5, 'verbose': True},
+        {'timeout': 5, 'desc': 'Description', 'verbose': True}
+    ]
+
+    TEST_ARGS_LIST = [
+        [None],
+        [5]
+    ]
+
+    @staticmethod
+    def expected_timeout(timeout, desc='', verbose=True):
+        _ = (desc, verbose) # for pylint
+        if timeout is None:
+            return 0
+        return timeout
+
+    def test_timeout_timer_is_set_correctly_with_args(self):
+        for test_args in self.TEST_ARGS_LIST:
+            with TimeoutTimer(*test_args):
+                # turns off timer and returns the previous setting in seconds
+                set_number_of_seconds = signal.alarm(0)
+            expected_number_of_seconds = self.expected_timeout(*test_args)
+            self.assertEqual(set_number_of_seconds, expected_number_of_seconds)
+
+    def test_timeout_timer_is_set_correctly_with_kwargs(self):
+        for test_kwargs in self.TEST_KWARGS_LIST:
+            with TimeoutTimer(**test_kwargs):
+                # turns off timer and returns the previous setting in seconds
+                set_number_of_seconds = signal.alarm(0)
+            expected_number_of_seconds = self.expected_timeout(**test_kwargs)
+            self.assertEqual(set_number_of_seconds, expected_number_of_seconds)
+
+    def test_timeout_raises_timeout_error(self):
+        def will_time_out():
+            from time import sleep
+            with TimeoutTimer(timeout=1, desc='Not enough time'):
+                sleep(2) # Seems appropriate to have at least one real timeout
+        self.assertRaises(TimeoutError, will_time_out)
+
+    def test_timeout_does_not_raise_on_clean_exit(self):
+        def will_not_time_out():
+            with TimeoutTimer(timeout=1, desc='Plenty of time'):
+                return 'a_random_return_value'
+        self.assertEqual('a_random_return_value', will_not_time_out())
+
+    def test_two_timeouts_raises(self):
+        def uses_nested_timeouts():
+            with TimeoutTimer(5):
+                with TimeoutTimer(3):
+                    pass
+        self.assertRaises(NotImplementedError, uses_nested_timeouts)
+        # And let's make sure that the alarm is cancelled too
+        leftover_timeout = signal.alarm(0)
+        self.assertEqual(leftover_timeout, 0)
diff --git a/harrison/timer.py b/harrison/timer.py
index aa36753..ba63873 100644
--- a/harrison/timer.py
+++ b/harrison/timer.py
@@ -45,3 +45,42 @@ def __exit__(self, *args):
         if self._verbose:
             desc = '{}: '.format(self.description) if self.description else ''
             print '{}{:.2f} ms'.format(desc, self.elapsed_time_ms)
+
+class TimeoutError(Exception):
+    pass
+
+class TimeoutTimer(Timer):
+    '''
+    Same as Timer but takes a timeout argument. It will raise a
+    TimeoutError if not exitted within that number of seconds.
+    Timer class arguments must be passed using keyword arguments.
+
+    If another TimeoutTimer is already set (or something else that uses alarm signals)
+    when the timer is started, a NotImplementedError will be raised.
+    '''
+    def __init__(self, timeout, **kwargs):
+        self.timeout = timeout
+        super(TimeoutTimer, self).__init__(**kwargs)
+
+    def raise_timeout(self, signal=None, stack_frame=None, msg=None):
+        _ = (signal, stack_frame) # for pylint
+        if msg is None:
+            msg = self.description + ' - Timed out after {} seconds.'.format(
+                self.timeout
+            )
+        raise TimeoutError(msg)
+
+    def start(self):
+        import signal
+        if self.timeout:
+            old_timeout = signal.alarm(self.timeout)
+            if old_timeout:
+                raise NotImplementedError('Nested TimeoutTimers are not supported.')
+            signal.signal(signal.SIGALRM, self.raise_timeout)
+        return super(TimeoutTimer, self).start()
+
+    def stop(self):
+        import signal
+        if self.timeout:
+            signal.alarm(0)
+        return super(TimeoutTimer, self).stop()