From 1b5d993ff5f8d1816264e3e12dceb68d00dc4f72 Mon Sep 17 00:00:00 2001 From: fourjr <28086837+fourjr@users.noreply.github.com> Date: Tue, 10 Nov 2020 21:31:56 +0800 Subject: [PATCH] Use enums in config, resolve #2821 --- bot.py | 6 +++--- cogs/modmail.py | 20 ++++++++++---------- core/config.py | 43 ++++++++++++++++++------------------------- core/models.py | 6 ++++++ core/thread.py | 2 +- 5 files changed, 38 insertions(+), 39 deletions(-) diff --git a/bot.py b/bot.py index 901999f2e9..08782084f1 100644 --- a/bot.py +++ b/bot.py @@ -34,7 +34,7 @@ from core.clients import ApiClient, PluginDatabaseClient, MongoDBClient from core.config import ConfigManager from core.utils import human_join, match_title, normalize_alias -from core.models import PermissionLevel, SafeFormatter, getLogger, configure_logging +from core.models import DMDisabled, PermissionLevel, SafeFormatter, getLogger, configure_logging from core.thread import ThreadManager from core.time import human_timedelta @@ -770,7 +770,7 @@ async def process_dm_modmail(self, message: discord.Message) -> None: ) return - if self.config["dm_disabled"] >= 1: + if self.config["dm_disabled"] in (DMDisabled.NEW_THREADS, DMDisabled.ALL_THREADS): embed = discord.Embed( title=self.config["disabled_new_thread_title"], color=self.error_color, @@ -787,7 +787,7 @@ async def process_dm_modmail(self, message: discord.Message) -> None: thread = await self.threads.create(message.author, message=message) else: - if self.config["dm_disabled"] == 2: + if self.config["dm_disabled"] == DMDisabled.ALL_THREADS: embed = discord.Embed( title=self.config["disabled_current_thread_title"], color=self.error_color, diff --git a/cogs/modmail.py b/cogs/modmail.py index 01bdb4d341..638ae2e28d 100644 --- a/cogs/modmail.py +++ b/cogs/modmail.py @@ -14,7 +14,7 @@ from natural.date import duration from core import checks -from core.models import PermissionLevel, SimilarCategoryConverter, getLogger +from core.models import DMDisabled, PermissionLevel, SimilarCategoryConverter, getLogger from core.paginator import EmbedPaginatorSession from core.thread import Thread from core.time import UserFriendlyTime, human_timedelta @@ -964,7 +964,7 @@ async def contact( else: thread = await self.bot.threads.create(user, creator=ctx.author, category=category) - if self.bot.config["dm_disabled"] >= 1: + if self.bot.config["dm_disabled"] in (DMDisabled.NEW_THREADS, DMDisabled.ALL_THREADS): logger.info("Contacting user %s when Modmail DM is disabled.", user) if not silent: @@ -1450,8 +1450,8 @@ async def enable(self, ctx): color=self.bot.main_color, ) - if self.bot.config["dm_disabled"] != 0: - self.bot.config["dm_disabled"] = 0 + if self.bot.config["dm_disabled"] != DMDisabled.NONE: + self.bot.config["dm_disabled"] = DMDisabled.NONE await self.bot.config.update() return await ctx.send(embed=embed) @@ -1481,8 +1481,8 @@ async def disable_new(self, ctx): description="Modmail will not create any new threads.", color=self.bot.main_color, ) - if self.bot.config["dm_disabled"] < 1: - self.bot.config["dm_disabled"] = 1 + if self.bot.config["dm_disabled"] < DMDisabled.NEW_THREADS: + self.bot.config["dm_disabled"] = DMDisabled.NEW_THREADS await self.bot.config.update() return await ctx.send(embed=embed) @@ -1501,8 +1501,8 @@ async def disable_all(self, ctx): color=self.bot.main_color, ) - if self.bot.config["dm_disabled"] != 2: - self.bot.config["dm_disabled"] = 2 + if self.bot.config["dm_disabled"] != DMDisabled.ALL_THREADS: + self.bot.config["dm_disabled"] = DMDisabled.ALL_THREADS await self.bot.config.update() return await ctx.send(embed=embed) @@ -1514,13 +1514,13 @@ async def isenable(self, ctx): Check if the DM functionalities of Modmail is enabled. """ - if self.bot.config["dm_disabled"] == 1: + if self.bot.config["dm_disabled"] == DMDisabled.NEW_THREADS: embed = discord.Embed( title="New Threads Disabled", description="Modmail is not creating new threads.", color=self.bot.error_color, ) - elif self.bot.config["dm_disabled"] == 2: + elif self.bot.config["dm_disabled"] == DMDisabled.ALL_THREADS: embed = discord.Embed( title="All DM Disabled", description="Modmail is not accepting any DM messages for new and existing threads.", diff --git a/core/config.py b/core/config.py index 66bb7a69ba..be53807f3f 100644 --- a/core/config.py +++ b/core/config.py @@ -12,7 +12,7 @@ from discord.ext.commands import BadArgument from core._color_data import ALL_COLORS -from core.models import InvalidConfigError, Default, getLogger +from core.models import DMDisabled, InvalidConfigError, Default, getLogger from core.time import UserFriendlyTimeSync from core.utils import strtobool @@ -98,9 +98,7 @@ class ConfigManager: "activity_message": "", "activity_type": None, "status": None, - # dm_disabled 0 = none, 1 = new threads, 2 = all threads - # TODO: use enum - "dm_disabled": 0, + "dm_disabled": DMDisabled.NONE, "oauth_whitelist": [], # moderation "blocked": {}, @@ -165,7 +163,11 @@ class ConfigManager: "enable_eval", } - special_types = {"status", "activity_type"} + enums = { + "dm_disabled": DMDisabled, + "status": discord.Status, + "activity_type": discord.ActivityType + } defaults = {**public_keys, **private_keys, **protected_keys} all_keys = set(defaults.keys()) @@ -277,26 +279,15 @@ def get(self, key: str, convert=True) -> typing.Any: value = strtobool(value) except ValueError: value = self.remove(key) - - elif key in self.special_types: + + elif key in self.enums: if value is None: return None - - if key == "status": - try: - # noinspection PyArgumentList - value = discord.Status(value) - except ValueError: - logger.warning("Invalid status %s.", value) - value = self.remove(key) - - elif key == "activity_type": - try: - # noinspection PyArgumentList - value = discord.ActivityType(value) - except ValueError: - logger.warning("Invalid activity %s.", value) - value = self.remove(key) + try: + value = self.enums[key](value) + except ValueError: + logger.warning("Invalid %s %s.", key, value) + value = self.remove(key) return value @@ -355,8 +346,10 @@ def set(self, key: str, item: typing.Any, convert=True) -> None: except ValueError: raise InvalidConfigError("Must be a yes/no value.") - # elif key in self.special_types: - # if key == "status": + elif key in self.enums: + if isinstance(item, self.enums[key]): + # value is an enum type + item = item.value return self.__setitem__(key, item) diff --git a/core/models.py b/core/models.py index 4bd62e5d90..8d1be4130a 100644 --- a/core/models.py +++ b/core/models.py @@ -257,3 +257,9 @@ async def publish(self): async def ack(self): return + + +class DMDisabled(IntEnum): + NONE = 0 + NEW_THREADS = 1 + ALL_THREADS = 2 diff --git a/core/thread.py b/core/thread.py index eca0e19220..4beacf31bc 100644 --- a/core/thread.py +++ b/core/thread.py @@ -896,7 +896,7 @@ async def send( except Exception as e: logger.warning("Cannot delete message: %s.", e) - if from_mod and self.bot.config["dm_disabled"] == 2 and destination != self.channel: + if from_mod and self.bot.config["dm_disabled"] == DMDisabled.ALL_THREADS and destination != self.channel: logger.info("Sending a message to %s when DM disabled is set.", self.recipient) try: