Add waiting() context manager to wait reminder of timeout after exiting

easypy/timing.py

@@ -334,6 +334,27 @@ def wait(*args, **kwargs):
     return ret
 
 
+@contextmanager
+@wraps(wait)
+def waiting(timeout, *args, **kwargs):
+    """
+    Wait the reminder of the timeout after exiting the context manager.
+
+    >>> with timing() as t:
+    >>>     with waiting(2):
+    >>>         print('Time after entering waiting():', t.duration)
+    >>>         wait(1)
+    >>>         print('Time after calling wait():', t.duration)
+    >>>     print('Time after exiting waiting():', t.duration)
+    Time after entering waiting(): 9e-06
+    Time after calling wait(): 1.000846
+    Time after exiting waiting(): 2.000853
+    """
+    timer = Timer(expiration=timeout)
+    yield
+    wait(timer.remain, *args, **kwargs)
+
+
 def repeat(timeout, callback, sleep=0.5, progressbar=True):
     pred = lambda: callback() and False  # prevent 'wait' from stopping when the callback returns a nonzero
     return wait(timeout, pred=pred, sleep=sleep, progressbar=progressbar, throw=False)

tests/test_timing.py

@@ -2,6 +2,7 @@
 import pytest
 
 from easypy.timing import iter_wait, wait, repeat, iter_wait_progress, Timer, TimeoutException, PredicateNotSatisfied, timing
+from easypy.timing import waiting
 from easypy.bunch import Bunch
 
 
@@ -179,3 +180,11 @@ def pred(is_final_attempt):
     assert all(iteration == 'regular' for iteration in data[:-2])
     assert data[-2] == 'final'
     assert data[-1] == 'regular'
+
+
+def test_waiting():
+    with timing() as t:
+        with waiting(.2):
+            wait(.1)
+    assert .2 <= t.duration, 'did not wait the reminder of the `waiting` timeout'
+    assert t.duration < .3, 'waited the entire `waiting` timeout without considering time spend inside context manager'