diff --git a/cogs/chat.py b/cogs/chat.py new file mode 100644 index 0000000..ee34e08 --- /dev/null +++ b/cogs/chat.py @@ -0,0 +1,63 @@ +from util.toru_rpg import ToruRpg +from util.error_handler import async_ignore_multiple_errors + +from discord import Member, Message +from discord.ext import commands +from discord.ext.commands import Context, errors + +from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError + + +class RpgChat(commands.Cog): + def __init__(self, client) -> None: + self.client = client + self.handler = ToruRpg() + + @commands.Cog.listener() + @async_ignore_multiple_errors([ServerSelectionTimeoutError, ConnectionFailure]) + async def on_member_join(self, member: Member): + self.handler.register(member.id, member.guild.id) + + @commands.Cog.listener() + @async_ignore_multiple_errors([ServerSelectionTimeoutError, ConnectionFailure]) + async def on_member_remove(self, member: Member): + self.handler.unregister(member.id, member.guild.id) + + @commands.Cog.listener() + @async_ignore_multiple_errors([ServerSelectionTimeoutError, ConnectionFailure]) + async def on_message(self, message: Message): + ctx = await self.client.get_context(message) + # check if the message is a command + if not ctx.valid: + self.handler.add_exp( + message.author.id, message.guild.id, len(message.content) + ) + + @commands.command() + async def exp(self, ctx: Context): + exp = self.handler.get_exp(ctx.author.id, ctx.guild.id) + await ctx.send( + f"{ctx.author.mention} has: `{exp['current_exp']}/{exp['required_exp']}` xp" + ) + + @commands.command() + async def level(self, ctx: Context): + level = self.handler.get_level(ctx.author.id, ctx.guild.id) + await ctx.send(f"{ctx.author.mention}'s level is: `{level}`") + + @exp.error + @level.error + async def chat_error(self, ctx: Context, error): + message = ":x: " + if isinstance(error, errors.CommandInvokeError): + cause = error.__cause__ + if isinstance(cause, (ServerSelectionTimeoutError, ConnectionFailure)): + message += "The database is currently taking a nap, try again later!" + else: + message += f"Command failed to execute due to: ```\n{error}\n```" + + await ctx.send(message) + + +def setup(client): + client.add_cog(RpgChat(client)) diff --git a/cogs/rpgchat.py b/cogs/rpgchat.py deleted file mode 100644 index 6c15e93..0000000 --- a/cogs/rpgchat.py +++ /dev/null @@ -1,111 +0,0 @@ -import math -import functools -from util.torudb import ToruDb -from discord.ext import commands - - -class RpgChat(commands.Cog): - def __init__(self, client) -> None: - self.client = client - - self.try_connect() - - def try_connect(self): - # if the server is non-existent then we - # just pretend everything is not working - try: - self.db = ToruDb() - self.is_connected = True - except: - self.db = None - self.is_connected = False - - return self.is_connected - - # decorator for handling an error, created for - # those pesky listeners since i don't know if - # there is a built-in error handling function - def handle_error(func): - @functools.wraps(func) - async def inner(self, *args, **kwargs): - # TODO: this just ignores the error for now, change later - try: - await func(self, *args, **kwargs) - except: - pass - - return inner - - def register(self, user, server): - self.db.update(user, server) - - def unregister(self, user, server): - self.db.remove(user, server) - - def chat(self, user, server, msg): - new_exp = RpgChat.expof(msg) + self.get_exp(user, server) - new_level = RpgChat.calc_level(new_exp) - new_info = {"chat_exp": new_exp, "level": new_level} - - self.db.update(user, server, new_info) - - def get_chat_info(self, user, server): - chat_info = self.db.get_chat_info(user, server) - if chat_info is None: - chat_info = {"chat_exp": 0, "level": 1} - - return chat_info - - def get_exp(self, user, server): - return self.get_chat_info(user, server).get("chat_exp", 0) - - def get_level(self, user, server): - return self.get_chat_info(user, server).get("level", 1) - - @commands.Cog.listener() - @handle_error - async def on_member_join(self, member): - self.register(member.id, member.guild.id) - - @commands.Cog.listener() - @handle_error - async def on_member_remove(self, member): - self.unregister(member.id, member.guild.id) - - @commands.Cog.listener() - @handle_error - async def on_message(self, message): - ctx = await self.client.get_context(message) - if not ctx.valid: - self.chat(message.author.id, message.guild.id, message.content) - - @commands.command() - async def exp(self, ctx): - await ctx.send( - f"Your current exp is: {self.get_exp(ctx.author.id, ctx.guild.id)}" - ) - - @commands.command() - async def level(self, ctx): - await ctx.send(f"Your level is: {self.get_level(ctx.author.id, ctx.guild.id)}") - - @level.error - @exp.error - async def error_handler(self, ctx, error): - # try to connect to the database, if success - # then next command invoke will succeed - self.try_connect() - - await ctx.send("Unable to connect to database, please try again later") - - def calc_level(exp): - # yes, the exp required to level up is literally just x^2 - return math.ceil(math.sqrt(exp)) - - def expof(msg): - # let's say you can get at most 50 exp - return min(len(msg), 50) - - -def setup(client): - client.add_cog(RpgChat(client)) diff --git a/util/error_handler.py b/util/error_handler.py new file mode 100644 index 0000000..828247c --- /dev/null +++ b/util/error_handler.py @@ -0,0 +1,98 @@ +"""Error Handler + +!!WORK IN PROGRESS!! + +This script provides useful decorators/methods for handling errors. +""" + +import logging +import functools +from typing import Type + + +def async_ignore_an_error(error_to_ignore: Type[Exception]): + """Decorator used to ignore a specific error for an async function + + Parameters + ---------- + error_to_ignore : Type[Exception] + The type of the error to be ignored + """ + + def decorator(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + try: + await func(*args, **kwargs) + except error_to_ignore: + logging.info(f"Ignored an error: {error_to_ignore}") + + return wrapper + + return decorator + + +def ignore_an_error(error_to_ignore: Type[Exception]): + """Decorator used to ignore a specific error for a blocking function + + Parameters + ---------- + error_to_ignore : Type[Exception] + The type of the error to be ignored + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + try: + func(*args, **kwargs) + except error_to_ignore: + logging.info(f"Ignored an error: {error_to_ignore}") + + return wrapper + + return decorator + + +def async_ignore_multiple_errors(errors_to_ignore: list[Type[Exception]]): + """Decorator used to ignore a list of errors for an async function + + Parameters + ---------- + errors_to_ignore : list[Type[Exception]] + The list of error types to ignore + """ + + def decorator(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + try: + await func(*args, **kwargs) + except tuple(errors_to_ignore) as e: + logging.info(f"Ignored an error: {e}") + + return wrapper + + return decorator + + +def ignore_multiple_errors(errors_to_ignore: list[Type[Exception]]): + """Decorator used to ignore a list of errors for a blocking function + + Parameters + ---------- + errors_to_ignore : list[Type[Exception]] + The list of error types to ignore + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + try: + func(*args, **kwargs) + except tuple(errors_to_ignore) as e: + logging.info(f"Ignored an error: {e}") + + return wrapper + + return decorator diff --git a/util/toru_rpg.py b/util/toru_rpg.py new file mode 100644 index 0000000..2198dfa --- /dev/null +++ b/util/toru_rpg.py @@ -0,0 +1,227 @@ +import util.torudb as db + + +class ToruRpg: + """ + A server class that represents a simple RPG experience but designed + directly for Toru-chan. + + This class communicates directly with the backend `torudb` and + introduces a simplified interface. + + Note that most methods in this class can raise a + `ServerSelectionTimeoutError` if the server is too busy or simply + offline + + Methods + ------- + register(user) + Registers the given user in the given server + + unregister(user) + Unregisters the given user from the given server + + get_detail(user) + Retrieves the details (such as exp, level) of the given user + + get_exp(user) + Retrieves the current exp and required exp of the given user + in a dictionary + + get_level(user) + Retrieves the level of the given user + + add_exp(user, exp) + Adds the amount of exp to the user and levels the user up + accordingly + + calc_exp(level): + Calculates the exp needed for the given level + """ + + def register(self, user: int, server: int) -> bool: + """Registers the given user in the given server + + Parameters + ---------- + user : int + The unique user id + server : int + The unique server id + + Returns + ------- + bool + True if registers successfully, False if the user is + already registered + + Raises + ------ + ServerSelectionTimeoutError + If the server is too busy or not up + """ + return db.register(user, server) + + def unregister(self, user: int, server: int) -> bool: + """Unregisters the given user from the given server + + Parameters + ---------- + user : int + The unique user id + server : int + The unique server id + + Returns + ------- + bool + True if unregisters successfully, False if the user is + not registered to begin with + + Raises + ------ + ServerSelectionTimeoutError + If the server is too busy or not up + """ + return db.unregister(user, server) + + def get_detail(self, user: int, server: int) -> dict[str, int]: + """Retrieves the details (such as exp, level) of the given user + + Parameters + ---------- + user : int + The unique user id + server : int + The unique server id + + Returns + ------- + dict[str, int] + A dictionary containing details about the user, see below + for an exmaple: + + { + "server": server_id1 + "current_exp": 69 + "required_exp": 420 + "level": 2 + } + + Raises + ------ + ServerSelectionTimeoutError + If the server is too busy or not up + """ + + # registers the user if not already done so + if not db.user_has_server(user, server): + db.register(user, server) + + return db.get_detail(user, server) + + def get_exp(self, user: int, server: int) -> dict[str, int]: + """Retrieves the current exp and required exp of the given user + in a dictionary + + Parameters + ---------- + user : int + The unique user id + server : int + The unique server id + + Returns + ------- + dict[str, int] + A dictionary containing the keys "current_exp" and + "required_exp" + + Raises + ------ + ServerSelectionTimeoutError + If the server is too busy or not up + """ + + detail = self.get_detail(user, server) + keys_to_extract = ["current_exp", "required_exp"] + + return {key: detail[key] for key in keys_to_extract} + + def get_level(self, user: int, server: int) -> int: + """Retrieves the level of the given user has in a server + + Parameters + ---------- + user : int + The unique user id + server : int + The unique server id + + Returns + ------- + int + The level of the given user + + Raises + ------ + ServerSelectionTimeoutError + If the server is too busy or not up + """ + return self.get_detail(user, server)["level"] + + def add_exp(self, user: int, server: int, exp: int) -> bool: + """Adds the amount of exp to the user and levels the user up + accordingly + + Parameters + ---------- + user : int + The unique user id + server : int + The unique server id + exp: int + The amount of exp to add + + Returns + ------- + bool + True if the exp updated successfully, else False + + Raises + ------ + ServerSelectionTimeoutError + If the server is too busy or not up + """ + + detail = self.get_detail(user, server) + + detail["current_exp"] += exp + + if detail["current_exp"] >= detail["required_exp"]: + detail["level"] += 1 + detail["current_exp"] = 0 + detail["required_exp"] = self.calc_exp(detail["level"]) + + return db.update(user, server, detail) + + def calc_exp(self, level: int) -> int: + """Calculates the exp needed for the given level + + Parameters + ---------- + level : int + The level to calculate exp for + + Returns + ------- + int + The exp need for the given level, 0 if the given level is + invalid (if the level is less or equal to 0) + """ + + # let's pretend this is a super fancy algorithm + if level < 1: + return 0 + + return int((level * 80) ** 1.051) diff --git a/util/torudb.py b/util/torudb.py index 344fabc..0b4d627 100644 --- a/util/torudb.py +++ b/util/torudb.py @@ -1,95 +1,239 @@ +"""Toru Database + +This script provides simple backend functions for Toru-chan's RPG chat +feature using MongoDB. + +Currently the functionalities are limited, and only basic CRUD +operations for the users are provided. + +Below shows an example of the schema for the only collection the +database utilises, for now: + + { + "user": user_id + "servers": [ + { + "server": server_id1 + "current_exp": 69 + "required_exp": 420 + "level": 2 + }, + { + "server": server_id2 + "current_exp": 0 + "required_exp": 1377 + "level": 4 + } + ] + } +""" + import os +import dotenv import logging -from dotenv import load_dotenv +import pymongo +from typing import Dict, Union from pymongo import MongoClient +from pymongo.errors import ServerSelectionTimeoutError -load_dotenv() -_MONGODB_URL = os.getenv("MONGODB_URL") - -""" -the intended scheme for the collection 'users' -{ - "user": user_id - "servers": [ - { - "server": server_id1 - "chat_exp": 69 - "level": 2 - }, - { - "server": server_id2 - "chat_exp": 420 - "level": 4 - } - ] -} -""" - -# TODO: error handling for those database calls and risky None values -class ToruDb: - def __init__(self) -> None: - # connect to the database - self.client = MongoClient( - host=f"mongodb://{_MONGODB_URL}:27017/", - serverSelectionTimeOutMS=500, - ) - logging.info(f"Server info: {self.client.server_info}") - - # the database object - self.db = self.client.toru - # the connection object - self.users = self.db.users - - # create indices to speed up the queries - self.users.create_index([("user", 1)], unique=True) - self.users.create_index([("servers.server", 1)]) - - def user_exists(self, user_id): - return self.users.find_one({"user": user_id}) is not None - - def user_has_server(self, user_id, server_id): - query = self.users.find_one({"user": user_id, "servers.server": server_id}) - return query is not None - - def get_chat_info(self, user_id, server_id): - if not self.user_has_server(user_id, server_id): - return None - - query = self.users.find_one({"user": user_id, "servers.server": server_id}) - - # wtf - return list( - filter( - lambda server: server.get("server") == server_id, query.get("servers") - ) - )[0] - - # inserts the user into database if the user does not exists, - # creates a server for the user if the user does not have the server, - # if chat_info is specified then also update the server info for the user - def update(self, user_id, server_id, chat_info=None): - if not self.user_exists(user_id): - user = {"user": user_id, "servers": []} - - self.users.insert_one(user) - - if not self.user_has_server(user_id, server_id): - server = {"server": server_id, "chat_exp": 0, "level": 1} +dotenv.load_dotenv() +MONGODB_URL = os.getenv("MONGODB_URL") - self.users.update_one({"user": user_id}, {"$push": {"servers": server}}) - if chat_info is not None: - chat_info.setdefault("server", server_id) +# create a connection to the database +client = MongoClient(MONGODB_URL, 27017, serverSelectionTimeoutMS=1000) +logging.info(f"Server info: {client.server_info}") - self.users.update_one( - {"user": user_id, "servers.server": server_id}, - {"$set": {"servers.$": chat_info}}, - ) +# retrieve/create the database object +db = client.toru - # used for when the user leaves a server - # why do we reset everything? - # becuase fuck you that's why - def remove(self, user_id, server_id): - self.users.update_one( - {"user": user_id}, {"$pull": {"servers": {"server": server_id}}} +# retrieve/create the primary collection object +users = db.users + +# in case the server wasn't running, just ignore it +try: + # create the indices to speed up the queries + users.create_index([("user", pymongo.ASCENDING)], unique=True) + users.create_index( + [("user", pymongo.ASCENDING), ("servers.server", pymongo.ASCENDING)], + unique=True, + ) +except ServerSelectionTimeoutError: + pass + + +def user_exists(user: int) -> bool: + """Checks if a given user exists in the database + + Parameters + ---------- + user : int + The unique user id + + Returns + ------- + bool + True if the user does exist, else False + """ + return users.find_one({"user": user}) is not None + + +def user_has_server(user: int, server: int) -> bool: + """Checks if a given user is registered on the given server + + Parameters + ---------- + user : int + The unique user id + + server : int + The unique server id + + Returns + ------- + bool + If the user does not exist or is not registered on the given + server then False, else True + """ + return users.find_one({"user": user, "servers.server": server}) is not None + + +def get_detail(user: int, server: int) -> Union[dict[str, int], None]: + """Retrieves the details the given user has on the given server + + Parameters + ---------- + user : int + The unique user id + + server : int + The unique server id + + Returns + ------- + Union[dict[str, int], None] + A dictionary representing the detail if the user has the given + server, else None + """ + if not user_has_server(user, server): + return None + + result = users.find_one({"user": user, "servers.server": server}) + + for detail in result["servers"]: + if detail["server"] == server: + return detail + + return None + + +def register(user: int, server: int) -> bool: + """Registers a given user on the given server + + Parameters + ---------- + user : int + The unique user id + + server : int + The unique server id + + Returns + ------- + bool + True if successfully registers the user, else False when the + user is already registered on the given server + """ + if not user_has_server(user, server): + users.find_one_and_update( + {"user": user}, + { + "$push": { + "servers": { + "server": server, + "current_exp": 0, + "required_exp": 100, + "level": 1, + } + } + }, + upsert=True, ) + return True + + return False + + +def unregister(user: int, server: int) -> bool: + """Unregisters a given user on the given server + + Parameters + ---------- + user : int + The unique user id + + server : int + The unique server id + + Returns + ------- + bool + True if sucessfully unregisters the user, else False when the + user is not registered to begin with + """ + return ( + users.update_one( + {"user": user}, {"$pull": {"servers": {"server": server}}} + ).modified_count + ) > 0 + + +def update(user: int, server: int, detail: Dict[str, int]) -> bool: + """Updates the given user's detail in a server with the given + details + + Parameters + ---------- + user : int + The unique user id + + server : int + The unique server id + + detail : Dict[str, int] + The detail to be updated to + + Returns + ------- + bool + True if update executed successfully, else False when the user + does not have the server, validate_detail() returns False on + the given detail or other unknown reasons + """ + if not (user_has_server(user, server) and validate_detail(detail)): + return False + + return ( + users.update_one( + {"user": user, "servers.server": server}, + {"$set": {"servers.$": detail}}, + upsert=True, + ).modified_count + ) > 0 + + +def validate_detail(detail: Dict[str, int]) -> bool: + """Validates the given detail + + Parameters + ---------- + detail : Dict[str, int] + The detail to be validated + + Returns + ------- + bool + True if the detail follows the schema, else False + """ + detail_keys = {"server", "current_exp", "required_exp", "level"} + return detail_keys <= detail.keys()