schedule: Add the 'count' argument.

src/schedule.py

@@ -36,6 +36,7 @@ from __future__ import with_statement
 
 import time
 import heapq
+import functools
 from threading import Lock
 
 import supybot.log as log
@@ -109,14 +110,19 @@ class Schedule(drivers.IrcDriver):
         f = self.removeEvent(name)
         self.addEvent(f, t, name=name)
 
-    def addPeriodicEvent(self, f, t, name=None, now=True, args=[], kwargs={}):
+    def addPeriodicEvent(self, f, t, name=None, now=True, args=[], kwargs={},
+            count=None):
         """Adds a periodic event that is called every t seconds."""
-        def wrapper():
+        def wrapper(count):
             try:
                 f(*args, **kwargs)
             finally:
                 # Even if it raises an exception, let's schedule it.
-                return self.addEvent(wrapper, time.time() + t, name)
+                if count[0] is not None:
+                    count[0] -= 1
+                if count[0] is None or count[0] > 0:
+                    return self.addEvent(wrapper, time.time() + t, name)
+        wrapper = functools.partial(wrapper, [count])
         if now:
             return wrapper()
         else:

test/test_schedule.py

@@ -96,6 +96,28 @@ class TestSchedule(SupyTestCase):
         sched.run() # 3.4
         self.assertEqual(i[0], 3)
 
+    def testCountedPeriodic(self):
+        sched = schedule.Schedule()
+        i = [0]
+        def inc():
+            i[0] += 1
+        n = sched.addPeriodicEvent(inc, 1, name='test_periodic', count=3)
+        time.sleep(0.6)
+        sched.run() # 0.6
+        self.assertEqual(i[0], 1)
+        time.sleep(0.6)
+        sched.run() # 1.2
+        self.assertEqual(i[0], 2)
+        time.sleep(0.6)
+        sched.run() # 1.8
+        self.assertEqual(i[0], 2)
+        time.sleep(0.6)
+        sched.run() # 2.4
+        self.assertEqual(i[0], 3)
+        time.sleep(1)
+        sched.run() # 3.4
+        self.assertEqual(i[0], 3)
+
 
 # vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79: