Discord/classes/holo_user.py
2024-12-16 16:25:35 +01:00

151 lines
4.5 KiB
Python

import logging
from typing import Any, Union
import discord
import discord.member
from libbot import config_get
from modules.database import col_warnings, sync_col_users, sync_col_warnings, col_users
logger = logging.getLogger(__name__)
class UserNotFoundError(Exception):
"""HoloUser could not find user with such an ID in database"""
def __init__(self, user, user_id):
self.user = user
self.user_id = user_id
super().__init__(
f"User of type {type(self.user)} with id {self.user_id} was not found"
)
class HoloUser:
def __init__(
self, user: Union[discord.User, discord.Member, discord.member.Member, int]
) -> None:
"""Get an object that has a proper binding between Discord ID and database
### Args:
* `user` (Union[discord.User, discord.Member, discord.member.Member, int]): Object from which ID can be extracted
### Raises:
* `UserNotFoundError`: User with such ID does not seem to exist in database
"""
if hasattr(user, "id"):
self.id = user.id # type: ignore
else:
self.id = user
jav_user = sync_col_users.find_one({"user": self.id})
if jav_user is None:
raise UserNotFoundError(user=user, user_id=self.id)
self.db_id = jav_user["_id"]
self.customrole = jav_user["customrole"]
self.customchannel = jav_user["customchannel"]
self.warnings = self.warns()
def warns(self) -> int:
"""Get number of warnings user has
### Returns:
* `int`: Number of warnings
"""
warns = sync_col_warnings.find_one({"user": self.id})
return 0 if warns is None else warns["warns"]
async def warn(self, count=1, reason: str = "Not provided") -> None:
"""Warn and add count to warns number
### Args:
* `count` (int, optional): Count of warnings to be added. Defaults to 1.
"""
warns = await col_warnings.find_one({"user": self.id})
if warns is not None:
await col_warnings.update_one(
{"_id": self.db_id},
{"$set": {"warns": warns["warns"] + count}},
)
else:
await col_warnings.insert_one(document={"user": self.id, "warns": count})
logger.info(f"User {self.id} was warned {count} times due to: {reason}")
async def set(self, key: str, value: Any) -> None:
"""Set attribute data and save it into database
### Args:
* `key` (str): Attribute to be changed
* `value` (Any): Value to set
"""
if not hasattr(self, key):
raise AttributeError()
setattr(self, key, value)
await col_users.update_one(
{"_id": self.db_id}, {"$set": {key: value}}, upsert=True
)
logger.info(f"Set attribute {key} of user {self.id} to {value}")
@staticmethod
async def is_moderator(
member: Union[discord.User, discord.Member, discord.member.Member]
) -> bool:
"""Check if user is moderator or council member
### Args:
* `member` (Union[discord.User, discord.Member, discord.member.Member]): Member object
### Returns:
`bool`: `True` if member is a moderator or member of council and `False` if not
"""
if isinstance(member, discord.User):
return False
moderator_role = await config_get("moderators", "roles")
council_role = await config_get("council", "roles")
for role in member.roles:
if role.id == moderator_role or role.id == council_role:
return True
return False
@staticmethod
async def is_council(
member: Union[discord.User, discord.Member, discord.member.Member]
) -> bool:
"""Check if user is a member of council
### Args:
* `member` (Union[discord.User, discord.Member, discord.member.Member]): Member object
### Returns:
`bool`: `True` if member is a member of council and `False` if not
"""
if isinstance(member, discord.User):
return False
council_role = await config_get("council", "roles")
for role in member.roles:
if role.id == council_role:
return True
return False
# def purge(self) -> None:
# """Completely remove data from database. Will not remove transactions logs and warnings."""
# col_users.delete_one(filter={"_id": self.db_id})
# self.unauthorize()