From 39479eb216bac3e9c0f9201ccd5ab6cffa95dbe2 Mon Sep 17 00:00:00 2001 From: SeaswimmerTheFsh Date: Sat, 4 May 2024 18:05:20 -0400 Subject: [PATCH] fix(aurora): fixing up some model stuff --- aurora/models.py | 61 +++++++++++++++++++-------------------- aurora/utilities/utils.py | 5 ++-- 2 files changed, 32 insertions(+), 34 deletions(-) diff --git a/aurora/models.py b/aurora/models.py index 6fd3862..c21ff4c 100644 --- a/aurora/models.py +++ b/aurora/models.py @@ -10,29 +10,29 @@ from aurora.utilities.utils import generate_dict class AuroraBaseModel(BaseModel): """Base class for all models in Aurora.""" + bot: ClassVar[Red] + guild_id: int def to_json(self, indent: int = None, file: Any = None, **kwargs): from aurora.utilities.json import dump, dumps - return dump(self.model_dump(), file, indent=indent, **kwargs) if file else dumps(self.model_dump(), indent=indent, **kwargs) + return dump(self.model_dump(exclude={"bot", "guild_id"}), file, indent=indent, **kwargs) if file else dumps(self.model_dump(exclude={"bot", "guild_id"}), indent=indent, **kwargs) class Moderation(AuroraBaseModel): - bot: ClassVar[Red] moderation_id: int - guild_id: int timestamp: datetime moderation_type: str target_type: str target_id: int moderator_id: int - role_id: Optional[int] - duration: Optional[timedelta] - end_timestamp: Optional[datetime] - reason: Optional[str] + role_id: Optional[int] = None + duration: Optional[timedelta] = None + end_timestamp: Optional[datetime] = None + reason: Optional[str] = None resolved: bool - resolved_by: Optional[int] - resolve_reason: Optional[str] + resolved_by: Optional[int] = None + resolve_reason: Optional[str] = None expired: bool - changes: List[Dict] + changes: List["Change"] metadata: Dict @property @@ -76,7 +76,7 @@ class Moderation(AuroraBaseModel): result = cursor.fetchone() if result: - case = generate_dict(result, guild_id) + case = generate_dict(bot, result, guild_id) cursor.close() return cls.from_dict(bot, case) @@ -86,25 +86,23 @@ class Moderation(AuroraBaseModel): def from_dict(cls, bot: Red, data: dict) -> "Moderation": return cls(bot=bot, **data) - def to_json(self, indent: int = None, file: Any = None, **kwargs): - from aurora.utilities.json import dump, dumps - return dump(self.model_dump(exclude={"bot", "guild_id"}), file, indent=indent, **kwargs) if file else dumps(self.model_dump(exclude={"bot", "guild_id"}), indent=indent, **kwargs) - - class Change(AuroraBaseModel): type: Literal["ORIGINAL", "RESOLVE", "EDIT"] timestamp: datetime reason: str user_id: int - duration: Optional[timedelta] - end_timestamp: Optional[datetime] + duration: Optional[timedelta] = None + end_timestamp: Optional[datetime] = None def __str__(self): return f"{self.type} {self.user_id} {self.reason}" + async def get_user(self) -> "PartialUser": + return await PartialUser.from_id(self.bot, self.user_id) + @classmethod - def from_dict(cls, data: dict) -> "Change": - return cls(**data) + def from_dict(cls, bot: Red, data: dict) -> "Change": + return cls(bot=bot, **data) class PartialUser(AuroraBaseModel): id: int @@ -124,9 +122,9 @@ class PartialUser(AuroraBaseModel): if not user: try: user = await bot.fetch_user(user_id) - return cls(id=user.id, username=user.name, discriminator=user.discriminator) + return cls(bot=bot, guild_id=0, id=user.id, username=user.name, discriminator=user.discriminator) except NotFound: - return cls(id=user_id, username="Deleted User", discriminator=0) + return cls(bot=bot, guild_id=0, id=user_id, username="Deleted User", discriminator=0) class PartialChannel(AuroraBaseModel): @@ -144,19 +142,18 @@ class PartialChannel(AuroraBaseModel): @classmethod async def from_id(cls, bot: Red, channel_id: int) -> "PartialChannel": - user = bot.get_channel(channel_id) - if not user: + channel = bot.get_channel(channel_id) + if not channel: try: - user = await bot.fetch_channel(channel_id) - return cls(id=user.id, username=user.name, discriminator=user.discriminator) + channel = await bot.fetch_channel(channel_id) + return cls(bot=bot, guild_id=channel.guild.id, id=channel.id, username=channel.name) except (NotFound, InvalidData, HTTPException, Forbidden) as e: if e == Forbidden: - return cls(id=channel_id, name="Forbidden Channel") - return cls(id=channel_id, name="Deleted Channel") + return cls(bot=bot, guild_id=0, id=channel_id, name="Forbidden Channel") + return cls(bot=bot, guild_id=0, id=channel_id, name="Deleted Channel") class PartialRole(AuroraBaseModel): id: int - guild_id: int name: str @property @@ -171,8 +168,8 @@ class PartialRole(AuroraBaseModel): try: guild = await bot.fetch_guild(guild_id, with_counts=False) except (Forbidden, HTTPException): - return cls(id=role_id, guild_id=guild_id, name="Forbidden Role") + return cls(bot=bot, guild_id=guild_id, id=role_id, name="Forbidden Role") role = guild.get_role(role_id) if not role: - return cls(id=role_id, guild_id=guild_id, name="Deleted Role") - return cls(id=role.id, guild_id=guild_id, name=role.name) + return cls(bot=bot, guild_id=guild_id, id=role_id, name="Deleted Role") + return cls(bot=bot, guild_id=guild_id, id=role.id, name=role.name) diff --git a/aurora/utilities/utils.py b/aurora/utilities/utils.py index 874ade3..f6e69bf 100644 --- a/aurora/utilities/utils.py +++ b/aurora/utilities/utils.py @@ -7,6 +7,7 @@ from dateutil.relativedelta import relativedelta as rd from discord import File, Guild, Interaction, Member, SelectOption, User from discord.errors import Forbidden, NotFound from redbot.core import commands, data_manager +from redbot.core.bot import Red from redbot.core.utils.chat_formatting import error from aurora.utilities.config import config @@ -124,7 +125,7 @@ async def get_next_case_number(guild_id: str, cursor=None) -> int: return (result[0] + 1) if result else 1 -def generate_dict(result: dict, guild_id: int) -> dict: +def generate_dict(bot: Red, result: dict, guild_id: int) -> dict: from aurora.models import Change if result[7] is not None: hours, minutes, seconds = map(int, result[7].split(':')) @@ -136,7 +137,7 @@ def generate_dict(result: dict, guild_id: int) -> dict: changes = json.loads(result[14].strip('"').replace('\\"', '"')) change_obj_list = [] for change in changes: - change_obj_list.append(Change(**change)) + change_obj_list.append(Change.from_dict(bot=bot, **change)) case = { "moderation_id": int(result[0]),