74 lines
1.9 KiB
Python
74 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Allows the bot to scheduler tasks to be performed later.
|
|
"""
|
|
import pickle
|
|
import threading
|
|
from datetime import datetime
|
|
|
|
stop = threading.Event()
|
|
|
|
class CronThread(threading.Thread):
|
|
def __init__(self, scheduler):
|
|
super().__init__()
|
|
self.scheduler = scheduler
|
|
|
|
def run(self):
|
|
while not self.scheduler.bot.stillConnected():
|
|
stop.wait(1)
|
|
while not stop.is_set():
|
|
self.scheduler.run_tasks()
|
|
stop.wait(5)
|
|
|
|
class Scheduler:
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.tasks = []
|
|
self.lock = threading.Lock()
|
|
self.init_database()
|
|
self.load_database()
|
|
self.bot.add_shutdown(stop.set)
|
|
self.loop = CronThread(self)
|
|
self.loop.start()
|
|
|
|
def run_tasks(self):
|
|
self.lock.acquire()
|
|
tasks_due = [t for t in self.tasks if t[1] <= datetime.now()]
|
|
for task in tasks_due:
|
|
args = (self.bot,) + (task[1],) + task[2]
|
|
t = threading.Thread(target=task[0], args=args)
|
|
t.start()
|
|
self.tasks.remove(task)
|
|
self.bot.db.execute(
|
|
"DELETE FROM scheduled_task WHERE dt = ?",
|
|
(pickle.dumps(task[1]),))
|
|
self.lock.release()
|
|
|
|
|
|
def add_task(self, func, dt, args):
|
|
"""
|
|
`func` - The function to call. Must accept `bot` as the first
|
|
argument. Must be picklable.
|
|
`dt` - A datetime object representing when to call the function.
|
|
`args` - Arguments to call the function with, not including `bot`.
|
|
"""
|
|
self.lock.acquire()
|
|
self.tasks.append((func, dt, args))
|
|
self.lock.release()
|
|
|
|
t = tuple(pickle.dumps(i) for i in (func, dt, args))
|
|
self.bot.db.execute("INSERT INTO scheduled_task VALUES (?,?,?)", t)
|
|
|
|
def init_database(self):
|
|
self.bot.db.execute("CREATE TABLE IF NOT EXISTS scheduled_task ("
|
|
"func BLOB,"
|
|
"dt BLOB,"
|
|
"args BLOB"
|
|
")")
|
|
|
|
def load_database(self):
|
|
tasks = self.bot.db.execute("SELECT * FROM scheduled_task").fetchall()
|
|
for task in tasks:
|
|
t = tuple(pickle.loads(i) for i in task)
|
|
self.tasks.append(t)
|