SeaCogs/aurora/models/base.py

32 lines
1.4 KiB
Python
Raw Normal View History

2024-06-05 23:13:23 -04:00
from typing import Any, Optional
2024-05-06 17:23:59 -04:00
2024-06-05 23:13:23 -04:00
from discord import Guild
2024-05-06 17:23:59 -04:00
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"})
2024-06-11 03:17:38 -04:00
def to_json(self, indent: int | None = None, file: Any | None = None, **kwargs) -> str:
2024-05-06 21:46:01 -04:00
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)
2024-05-06 17:23:59 -04:00
class AuroraGuildModel(AuroraBaseModel):
2024-06-05 23:13:23 -04:00
"""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)
2024-05-06 17:23:59 -04:00
guild_id: int
2024-06-05 23:13:23 -04:00
guild: Optional[Guild] = None
2024-05-06 17:23:59 -04:00
def dump(self) -> dict:
2024-06-05 23:13:23 -04:00
return self.model_dump(exclude={"bot", "guild_id", "guild"})
2024-06-11 03:17:38 -04:00
def to_json(self, indent: int | None = None, file: Any | None = None, **kwargs) -> str:
2024-05-06 21:46:01 -04:00
from ..utilities.json import dump, dumps # pylint: disable=cyclic-import
2024-06-05 23:13:23 -04:00
return dump(self.dump(), file, indent=indent, **kwargs) if file else dumps(self.dump(), indent=indent, **kwargs)