Files
agon/utilities/utils.py
T

399 lines
14 KiB
Python

"""Contains various utility functions for highscores bot."""
import time
import nextcord
import pandas as pd
from sqlalchemy import distinct, select
from sqlalchemy.orm import Session, aliased
# pylint: disable=import-error
from database_setup.db 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[str, str] = {cat.CategoryName: str(cat.CategoryId) for cat in session.scalars(stmt)}
# if string empty, return all category names
if not category_begin:
return dict(list(categories_dict.items())[:25])
# else find all categories that start with category_begin
# TODO is it better to have sql handle the starts with or python?
categories: dict[str, str] = {
cat_name: cat_id for cat_name, cat_id in categories_dict.items() if category_begin.lower() in cat_name.lower()
}
return dict(list(categories.items())[:25])
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[str, str] = {group.GroupName: str(group.GroupId) for group in session.scalars(stmt)}
# if string empty, return all group names
if not group_begin:
return dict(list(group_dict.items())[:25])
# else find all groups that start with group_begin
groups: dict[str, str] = {
group_name: group_id for group_name, group_id in group_dict.items() if group_begin.lower() in group_name.lower()
}
return dict(list(groups.items())[:25])
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[:25]
# else find all users that start with user_begin
users: list[str] = [user for user in users_list if user_begin.lower() in user.lower()]
return users[:25]
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[:25]
# else find all users that start with user_begin
users: list[str] = [user for user in users_list if user_begin.lower() in user.lower()]
return users[:25]
# TODO(imcompiling): #73 Address complexity warning.
# 073
def format_highscore_message(session: Session, group: CategoryGroup) -> str:
"""
Creates a formatted string for a category'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):
max_cat_size = max(cat.Size, max_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")
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:
username = processed_group["UserName"][ind]
score = processed_group["Score"][ind]
str_list[ind + 1] += format_score_string(
username,
score,
group.UsernameLength,
category.IsTimeRecord,
ind + 1,
pad_length,
)
# if not enough entries, pad out
elif 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 and not ret_string.endswith("\n\n"):
ret_string += "\n"
return ret_string + "```"
def format_score_string(
username: str,
score: float,
username_length: int,
is_time_record: bool,
rank: int,
pad_length: int,
) -> str:
"""
Builds the string to show a user's score.
Args:
username (str): username.
score (int): score in int form.
username_length (float): limit to how long username is allowed to be.
is_time_record (bool): if the score represents an number or a time record.
rank (int): the user's rank.
pad_length (int): limit to how long the string can be.
Returns:
str: string to show user's score.
"""
username = username[:username_length].ljust(username_length, " ")
# get score, pad to correct length
if is_time_record:
# MM:SS.ms is 7 characters long
# if there are single digit minutes, add a space for formatting.
score = get_time_str(score)
if score.find(":") and len(score.split(":", maxsplit=1)[0]) < 2:
score = " " + score
else:
score = str(int(score)).rjust(7, " ")
return f"{rank}: {username} - {score}".ljust(pad_length, " ")
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)
if group.MessageId is None:
# message doesn't exist.
message: nextcord.Message = await channel.send(highscore_string)
group.MessageId = message.id
session.commit()
else:
message: nextcord.Message = await channel.fetch_message(group.MessageId)
await message.edit(content=highscore_string)
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)