Allow tasks to remove themselves during execution

The scheduler class is now capable with dealing with tasks which remove
themselves from the scheduler during execution.

Additionally, some optimizations were applied by use of iterators and
some functions better suited for the purpose.

Please peer-review, all tests pass.

1 files changed, 17 insertions, 11 deletions

diff --git a/sleekxmpp/xmlstream/scheduler.py b/sleekxmpp/xmlstream/scheduler.py
index d98dc6c8..b3e50983 100644
--- a/sleekxmpp/xmlstream/scheduler.py
+++ b/sleekxmpp/xmlstream/scheduler.py
@@ -15,6 +15,7 @@
 import time
 import threading
 import logging
+import itertools
 
 from sleekxmpp.util import Queue, QueueEmpty
 
@@ -156,17 +157,23 @@ class Scheduler(object):
                             newtask = self.addq.get(True, 0.1)
                             elapsed += 0.1
                 except QueueEmpty:
-                    cleanup = []
                     self.schedule_lock.acquire()
-                    for task in self.schedule:
-                        if time.time() >= task.next:
-                            updated = True
-                            if not task.run():
-                                cleanup.append(task)
+                    # select only those tasks which are to be executed now
+                    relevant = itertools.takewhile(
+                        lambda task: time.time() >= task.next, self.schedule)
+                    # run the tasks and keep the return value in a tuple
+                    status = map(lambda task: (task, task.run()), relevant)
+                    # remove non-repeating tasks
+                    for task, doRepeat in status:
+                        if not doRepeat:
+                            try:
+                                self.schedule.remove(task)
+                            except ValueError:
+                                pass
                         else:
-                            break
-                    for task in cleanup:
-                        self.schedule.pop(self.schedule.index(task))
+                            # only need to resort tasks if a repeated task has
+                            # been kept in the list.
+                            updated = True
                 else:
                     updated = True
                     self.schedule_lock.acquire()
@@ -174,8 +181,7 @@ class Scheduler(object):
                         self.schedule.append(newtask)
                 finally:
                     if updated:
-                        self.schedule = sorted(self.schedule,
-                                               key=lambda task: task.next)
+                        self.schedule.sort(key=lambda task: task.next)
                     self.schedule_lock.release()
         except KeyboardInterrupt:
             self.run = False