25 lines
1.1 KiB
Python
25 lines
1.1 KiB
Python
|
from typing import Any
|
||
|
|
|
||
|
|
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 to_json(self, indent: int = None, file: Any = None, **kwargs):
|
||
|
|
from aurora.utilities.json import ( # pylint: disable=cyclic-import
|
||
|
|
dump, dumps)
|
||
|
|
return dump(self.model_dump(exclude={"bot"}), 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 modified to_json() method to match."""
|
||
|
|
guild_id: int
|
||
|
|
|
||
|
|
def to_json(self, indent: int = None, file: Any = None, **kwargs):
|
||
|
|
from aurora.utilities.json import ( # pylint: disable=cyclic-import
|
||
|
|
dump, dumps)
|
||
|
|
return dump(self.model_dump(exclude={"bot", "guild_id"}), file, indent=indent, **kwargs) if file else dumps(self.model_dump(exclude={"bot", "guild_id"}), indent=indent, **kwargs)
|