diff --git a/src/schedule.py b/src/schedule.py index 51ef275d4..620ee0b5a 100644 --- a/src/schedule.py +++ b/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: diff --git a/test/test_schedule.py b/test/test_schedule.py index 907a2e618..647bb6713 100644 --- a/test/test_schedule.py +++ b/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: