You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
409 lines
13 KiB
409 lines
13 KiB
from datetime import datetime, timedelta
|
|
from typing import Dict, Union
|
|
|
|
import discord
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.base import BaseTrigger
|
|
from apscheduler.triggers.combining import OrTrigger
|
|
from apscheduler.triggers.date import DateTrigger
|
|
from apscheduler.triggers.interval import IntervalTrigger
|
|
from dateutil import parser
|
|
from redbot.core import Config, checks, commands
|
|
from redbot.core.bot import Red
|
|
from redbot.core.commands import DictConverter, TimedeltaConverter, parse_timedelta
|
|
|
|
from .datetimeconverter import DatetimeConverter
|
|
from .redconfigjobstore import RedConfigJobStore
|
|
|
|
|
|
def get_trigger(data):
|
|
if data["type"] == "interval":
|
|
parsed_time = data["time_data"]
|
|
return IntervalTrigger(days=parsed_time.days, seconds=parsed_time.seconds)
|
|
|
|
if data["type"] == "date":
|
|
return DateTrigger(data["time_data"])
|
|
|
|
if data["type"] == "cron":
|
|
return None # TODO: Cron parsing
|
|
|
|
return False
|
|
|
|
|
|
def parse_triggers(data: Union[Dict, None]):
|
|
if data is None or not data.get("triggers", False): # No triggers
|
|
return None
|
|
|
|
if len(data["triggers"]) > 1: # Multiple triggers
|
|
return OrTrigger(get_trigger(t_data) for t_data in data["triggers"])
|
|
|
|
return get_trigger(data[0])
|
|
|
|
|
|
class FakeMessage:
|
|
_state = None
|
|
|
|
|
|
# class FakeMessage(discord.Message):
|
|
# def __init__(self):
|
|
# super().__init__(state=None, channel=None, data=None)
|
|
|
|
|
|
class Task:
|
|
default_task_data = {"triggers": [], "command_str": ""}
|
|
|
|
default_trigger = {
|
|
"type": "",
|
|
"time_data": None, # Used for Interval and Date Triggers
|
|
}
|
|
|
|
def __init__(self, name: str, guild_id, config: Config, author_id=None, bot: Red = None):
|
|
self.name = name
|
|
self.guild_id = guild_id
|
|
self.config = config
|
|
self.bot = bot
|
|
self.author_id = author_id
|
|
self.data = None
|
|
|
|
async def _encode_time_data(self):
|
|
if not self.data or not self.data.get("triggers", None):
|
|
return None
|
|
|
|
triggers = []
|
|
for t in self.data["triggers"]:
|
|
if t["type"] == "interval": # Convert into timedelta
|
|
td: timedelta = t["time_data"]
|
|
|
|
triggers.append({"type": t["type"], "time_data": {"days": td.days, "seconds": td.seconds} })
|
|
|
|
if t["type"] == "date": # Convert into datetime
|
|
dt: datetime = t["time_data"]
|
|
triggers.append({"type": t["type"], "time_data": {
|
|
"year": dt.year,
|
|
"month": dt.month,
|
|
"day": dt.day,
|
|
"hour": dt.hour,
|
|
"minute": dt.minute,
|
|
"second": dt.second,
|
|
}})
|
|
|
|
if t["type"] == "cron":
|
|
raise NotImplemented
|
|
raise NotImplemented
|
|
|
|
return triggers
|
|
|
|
async def _decode_time_data(self):
|
|
if not self.data or not self.data.get("triggers", None):
|
|
return
|
|
|
|
for t in self.data["triggers"]:
|
|
if t["type"] == "interval": # Convert into timedelta
|
|
t["time_data"] = timedelta(**t["time_data"])
|
|
|
|
if t["type"] == "date": # Convert into datetime
|
|
t["time_data"] = datetime(**t["time_data"])
|
|
|
|
if t["type"] == "cron":
|
|
raise NotImplemented
|
|
raise NotImplemented
|
|
|
|
# async def load_from_data(self, data: Dict):
|
|
# self.data = data.copy()
|
|
|
|
async def load_from_config(self):
|
|
data = await self.config.guild_from_id(self.guild_id).tasks.get_raw(
|
|
self.name, default=None
|
|
)
|
|
|
|
if not data:
|
|
return
|
|
|
|
self.author_id = data["author_id"]
|
|
self.guild_id = data["guild_id"]
|
|
|
|
self.data = data["data"]
|
|
|
|
await self._decode_time_data()
|
|
return self.data
|
|
|
|
async def get_trigger(self) -> Union[BaseTrigger, None]:
|
|
if not self.data:
|
|
await self.load_from_config()
|
|
|
|
return parse_triggers(self.data)
|
|
|
|
# async def set_job_id(self, job_id):
|
|
# if self.data is None:
|
|
# await self.load_from_config()
|
|
#
|
|
# self.data["job_id"] = job_id
|
|
|
|
async def save_all(self):
|
|
"""To be used when creating an new task"""
|
|
|
|
data_to_save = self.default_task_data.copy()
|
|
if self.data:
|
|
data_to_save["command_str"] = self.data.get("command_str", "")
|
|
data_to_save["triggers"] = await self._encode_time_data()
|
|
|
|
to_save = {
|
|
"guild_id": self.guild_id,
|
|
"author_id": self.author_id,
|
|
"data": data_to_save,
|
|
}
|
|
await self.config.guild_from_id(self.guild_id).tasks.set_raw(self.name, value=to_save)
|
|
|
|
async def save_data(self):
|
|
"""To be used when updating triggers"""
|
|
if not self.data:
|
|
return
|
|
await self.config.guild_from_id(self.guild_id).tasks.set_raw(
|
|
self.name, "data", value=await self._encode_time_data()
|
|
)
|
|
|
|
async def execute(self):
|
|
if not self.data or self.data["command_str"]:
|
|
return False
|
|
message = FakeMessage()
|
|
message.guild = self.bot.get_guild(self.guild_id) # used for get_prefix
|
|
message.author = message.guild.get_member(self.author_id)
|
|
message.content = await self.bot.get_prefix(message) + self.data["command_str"]
|
|
|
|
if not message.guild or not message.author or not message.content:
|
|
return False
|
|
|
|
new_ctx: commands.Context = await self.bot.get_context(message)
|
|
if not new_ctx.valid:
|
|
return False
|
|
|
|
await self.bot.invoke(new_ctx)
|
|
return True
|
|
|
|
async def set_bot(self, bot: Red):
|
|
self.bot = bot
|
|
|
|
async def set_author(self, author: Union[discord.User, str]):
|
|
self.author_id = getattr(author, "id", None) or author
|
|
|
|
async def set_commmand_str(self, command_str):
|
|
if not self.data:
|
|
self.data = self.default_task_data.copy()
|
|
self.data["command_str"] = command_str
|
|
return True
|
|
|
|
async def add_trigger(self, param, parsed_time: Union[timedelta, datetime]):
|
|
trigger_data = {"type": param, "time_data": parsed_time}
|
|
if not get_trigger(trigger_data):
|
|
return False
|
|
|
|
if not self.data:
|
|
self.data = self.default_task_data.copy()
|
|
|
|
self.data["triggers"].append(trigger_data)
|
|
return True
|
|
|
|
|
|
class FIFO(commands.Cog):
|
|
"""
|
|
Simple Scheduling Cog
|
|
|
|
Named after the simplest scheduling algorithm: First In First Out
|
|
"""
|
|
|
|
def __init__(self, bot: Red):
|
|
super().__init__()
|
|
self.bot = bot
|
|
self.config = Config.get_conf(self, identifier=70737079, force_registration=True)
|
|
|
|
default_global = {"jobs_index": {}, "jobs": []}
|
|
default_guild = {"tasks": {}}
|
|
|
|
self.config.register_global(**default_global)
|
|
self.config.register_guild(**default_guild)
|
|
|
|
jobstores = {"default": RedConfigJobStore(self.config, self.bot)}
|
|
|
|
job_defaults = {"coalesce": False, "max_instances": 1}
|
|
|
|
# executors = {"default": AsyncIOExecutor()}
|
|
|
|
# Default executor is already AsyncIOExecutor
|
|
self.scheduler = AsyncIOScheduler(jobstores=jobstores, job_defaults=job_defaults)
|
|
|
|
self.scheduler.start()
|
|
|
|
async def red_delete_data_for_user(self, **kwargs):
|
|
"""Nothing to delete"""
|
|
return
|
|
|
|
def _assemble_job_id(self, task_name, guild_id):
|
|
return task_name + "_" + guild_id
|
|
|
|
async def _check_parsable_command(self, ctx: commands.Context, command_to_parse: str):
|
|
message = FakeMessage()
|
|
message.content = ctx.prefix + command_to_parse
|
|
message.author = ctx.author
|
|
message.guild = ctx.guild
|
|
|
|
new_ctx: commands.Context = await self.bot.get_context(message)
|
|
|
|
return new_ctx.valid
|
|
|
|
async def _get_job(self, task_name, guild_id):
|
|
return self.scheduler.get_job(self._assemble_job_id(task_name, guild_id))
|
|
|
|
async def _add_job(self, task):
|
|
return self.scheduler.add_job(
|
|
task.execute,
|
|
id=self._assemble_job_id(task.name, task.guild_id),
|
|
trigger=await task.get_trigger(),
|
|
)
|
|
|
|
@checks.is_owner()
|
|
@commands.command()
|
|
async def fifoclear(self, ctx: commands.Context):
|
|
"""Debug command to clear fifo config"""
|
|
await self.config.guild(ctx.guild).tasks.clear()
|
|
await ctx.tick()
|
|
|
|
@checks.is_owner() # Will be reduced when I figure out permissions later
|
|
@commands.group()
|
|
async def fifo(self, ctx: commands.Context):
|
|
"""
|
|
Base command for handling scheduling of tasks
|
|
"""
|
|
if ctx.invoked_subcommand is None:
|
|
pass
|
|
|
|
@fifo.command(name="list")
|
|
async def fifo_list(self, ctx: commands.Context, all_guilds: bool = False):
|
|
"""
|
|
Lists all current tasks and their triggers.
|
|
|
|
Do `[p]fifo list True` to see tasks from all guilds
|
|
"""
|
|
if all_guilds:
|
|
pass
|
|
else:
|
|
pass # TODO: parse and display tasks
|
|
|
|
@fifo.command(name="add")
|
|
async def fifo_add(self, ctx: commands.Context, task_name: str, *, command_to_execute: str):
|
|
"""
|
|
Add a new task to this guild's task list
|
|
"""
|
|
if (await self.config.guild(ctx.guild).tasks.get_raw(task_name, default=None)) is not None:
|
|
await ctx.maybe_send_embed(f"Task already exists with {task_name=}")
|
|
return
|
|
|
|
if not await self._check_parsable_command(ctx, command_to_execute):
|
|
await ctx.maybe_send_embed("Failed to parse command. Make sure to include the prefix")
|
|
return
|
|
|
|
task = Task(task_name, ctx.guild.id, self.config, ctx.author.id)
|
|
await task.set_commmand_str(command_to_execute)
|
|
await task.save_all()
|
|
await ctx.tick()
|
|
|
|
@fifo.command(name="delete")
|
|
async def fifo_delete(self, ctx: commands.Context, task_name: str):
|
|
"""
|
|
Deletes a task from this guild's task list
|
|
"""
|
|
pass
|
|
|
|
@fifo.group(name="trigger")
|
|
async def fifo_trigger(self, ctx: commands.Context):
|
|
"""
|
|
Add a new trigger for a task from the current guild.
|
|
"""
|
|
if ctx.invoked_subcommand is None:
|
|
pass
|
|
|
|
@fifo_trigger.command(name="interval")
|
|
async def fifo_trigger_interval(
|
|
self, ctx: commands.Context, task_name: str, interval_str: TimedeltaConverter
|
|
):
|
|
"""
|
|
Add an interval trigger to the specified task
|
|
"""
|
|
|
|
task = Task(task_name, ctx.guild.id, self.config)
|
|
await task.load_from_config()
|
|
|
|
if task.data is None:
|
|
await ctx.maybe_send_embed(
|
|
f"Task by the name of {task_name} is not found in this guild"
|
|
)
|
|
return
|
|
|
|
result = await task.add_trigger("interval", interval_str)
|
|
if not result:
|
|
await ctx.maybe_send_embed(
|
|
"Failed to add an interval trigger to this task, see console for logs"
|
|
)
|
|
return
|
|
await task.save_data()
|
|
await ctx.tick()
|
|
|
|
@fifo_trigger.command(name="date")
|
|
async def fifo_trigger_date(
|
|
self, ctx: commands.Context, task_name: str, datetime_str: DatetimeConverter
|
|
):
|
|
"""
|
|
Add a "run once" datetime trigger to the specified task
|
|
"""
|
|
|
|
task = Task(task_name, ctx.guild.id, self.config)
|
|
await task.load_from_config()
|
|
|
|
if task.data is None:
|
|
await ctx.maybe_send_embed(
|
|
f"Task by the name of {task_name} is not found in this guild"
|
|
)
|
|
return
|
|
|
|
result = await task.add_trigger("date", datetime_str)
|
|
if not result:
|
|
await ctx.maybe_send_embed(
|
|
"Failed to add a date trigger to this task, see console for logs"
|
|
)
|
|
return
|
|
|
|
await task.save_data()
|
|
await ctx.tick()
|
|
|
|
@fifo_trigger.command(name="cron")
|
|
async def fifo_trigger_cron(
|
|
self, ctx: commands.Context, task_name: str, cron_settings: DictConverter
|
|
):
|
|
"""
|
|
Add a "time of day" trigger to the specified task
|
|
"""
|
|
await ctx.maybe_send_embed("Not yet implemented")
|
|
|
|
# async def load_tasks(self):
|
|
# """
|
|
# Run once on cog load.
|
|
# """
|
|
# all_guilds = await self.config.all_guilds()
|
|
# async for guild_id, guild_data in AsyncIter(all_guilds["tasks"].items(), steps=100):
|
|
# for task_name, task_data in guild_data["tasks"].items():
|
|
# task = Task(task_name, guild_id, self.config)
|
|
# await task.load_from_data(task_data)
|
|
#
|
|
# job = self.scheduler.add_job(
|
|
# task.execute, id=task_name + "_" + guild_id, trigger=await task.get_trigger(),
|
|
# )
|
|
#
|
|
# self.scheduler.start()
|
|
|
|
# async def parent_loop(self):
|
|
# await asyncio.sleep(60)
|
|
# asyncio.create_task(self.process_tasks(datetime.datetime.utcnow()))
|
|
#
|
|
# async def process_tasks(self, now: datetime.datetime):
|
|
# for task in self.tasks:
|
|
# pass
|