diff --git a/aurora/models/partials.py b/aurora/models/partials.py index 68b0261..05c3cc8 100644 --- a/aurora/models/partials.py +++ b/aurora/models/partials.py @@ -1,4 +1,5 @@ -from discord import ChannelType, Forbidden, Guild, HTTPException, InvalidData, NotFound +from discord import ChannelType, Forbidden, Guild, HTTPException, InvalidData, NotFound, Role, User +from discord.abc import Messageable from redbot.core.bot import Red from .base import AuroraBaseModel, AuroraGuildModel @@ -8,6 +9,7 @@ class PartialUser(AuroraBaseModel): id: int username: str discriminator: int + _obj: User | None @property def name(self): @@ -22,16 +24,17 @@ class PartialUser(AuroraBaseModel): if not user: try: user = await bot.fetch_user(user_id) - return cls(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator) + return cls(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator, _obj=user) except NotFound: - return cls(bot=bot, id=user_id, username="Deleted User", discriminator=0) - return cls(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator) + return cls(bot=bot, id=user_id, username="Deleted User", discriminator=0, _obj=None) + return cls(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator, _obj=user) class PartialChannel(AuroraGuildModel): id: int name: str type: ChannelType + _obj: Messageable | None @property def mention(self): @@ -48,16 +51,17 @@ class PartialChannel(AuroraGuildModel): if not channel: try: channel = await bot.fetch_channel(channel_id) - return cls(bot=bot, guild_id=channel.guild.id, guild=guild, id=channel.id, name=channel.name, type=channel.type) + return cls(bot=bot, guild_id=channel.guild.id, guild=guild, id=channel.id, name=channel.name, type=channel.type, _obj=channel) except (NotFound, InvalidData, HTTPException, Forbidden) as e: if e == Forbidden: 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", type=ChannelType.text) - return cls(bot=bot, guild_id=channel.guild.id, guild=guild, id=channel.id, name=channel.name, type=channel.type) + return cls(bot=bot, guild_id=0, id=channel_id, name="Deleted Channel", type=ChannelType.text, _obj=None) + return cls(bot=bot, guild_id=channel.guild.id, guild=guild, id=channel.id, name=channel.name, type=channel.type, _obj=channel) class PartialRole(AuroraGuildModel): id: int name: str + _obj: Role | None @property def mention(self): @@ -72,5 +76,5 @@ class PartialRole(AuroraGuildModel): async def from_id(cls, bot: Red, guild: Guild, role_id: int) -> "PartialRole": role = guild.get_role(role_id) if not role: - 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) + return cls(bot=bot, guild_id=guild.id, id=role_id, name="Deleted Role", _obj=None) + return cls(bot=bot, guild_id=guild.id, id=role.id, name=role.name, _obj=role)