SeaCogs/aurora/models/base.py

from typing import Any, Optional

from discord import Guild
from pydantic import BaseModel, ConfigDict
from redbot.core.bot import Red


class AuroraBaseModel(BaseModel):
    """Base class for all models in Aurora."""
    model_config = ConfigDict(ignored_types=(Red,), arbitrary_types_allowed=True)
    bot: Red

    def dump(self) -> dict:
        return self.model_dump(exclude={"bot"})

    def to_json(self, indent: int | None = None, file: Any | None = None, **kwargs):
        from ..utilities.json import dump, dumps  # pylint: disable=cyclic-import
        return dump(self.dump(), file, indent=indent, **kwargs) if file else dumps(self.model_dump(exclude={"bot"}), indent=indent, **kwargs)

class AuroraGuildModel(AuroraBaseModel):
    """Subclass of AuroraBaseModel that includes a guild_id attribute and a guild attribute, and a modified to_json() method to match."""
    model_config = ConfigDict(ignored_types=(Red, Guild), arbitrary_types_allowed=True)
    guild_id: int
    guild: Optional[Guild] = None

    def dump(self) -> dict:
        return self.model_dump(exclude={"bot", "guild_id", "guild"})

    def to_json(self, indent: int | None = None, file: Any | None = None, **kwargs):
        from ..utilities.json import dump, dumps  # pylint: disable=cyclic-import
        return dump(self.dump(), file, indent=indent, **kwargs) if file else dumps(self.dump(), indent=indent, **kwargs)