Files
agon/utilities/utils.py
T

399 lines
15 KiB
Python

"""
Contains various utlity functions for highscores bot
"""
import json
import time
import nextcord
import pandas as pd
from sqlalchemy import distinct, select
from sqlalchemy.orm import Session, aliased
# pylint: disable=import-error
from utilities.sql_data_classes import Category, CategoryGroup, Guild, Score, engine
MAX_LINE_LENGTH = 142
def autofill_categories(guild_id: int, group_id: int, category_begin: str) -> "dict[str, str]":
"""
Returns all category names that start with category_begin.
If category_begin is empty, returns list of all category names.
Args:
guild_id (int): Guild id that holds the group and categories.
group_id (int): Group id to search in.
category_begin (str): The beginning of a category name to search with. Can be empty.
Returns:
dict[str, str]: List of all category names that start with category_begin.
"""
with Session(engine) as session:
# search database for all category names
subq = (
select(CategoryGroup)
.where(CategoryGroup.GuildId == guild_id)
.where(CategoryGroup.GroupId == group_id)
.subquery()
)
cat_group_alias = aliased(CategoryGroup, subq)
stmt = select(Category).where(Category.GroupId == cat_group_alias.GroupId)
categories_dict: dict = {cat.CategoryName: str(cat.CategoryId) for cat in session.scalars(stmt)}
# if string empty, return all category names
if not category_begin:
return categories_dict
# else find all categories that start with category_begin
# TODO is it better to have sql handle the starts with or python?
categories: dict = {
cat_name: cat_id for cat_name, cat_id in categories_dict.items() if category_begin.lower() in cat_name.lower()
}
return categories
def autofill_groups(guild_id: int, group_begin: str) -> "dict[str, str]":
"""
Returns all group names that start with group_begin. If group_begin is empty, returns list of all group names.
Args:
guild_id (int): Guild id that holds the groups.
group_begin (str): The beginning of a group name to search with. Can be empty.
Returns:
dict[str, str]: Dict of group names and ids where group name start with group_begin. Key is name, value is id.
"""
with Session(engine) as session:
# Select all groups in guild.
stmt = select(CategoryGroup).where(CategoryGroup.GuildId == guild_id)
# dict has to have key values of same type for discord.
group_dict: dict = {group.GroupName: str(group.GroupId) for group in session.scalars(stmt)}
# if string empty, return all group names
if not group_begin:
return group_dict
# else find all groups that start with group_begin
groups: dict = {
group_name: group_id for group_name, group_id in group_dict.items() if group_begin.lower() in group_name.lower()
}
return groups
def autofill_users_in_group_category(guild_id: int, group_id: int, category_id: int, user_begin: str) -> "list[str]":
"""
Returns list of users that match user_begin. Either all users in group or just users in category if category is
not None. Might be expensive due to unioning all users in group.
Args:
guild_id (int): Guild id corresponding to group and category.
group_id (str): group to look in for category.
category_id (str): Category to look in for users. Can be None if not looking for category.
user_begin (str): Beginning of user name to search with. Can be empty.
Returns:
list[str]: List of users that start with user_begin.
"""
with Session(engine) as session:
# get all categories in the group
subq = (
select(CategoryGroup)
.where(CategoryGroup.GuildId == guild_id)
.where(CategoryGroup.GroupId == group_id)
.subquery()
)
category_group = aliased(CategoryGroup, subq)
# If no category id, get all users in group.
if category_id == 0:
subq = select(Category.CategoryId).where(Category.GroupId == category_group.GroupId)
stmt = select(distinct(Score.UserName)).where(Score.CategoryId.in_(subq))
else:
# get all users in category_id
subq = (
select(Category)
.where(Category.GroupId == category_group.GroupId)
.where(Category.CategoryId == category_id)
.subquery()
)
category = aliased(Category, subq)
stmt = select(distinct(Score.UserName)).where(Score.CategoryId == category.CategoryId)
users_list: "list[str]" = list(session.scalars(stmt))
# if string empty, return all users
if not user_begin:
return users_list
# else find all users that start with user_begin
users = [user for user in users_list if user_begin.lower() in user.lower()]
return users
def autofill_users_in_guild(guild_id: int, user_begin: str) -> "list[str]":
"""
Obtains and returns a list of all the users who have submitted to the guild's highscores.
Args:
guild_id (int): Id of the guild.
user_begin (str): Beginning of user name to search with. Can be empty.
Returns:
list[str]: All users with submissions in the guild.
"""
with Session(engine) as session:
group_ids = select(CategoryGroup.GroupId).where(CategoryGroup.GuildId == guild_id)
category_ids = select(Category.CategoryId).where(Category.GroupId.in_(group_ids))
stmt = select(distinct(Score.UserName)).where(Score.CategoryId.in_(category_ids))
users_list: "list[str]" = list(session.scalars(stmt))
# if string empty, return all users
if not user_begin:
return users_list
# else find all users that start with user_begin
users = [user for user in users_list if user_begin.lower() in user.lower()]
return users
def format_highscore_message(session: Session, group: CategoryGroup) -> str:
"""
Creates a formatted string for a boss's highscore discord message.
This formatted string is in discord's code block to ensure monospace text.
Loops through each category to build the message one place at a time.
Args:
session (Session): the sqlalchemy session used for this function call.
group (CategoryGroup): group object that contains the contents of the message.
Returns:
str: Formatted string for discord message
"""
ret_string = f"**{group.GroupName}**\n```"
stmt = select(Guild).where(Guild.GuildId == group.GuildId)
# Get all categories in the group
subq = (
select(CategoryGroup)
.where(CategoryGroup.GuildId == group.GuildId)
.where(CategoryGroup.GroupName == group.GroupName)
.subquery()
)
cat_group_alias = aliased(CategoryGroup, subq)
stmt = select(Category).where(Category.GroupId == cat_group_alias.GroupId)
# set up a map for the categories
# also grab the max category size
categories = []
category_ids = []
max_cat_size: int = 0
pad_length = group.PadLength
for cat in session.scalars(stmt):
if cat.Size > max_cat_size:
max_cat_size = cat.Size
categories.append(cat)
category_ids.append(cat.CategoryId)
# get all scores in category group, load into data frame.
d_f = pd.read_sql_query(sql=session.query(Score).where(Score.CategoryId.in_(category_ids)).statement, con=engine)
# this list is what will be printed in the discord message
# first string will be category list, others will be score printouts
str_list: "list[str]" = ["" for i in range(max_cat_size + 1)]
groups_df = d_f.groupby("CategoryId")
# TODO figure out how to have less nests.
for category in categories:
# category name string
str_list[0] += f"{category.CategoryName}".ljust(pad_length, " ")
# ensure there is data to print
if category.CategoryId in groups_df.indices:
group_df = groups_df.get_group(category.CategoryId)
# process data. Want each group sorted correctly and only top X results
processed_group = (
group_df.sort_values("Score", ascending=category.IsAscending).head(category.Size).reset_index()
)
# add each score to list
for ind in range(max_cat_size): # need to subtract to deal with category name string
if ind in processed_group.index:
user_length = group.UsernameLength
username = processed_group["UserName"][ind][:user_length].ljust(user_length, " ")
# get score, pad to correct length
if category.IsTimeRecord:
# MM:SS.ms is 7 characters long
# if there are single digit minutes, add a space for formatting.
score = get_time_str(processed_group["Score"][ind])
if score.find(":") and len(score.split(":", maxsplit=1)[0]) < 2:
score = " " + score
else:
score = str(int(processed_group["Score"][ind])).rjust(7, " ")
str_list[ind + 1] += f"{ind + 1}: {username} - {score}".ljust(pad_length, " ")
else:
# if not enough entries, pad out
if ind + 1 <= category.Size:
str_list[ind + 1] += f"{ind + 1}: ".ljust(pad_length, " ")
else:
str_list[ind + 1] += "".ljust(pad_length, " ")
else:
for ind in range(max_cat_size):
if ind + 1 <= category.Size:
str_list[ind + 1] += f"{ind + 1}: ".ljust(pad_length, " ")
else:
str_list[ind + 1] += "".ljust(pad_length, " ")
# next step is to cut down on the length of each line if required
# need to use index to allow for modifying the list
for ind, line in enumerate(str_list):
if len(line) > MAX_LINE_LENGTH:
split_loc = int(MAX_LINE_LENGTH / pad_length) * pad_length
new_line = line[split_loc:]
str_list[ind] = line[:split_loc]
str_list.append(new_line)
# this list can now be placed into the format message function in the bot
for ind, string in enumerate(str_list):
# rstrip to clear unneeded spaces
# add newline here for each string
ret_string += string.rstrip() + "\n"
# add extra newline at category split unless there is already a space.
if ((ind + 1) % (max_cat_size + 1)) == 0:
# checking inside the other if to ensure called as little as possible.
if not ret_string.endswith("\n\n"):
ret_string += "\n"
return ret_string + "```"
def get_time_str(time_in: float) -> str:
"""
Formats a float value into the desired string format for a timestamp.
Args:
time_in (float): A time record in seconds
Returns:
str: Describes the time record in string format.
"""
# TODO see if theres a better way to do this
milliseconds = round((time_in - int(time_in)) * 10)
return time.strftime(f"%-M:%S.{milliseconds}", time.gmtime(time_in))
async def send_highscore_message(session: Session, guild: nextcord.Guild, group: CategoryGroup) -> nextcord.Message:
"""
Function takes care of checking if a message exists for a group.
Creates and sends/edits the message.
Args:
session (Session): The sqlalchemy session being used for this function call.
guild (nextcord.Guild): the nextcord.Guild object requesting this message.
group (CategoryGroup): group to post to the message.
Returns:
nextcord.Message: Message object that was posted with group contents.
"""
highscore_string = format_highscore_message(session, group)
channel = guild.get_channel(group.ChannelId)
try:
message: nextcord.Message = await channel.fetch_message(group.MessageId)
await message.edit(content=highscore_string)
except nextcord.NotFound:
# message doesn't exist.
message: nextcord.Message = await channel.send(highscore_string)
group.MessageId = message.id
session.commit()
return message
async def submit_score(
interaction: nextcord.Interaction, group_id: int, category_id: int, username: str, score: str
) -> None:
"""
Adds passed score to group. Then updates group message by calling send_highscore_message.
Ensures that a highscores channel exists.
Args:
interaction (nextcord.Interaction): Interaction that caused the score submission.
group (CategoryGroup): group to submit score to.
category_id (Category): Category to submit score to.
username (str): The user who submitted the score.
score (str): Score to submit. Need to convert to correct type.
"""
with Session(engine) as session:
stmt = select(Category.IsTimeRecord).where(Category.CategoryId == category_id)
is_time_record = session.scalar(stmt)
if is_time_record:
minutes = float(score.split(":", maxsplit=1)[0])
seconds = float(score.split(":", maxsplit=1)[1])
score = (minutes * 60) + seconds
else:
score = float(score)
stmt = select(Score).where(Score.UserName == username).where(Score.CategoryId == category_id)
existing_score = session.scalar(stmt)
if existing_score:
existing_score.Score = score
else:
score_obj = Score(UserName=username, Score=score, CategoryId=category_id)
session.add(score_obj)
session.commit()
stmt = select(CategoryGroup).where(CategoryGroup.GroupId == group_id)
group = session.scalar(stmt)
await send_highscore_message(session, interaction.guild, group)
def add_preset_configuration(guild_id: int, config_str: str) -> None:
"""
Adds an example configuration to a Guild. Uses guild id to know which guild and
config name to open a file in example_config.
Args:
guild_id (int): Guild id to add config to
config_str (str): Config name to add
"""
with Session(engine) as session:
with open(f"example_config/{config_str}.json", encoding="utf-8", mode="r") as file:
config = json.load(file)
for group_name, categorys in config["tables"].items():
group_obj = CategoryGroup(GroupName=group_name, GuildId=guild_id)
session.add(group_obj)
for category_name in categorys:
cat_obj = Category(
CategoryName=category_name,
Size=config["table_size"],
IsTimeRecord=config["categories"][category_name]["is_time_record"],
IsAscending=config["categories"][category_name]["ascending"],
)
session.add(cat_obj)
session.commit()