diff --git a/Lib/sched.py b/Lib/sched.py --- a/Lib/sched.py +++ b/Lib/sched.py @@ -35,6 +35,8 @@ __all__ = ["scheduler"] class Event(namedtuple('Event', 'time, priority, action, argument, kwargs')): + def __init__(self, *args, **kwargs): + self._cancelled = False def __eq__(s, o): return (s.time, s.priority) == (o.time, o.priority) def __ne__(s, o): return (s.time, s.priority) != (o.time, o.priority) def __lt__(s, o): return (s.time, s.priority) < (o.time, o.priority) @@ -48,6 +50,7 @@ """Initialize a new instance, passing the time and delay functions""" self._queue = [] + self._cancellations = 0 self.timefunc = timefunc self.delayfunc = delayfunc @@ -78,12 +81,15 @@ If the event is not in the queue, this raises ValueError. """ - self._queue.remove(event) - heapq.heapify(self._queue) + assert isinstance(event, Event), event + if event._cancelled: + raise ValueError("this event was already cancelled") + event._cancelled = True + self._cancellations += 1 def empty(self): """Check whether the queue is empty.""" - return not self._queue + return (len(self._queue) - self._cancellations) <= 0 def run(self): """Execute events until the queue is empty. @@ -112,13 +118,22 @@ delayfunc = self.delayfunc timefunc = self.timefunc pop = heapq.heappop + heapify = heapq.heapify while q: time, priority, action, argument, kwargs = checked_event = q[0] now = timefunc() if now < time: + if self._cancellations > 50 \ + and self._cancellations > (len(self._queue) >> 1): + self._cancellations = 0 + self._queue = [x for x in self._queue if not x._cancelled] + heapify(self._queue) delayfunc(time - now) else: event = pop(q) + if event._cancelled: + self._cancellations -= 1 + continue # Verify that the event was not removed or altered # by another thread after we last looked at q[0]. if event is checked_event: @@ -139,4 +154,8 @@ # With heapq, two events scheduled at the same time will show in # the actual order they would be retrieved. events = self._queue[:] - return map(heapq.heappop, [events]*len(events)) + heappop = heapq.heappop + while events: + ev = heappop(events) + if not ev._cancelled: + yield ev diff --git a/Lib/test/test_sched.py b/Lib/test/test_sched.py --- a/Lib/test/test_sched.py +++ b/Lib/test/test_sched.py @@ -6,7 +6,7 @@ from test import support -class TestCase(unittest.TestCase): +class TestScheduler(unittest.TestCase): def test_enter(self): l = [] @@ -86,8 +86,29 @@ scheduler.run() self.assertEqual(flag, [None]) + +class TestCancellations(unittest.TestCase): + + def test_cancel_twice(self): + scheduler = sched.scheduler() + ev = scheduler.enter(0.01, 1, lambda: 0) + scheduler.cancel(ev) + self.assertRaises(ValueError, scheduler.cancel, ev) + + def test_queue(self): + scheduler = sched.scheduler() + ev = scheduler.enter(0.01, 1, lambda: 0) + scheduler.cancel(ev) + self.assertEqual(list(scheduler.queue), []) + + def test_empty(self): + scheduler = sched.scheduler() + ev = scheduler.enter(0.01, 1, lambda: 0) + scheduler.cancel(ev) + self.assertTrue(scheduler.empty()) + def test_main(): - support.run_unittest(TestCase) + support.run_unittest(TestScheduler, TestCancellations) if __name__ == "__main__": test_main()