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.dump(), 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)