schedule: Add the 'count' argument.

This commit is contained in:
Valentin Lorentz 2013-08-04 09:16:07 +02:00
parent 61e5edefd3
commit 6a4750015e
2 changed files with 31 additions and 3 deletions

View File

@ -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:

View File

@ -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: