diff --git a/maintenance/__init__.py b/maintenance/__init__.py index 4811660..4f8010c 100644 --- a/maintenance/__init__.py +++ b/maintenance/__init__.py @@ -10,3 +10,4 @@ async def setup(bot): cog = Maintenance(bot) await bot.add_cog(cog) + await cog.initialize() diff --git a/maintenance/classes.py b/maintenance/classes.py index 0ca2a5d..32ceb7d 100644 --- a/maintenance/classes.py +++ b/maintenance/classes.py @@ -6,7 +6,9 @@ def __init__(self, start=0, end=0, after=True, whitelist=None): if not whitelist: whitelist = [] self.start = start + time.time() - if end: + if end is True: + self.end = True + elif end: if after: self.end = end + self.start else: diff --git a/maintenance/converters.py b/maintenance/converters.py index 62642ee..f1e095b 100644 --- a/maintenance/converters.py +++ b/maintenance/converters.py @@ -23,6 +23,7 @@ async def convert(self, ctx, argument): _end = parser.add_mutually_exclusive_group() _end.add_argument("--end-after", nargs="*", dest="end", default=[]) _end.add_argument("--end-in", nargs="*", dest="endin", default=[]) + _end.add_argument("--end-after-restart", dest="end_after_restart", action="store_true") try: vals = vars(parser.parse_args(argument.split(" "))) except Exception as exc: @@ -35,6 +36,8 @@ async def convert(self, ctx, argument): if end_seconds == None: end_seconds = convert_time(vals.get("endin", None)) after = False + if vals.get("end_after_restart", False): + end_seconds = True if start_seconds: if end_seconds: scheduled = ScheduledMaintenance( diff --git a/maintenance/maintenance.py b/maintenance/maintenance.py index c159436..d11108e 100644 --- a/maintenance/maintenance.py +++ b/maintenance/maintenance.py @@ -58,10 +58,21 @@ def __init__(self, bot): "delete": 3, "scheduledmaintenance": [], } - # When on maintenance, on will be set to [True, second of when it is off maintenance, list of people who can bypass the maintenance] + # When on maintenance, on will be set to [True, second of when it is off maintenance or True if it's planned to end after bot restart, list of people who can bypass the maintenance] self.conf.register_global(**default_global) self.bot.add_check(self.this_check) - self.task = self.bot.loop.create_task(self.bg_loop()) + self.task = None + + async def cog_load(self): + # disable maintenance scheduled to end after restart, + # if the cog was loaded during bot startup + if not self.bot.is_ready(): + on = await self.conf.on() + if on[0] and on[1] is True: + await self.conf.on.set([False, 0, []]) + + async def initialize(self): + self.task = asyncio.create_task(self.bg_loop()) @listener() async def on_command_error(self, ctx, error): @@ -73,11 +84,9 @@ async def on_command_error(self, ctx, error): await ctx.send(error) def cog_unload(self): - self.__unload() - - def __unload(self): self.bot.remove_check(self.this_check) - self.task.cancel() + if self.task is not None: + self.task.cancel() async def red_delete_data_for_user( self, @@ -121,7 +130,10 @@ async def this_check(self, ctx): on = await self.conf.on() if not on[0]: return True - if on[1] <= time.time() and on[1] != 0: + if on[1] is True: + # maintenance will be disabled after restart + pass + elif on[1] != 0 and on[1] <= time.time(): setting = [False, 0, []] await self.conf.on.set(setting) return True @@ -145,6 +157,7 @@ async def _on(self, ctx, *, args: Margs = None): --start-in: Makes the maintenace start in that long. --end-in: Schedules the maintenance to end in that long from the current second. --end-after: Schedules the maintenance to end in that long after the maitenance has started. + --end-after-restart: Schedules the maintenance to end after bot restart. --whitelist: Provide user IDs after this to whitelist people from the maintenance. Examples: @@ -165,8 +178,8 @@ async def _on(self, ctx, *, args: Margs = None): else: num = 0 whitelist = [] - start = time.time() - if start == time.time(): + start = None + if start is None: setting = [True, num, whitelist] await self.conf.on.set(setting) else: @@ -195,7 +208,9 @@ async def settings(self, ctx): sending += " • " + starting sending += "```" return await ctx.send(sending) - if on[1] != 0: + if on[1] is True: + done = "after bot restart" + elif on[1] != 0: done = str(datetime.fromtimestamp(on[1]).strftime("%A, %B %d, %Y %I:%M:%S")) done = "on " + done else: