Rosie-Val/SeaCogs
1
0
Fork 0
Code Pull requests Activity

fix(aurora): fixing up some model stuff

Browse source
This commit is contained in:
Seaswimmer 2024-05-04 18:05:20 -04:00
parent e46612dc08
commit 39479eb216
Signed by untrusted user: cswimr
GPG key ID: 5D671B5D03D65A7F
2 changed files with 32 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
aurora/models.py
View file

@ -10,29 +10,29 @@ from aurora.utilities.utils import generate_dict
class AuroraBaseModel(BaseModel): class AuroraBaseModel(BaseModel):
"""Base class for all models in Aurora.""" """Base class for all models in Aurora."""
bot: ClassVar[Red]
guild_id: int
def to_json(self, indent: int = None, file: Any = None, **kwargs): def to_json(self, indent: int = None, file: Any = None, **kwargs):
from aurora.utilities.json import dump, dumps 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): class Moderation(AuroraBaseModel):
bot: ClassVar[Red]
moderation_id: int moderation_id: int
guild_id: int
timestamp: datetime timestamp: datetime
moderation_type: str moderation_type: str
target_type: str target_type: str
target_id: int target_id: int
moderator_id: int moderator_id: int
role_id: Optional[int] role_id: Optional[int] = None
duration: Optional[timedelta] duration: Optional[timedelta] = None
end_timestamp: Optional[datetime] end_timestamp: Optional[datetime] = None
reason: Optional[str] reason: Optional[str] = None
resolved: bool resolved: bool
resolved_by: Optional[int] resolved_by: Optional[int] = None
resolve_reason: Optional[str] resolve_reason: Optional[str] = None
expired: bool expired: bool
changes: List[Dict] changes: List["Change"]
metadata: Dict metadata: Dict
@property @property
@ -76,7 +76,7 @@ class Moderation(AuroraBaseModel):
result = cursor.fetchone() result = cursor.fetchone()
if result: if result:
case = generate_dict(result, guild_id) case = generate_dict(bot, result, guild_id)
cursor.close() cursor.close()
return cls.from_dict(bot, case) return cls.from_dict(bot, case)
@ -86,25 +86,23 @@ class Moderation(AuroraBaseModel):
def from_dict(cls, bot: Red, data: dict) -> "Moderation": def from_dict(cls, bot: Red, data: dict) -> "Moderation":
return cls(bot=bot, **data) 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): class Change(AuroraBaseModel):
type: Literal["ORIGINAL", "RESOLVE", "EDIT"] type: Literal["ORIGINAL", "RESOLVE", "EDIT"]
timestamp: datetime timestamp: datetime
reason: str reason: str
user_id: int user_id: int
duration: Optional[timedelta] duration: Optional[timedelta] = None
end_timestamp: Optional[datetime] end_timestamp: Optional[datetime] = None
def __str__(self): def __str__(self):
return f"{self.type} {self.user_id} {self.reason}" 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 @classmethod
def from_dict(cls, data: dict) -> "Change": def from_dict(cls, bot: Red, data: dict) -> "Change":
return cls(**data) return cls(bot=bot, **data)
class PartialUser(AuroraBaseModel): class PartialUser(AuroraBaseModel):
id: int id: int
@ -124,9 +122,9 @@ class PartialUser(AuroraBaseModel):
if not user: if not user:
try: try:
user = await bot.fetch_user(user_id) 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: 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): class PartialChannel(AuroraBaseModel):
@ -144,19 +142,18 @@ class PartialChannel(AuroraBaseModel):
@classmethod @classmethod
async def from_id(cls, bot: Red, channel_id: int) -> "PartialChannel": async def from_id(cls, bot: Red, channel_id: int) -> "PartialChannel":
user = bot.get_channel(channel_id) channel = bot.get_channel(channel_id)
if not user: if not channel:
try: try:
user = await bot.fetch_channel(channel_id) channel = await bot.fetch_channel(channel_id)
return cls(id=user.id, username=user.name, discriminator=user.discriminator) return cls(bot=bot, guild_id=channel.guild.id, id=channel.id, username=channel.name)
except (NotFound, InvalidData, HTTPException, Forbidden) as e: except (NotFound, InvalidData, HTTPException, Forbidden) as e:
if e == Forbidden: if e == Forbidden:
return cls(id=channel_id, name="Forbidden Channel") return cls(bot=bot, guild_id=0, 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="Deleted Channel")
class PartialRole(AuroraBaseModel): class PartialRole(AuroraBaseModel):
id: int id: int
guild_id: int
name: str name: str
@property @property
@ -171,8 +168,8 @@ class PartialRole(AuroraBaseModel):
try: try:
guild = await bot.fetch_guild(guild_id, with_counts=False) guild = await bot.fetch_guild(guild_id, with_counts=False)
except (Forbidden, HTTPException): 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) role = guild.get_role(role_id)
if not role: if not role:
return cls(id=role_id, guild_id=guild_id, name="Deleted Role") return cls(bot=bot, guild_id=guild_id, id=role_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=role.name)

5
aurora/utilities/utils.py
View file

@ -7,6 +7,7 @@ from dateutil.relativedelta import relativedelta as rd
from discord import File, Guild, Interaction, Member, SelectOption, User from discord import File, Guild, Interaction, Member, SelectOption, User
from discord.errors import Forbidden, NotFound from discord.errors import Forbidden, NotFound
from redbot.core import commands, data_manager from redbot.core import commands, data_manager
from redbot.core.bot import Red
from redbot.core.utils.chat_formatting import error from redbot.core.utils.chat_formatting import error
from aurora.utilities.config import config 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 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 from aurora.models import Change
if result[7] is not None: if result[7] is not None:
hours, minutes, seconds = map(int, result[7].split(':')) 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('\\"', '"')) changes = json.loads(result[14].strip('"').replace('\\"', '"'))
change_obj_list = [] change_obj_list = []
for change in changes: for change in changes:
change_obj_list.append(Change(**change)) change_obj_list.append(Change.from_dict(bot=bot, **change))
case = { case = {
"moderation_id": int(result[0]), "moderation_id": int(result[0]),