SeaCogs/aurora/models/base.py
cswimr fa27d12de5
Some checks failed
Actions / Build Documentation (MkDocs) (pull_request) Successful in 29s
Actions / Lint Code (Ruff & Pylint) (pull_request) Failing after 51s
feat(aurora): make all models compatible with repr
2024-08-22 15:33:36 -04:00

30 lines
1.1 KiB
Python

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) -> str:
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."""
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 __repr__(self) -> str:
return f"<{self.__class__.__name__} guild_id={self.guild_id}>"