Discord/classes/holo_user.py

148 lines
4.5 KiB
Python
Raw Normal View History

2023-05-06 18:08:52 +03:00
import logging
from typing import Any, Union
2023-05-04 17:09:47 +03:00
import discord
import discord.member
2024-12-16 00:36:48 +02:00
from libbot import config_get
2023-05-04 17:09:47 +03:00
from modules.database import col_users, col_warnings
2023-05-04 17:09:47 +03:00
2023-05-06 18:08:52 +03:00
logger = logging.getLogger(__name__)
2023-05-04 17:14:23 +03:00
2023-05-04 17:09:47 +03:00
class UserNotFoundError(Exception):
"""HoloUser could not find user with such an ID in database"""
2023-05-04 17:14:23 +03:00
2023-05-04 17:09:47 +03:00
def __init__(self, user, user_id):
self.user = user
self.user_id = user_id
2024-12-16 00:36:48 +02:00
2023-05-04 17:14:23 +03:00
super().__init__(
f"User of type {type(self.user)} with id {self.user_id} was not found"
)
2023-05-04 17:09:47 +03:00
2023-05-04 17:14:23 +03:00
class HoloUser:
def __init__(
self, user: Union[discord.User, discord.Member, discord.member.Member, int]
) -> None:
2023-05-04 17:09:47 +03:00
"""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
2023-05-04 17:14:23 +03:00
"""
2023-05-04 17:09:47 +03:00
if hasattr(user, "id"):
2023-05-04 17:14:23 +03:00
self.id = user.id # type: ignore
2023-05-04 17:09:47 +03:00
else:
self.id = user
jav_user = 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
2023-05-04 17:14:23 +03:00
"""
2023-05-04 17:09:47 +03:00
warns = col_warnings.find_one({"user": self.id})
2024-12-16 00:36:48 +02:00
return 0 if warns is None else warns["warns"]
2023-05-04 17:09:47 +03:00
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.
2023-05-04 17:14:23 +03:00
"""
2023-05-04 17:09:47 +03:00
warns = col_warnings.find_one({"user": self.id})
2024-12-16 00:36:48 +02:00
if warns is not None:
2023-05-04 17:14:23 +03:00
col_warnings.update_one(
filter={"_id": self.db_id},
update={"$set": {"warns": warns["warns"] + count}},
)
2023-05-04 17:09:47 +03:00
else:
2023-05-04 17:14:23 +03:00
col_warnings.insert_one(document={"user": self.id, "warns": count})
2024-12-16 00:36:48 +02:00
2023-05-06 18:08:52 +03:00
logger.info(f"User {self.id} was warned {count} times due to: {reason}")
2023-05-04 17:09:47 +03:00
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
2023-05-04 17:14:23 +03:00
"""
2023-05-04 17:09:47 +03:00
if not hasattr(self, key):
raise AttributeError()
2024-12-16 00:36:48 +02:00
2023-05-04 17:09:47 +03:00
setattr(self, key, value)
2023-05-04 17:14:23 +03:00
col_users.update_one(
filter={"_id": self.db_id}, update={"$set": {key: value}}, upsert=True
)
2024-12-16 00:36:48 +02:00
2023-05-06 18:08:52 +03:00
logger.info(f"Set attribute {key} of user {self.id} to {value}")
2023-05-04 17:09:47 +03:00
2023-05-08 16:45:00 +03:00
async def is_moderator(
self, 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
2024-12-16 00:36:48 +02:00
2023-05-08 16:48:15 +03:00
moderator_role = await config_get("moderators", "roles")
2023-05-08 16:45:00 +03:00
council_role = await config_get("council", "roles")
2024-12-16 00:36:48 +02:00
2023-05-08 16:45:00 +03:00
for role in member.roles:
if role.id == moderator_role or role.id == council_role:
return True
2024-12-16 00:36:48 +02:00
2023-05-08 16:45:00 +03:00
return False
async def is_council(
self, 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
2024-12-16 00:36:48 +02:00
2023-05-08 16:45:00 +03:00
council_role = await config_get("council", "roles")
2024-12-16 00:36:48 +02:00
2023-05-08 16:45:00 +03:00
for role in member.roles:
if role.id == council_role:
return True
2024-12-16 00:36:48 +02:00
2023-05-08 16:45:00 +03:00
return False
2023-05-06 18:08:52 +03:00
# 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()