SeaCogs/aurora/models/partials.py

from discord import ChannelType, Forbidden, Guild, HTTPException, InvalidData, NotFound
from redbot.core.bot import Red

from .base import AuroraBaseModel, AuroraGuildModel


class PartialUser(AuroraBaseModel):
    id: int
    username: str
    discriminator: int

    @property
    def name(self):
        return f"{self.username}#{self.discriminator}" if self.discriminator != 0 else self.username

    def __str__(self):
        return self.name

    @classmethod
    async def from_id(cls, bot: Red, user_id: int) -> "PartialUser":
        user = bot.get_user(user_id)
        if not user:
            try:
                user = await bot.fetch_user(user_id)
                return cls(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator)
            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)


class PartialChannel(AuroraGuildModel):
    id: int
    name: str
    type: ChannelType

    @property
    def mention(self):
        if self.name in ["Deleted Channel", "Forbidden Channel"]:
            return self.name
        return f"<#{self.id}>"

    def __str__(self):
        return self.mention

    @classmethod
    async def from_id(cls, bot: Red, channel_id: int, guild: Guild) -> "PartialChannel":
        channel = bot.get_channel(channel_id)
        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)
            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)

class PartialRole(AuroraGuildModel):
    id: int
    name: str

    @property
    def mention(self):
        if self.name in ["Deleted Role", "Forbidden Role"]:
            return self.name
        return f"<@&{self.id}>"

    def __str__(self):
        return self.mention

    @classmethod
    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)