diff --git a/maintenance/__init__.py b/maintenance/__init__.py index e761b33..443677e 100644 --- a/maintenance/__init__.py +++ b/maintenance/__init__.py @@ -10,3 +10,4 @@ async def setup(bot): cog = Maintenance(bot) 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 75e5b9d..20ed9bb 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 c3b00a5..fa884d7 100644 --- a/maintenance/maintenance.py +++ b/maintenance/maintenance.py @@ -34,9 +34,17 @@ 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) + + async def initialize(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, []]) self.task = self.bot.loop.create_task(self.bg_loop()) @listener() @@ -97,7 +105,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 @@ -121,6 +132,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: @@ -140,8 +152,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: @@ -170,7 +182,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: