diff --git a/cogs/finance.py b/cogs/finance.py index 212ca84..74fd1ad 100644 --- a/cogs/finance.py +++ b/cogs/finance.py @@ -3,13 +3,15 @@ import requests import json from discord.ext import commands +ALPHAVANTAGE_API_KEY = os.getenv('ALPHAVANTAGE_API_KEY') +LUNAR_CRUSH_API_KEY = os.getenv("LUNAR_CRUSH_API_KEY") + class Finance(commands.Cog): """Commands to query stock and crypto prices.""" @commands.command(usage="") async def stock(self, ctx, *tickers: str) -> None: """Gets prices for the given stock tickers. Will default tickers to AAPL, GOOG, MSFT and AMZN if none provided.""" - ALPHAVANTAGE_API_KEY = os.getenv('ALPHAVANTAGE_API_KEY') # default tickers if len(tickers) == 0: @@ -43,7 +45,6 @@ class Finance(commands.Cog): @commands.command(usage="") async def crypto(self, ctx, *tickers: str) -> None: """Gets prices for the given crypto tickers. Will default tickers to BTC, ETH and LTC if none provided.""" - LUNAR_CRUSH_API_KEY = os.getenv("LUNAR_CRUSH_API_KEY") # default tickers if len(tickers) == 0: diff --git a/cogs/fun.py b/cogs/fun.py index f26ac30..da0238e 100644 --- a/cogs/fun.py +++ b/cogs/fun.py @@ -1,21 +1,17 @@ -from os import name import random import discord -from discord.ext.commands.core import command import requests import json import requests -from data import poll_cache +from model import Poll +from data import polls from datetime import date from discord.ext import commands +from discord.raw_models import RawReactionActionEvent class Fun(commands.Cog): """Commands that are good for the soul!""" - OPTION_EMOJIS = ["1️⃣", "2️⃣", "3️⃣", "4️⃣", "5️⃣"] - FULL_CHAR = "█" - BAR_LENGTH = 20 - @commands.command() async def flip(self, ctx) -> None: """Flips a coin and return either heads or tails.""" @@ -65,27 +61,25 @@ class Fun(commands.Cog): await ctx.send("**Please provide the parameters for the poll. ex.** `!poll \"My question?\" option1 option2`") return + # parse the parameters question = params[0] options = params[1:] # create embed response poll_embed = discord.Embed(title=question, description="\u200b", color=0x3DCCDD) for i in range(len(options)): - poll_embed.add_field(value="\u200b", name=f"{self.OPTION_EMOJIS[i]} {options[i]}", inline=False) + poll_embed.add_field(value="\u200b", name=f"{Poll.OPTION_EMOJIS[i]} {options[i]}", inline=False) poll_embed.set_footer(text=f"Poll created on • {date.today().strftime('%m/%d/%y')}") message = await ctx.send(embed=poll_embed) # add poll to cache - poll_cache[message.id] = {} - for i in range(len(options)): - poll_cache[message.id][f"{self.OPTION_EMOJIS[i]}"] = [] - poll_cache[message.id]["question"] = question - poll_cache[message.id]["options"] = ",".join(options) + poll = Poll(message.id, question, options) + polls[message.id] = json.dumps(poll.__dict__) - # add option emojis to poll + # add option emoji reactions to poll for i in range(len(options)): - await message.add_reaction(self.OPTION_EMOJIS[i]) + await message.add_reaction(Poll.OPTION_EMOJIS[i]) @poll.command(usage="", name="results") async def poll_results(self, ctx, message_id: int = None) -> None: @@ -95,27 +89,39 @@ class Fun(commands.Cog): await ctx.send("**Please provide the message id of a poll. ex.** `!poll results 870733802190807050") return - if message_id not in poll_cache: + if message_id not in polls: await ctx.send("**Could not find poll with that message id.**") return - poll = poll_cache[message_id] - question = poll["question"] - options = poll["options"].split(",") - - total_votes = sum([0 if item[0] == "question" or item[0] == "options" else len(item[1]) for item in poll.items()]) - option_percentages = [] - for i in range(len(options)): - option_votes = len(poll[self.OPTION_EMOJIS[i]]) - option_percentages.append(0 if total_votes == 0 else option_votes / total_votes) + poll = Poll.create_from(json.loads(polls[message_id])) # construct embed response - results_embed = discord.Embed(title=question, description="🥁", color=0x3DCCDD) - for i in range(len(options)): - str_bar = "".join([self.FULL_CHAR for _ in range(int(self.BAR_LENGTH * option_percentages[i]))]) - str_percent = f"{option_percentages[i] * 100:.2f}%" - str_votes = f"({int(option_percentages[i] * total_votes)} votes)" - results_embed.add_field(name=f"{self.OPTION_EMOJIS[i]} {options[i]}", value=f"{str_bar} {str_percent} {str_votes}", inline=False) + results_embed = discord.Embed(title=poll.question, description="🥁", color=0x3DCCDD) + for option in poll.options: + str_bar = "".join([Poll.BAR_CHAR for _ in range(int(Poll.BAR_LENGTH * poll.get_vote_percentage(option)))]) + str_percent = f"{poll.get_vote_percentage(option) * 100:.2f}%" + str_votes = f"({poll.get_vote_count(option)} votes)" + results_embed.add_field(name=f"{poll.get_emoji(option)} {option}", value=f"{str_bar} {str_percent} {str_votes}", inline=False) results_embed.set_footer(text=f"Poll queried on • {date.today().strftime('%m/%d/%y')}") await ctx.send(embed=results_embed) + + @staticmethod + def poll_vote(event: RawReactionActionEvent) -> None: + """Handles a user voting on a poll.""" + + # check if message is a poll and emoji is a valid option + if event.message_id in polls and event.emoji.name in Poll.OPTION_EMOJIS: + # get poll from cache + poll = Poll.create_from(json.loads(polls[event.message_id])) + + print(event) + + # check reaction emoji type + if event.event_type == "REACTION_ADD": + poll.add_vote(event.user_id, event.emoji.name) + elif event.event_type == "REACTION_REMOVE": + poll.remove_vote(event.user_id, event.emoji.name) + + # update cache + polls[event.message_id] = json.dumps(poll.__dict__) diff --git a/data/__init__.py b/data/__init__.py index c091082..9da43fc 100644 --- a/data/__init__.py +++ b/data/__init__.py @@ -1 +1 @@ -poll_cache = {} # TODO: use a real cache \ No newline at end of file +from data.redis import polls \ No newline at end of file diff --git a/data/redis.py b/data/redis.py new file mode 100644 index 0000000..e5b0a0d --- /dev/null +++ b/data/redis.py @@ -0,0 +1,6 @@ +import os +from walrus import * + +_REDIS_URL = os.getenv('REDIS_URL') +_db = Database.from_url(_REDIS_URL) +polls = _db.Hash("polls") \ No newline at end of file diff --git a/model/__init__.py b/model/__init__.py new file mode 100644 index 0000000..46f4faf --- /dev/null +++ b/model/__init__.py @@ -0,0 +1 @@ +from model.poll import Poll \ No newline at end of file diff --git a/model/poll.py b/model/poll.py new file mode 100644 index 0000000..8dbf61d --- /dev/null +++ b/model/poll.py @@ -0,0 +1,65 @@ +from typing import Dict, List + +class Poll: + + OPTION_EMOJIS = ["1️⃣", "2️⃣", "3️⃣", "4️⃣", "5️⃣"] + BAR_CHAR = "█" + BAR_LENGTH = 20 + + def __init__(self, message_id: int, question: str, options: List[str], votes: Dict = None): + self.message_id = message_id + self.question = question + self.options = options + self.votes = {option: [] for option in options} if not votes else votes + + def get_vote_count(self, option: str) -> int: + if option not in self.options: + return 0 + return len(self.votes[option]) + + def get_total_vote_count(self) -> int: + return sum([self.get_vote_count(option) for option in self.options]) + + def get_vote_percentage(self, option: str) -> float: + vote_count = self.get_vote_count(option) + + # avoid divide by zero + return 0 if self.get_total_vote_count() == 0 else vote_count / self.get_total_vote_count() + + def add_vote(self, user_id: str, emoji: str) -> None: + option = self._get_option(emoji) + + # only add vote if user has not already voted + if user_id not in self.votes[option]: + self.votes[option].append(user_id) + + print(self.votes) # TODO: remove + + def remove_vote(self, user_id: str, emoji: str) -> None: + option = self._get_option(emoji) + + # only remove if user has voted + if user_id in self.votes[option]: + self.votes[option].remove(user_id) + + print(self.votes) # TODO: remove + + def get_emoji(self, option: str) -> str: + return Poll.OPTION_EMOJIS[self.options.index(option)] + + def _get_option(self, emoji: str) -> str: + # check if valid emoji + if emoji not in self.OPTION_EMOJIS: + return None + + option_index = Poll.OPTION_EMOJIS.index(emoji) + + # check if valid vote + if option_index >= len(self.options): + return None + + return self.options[option_index] + + @staticmethod + def create_from(poll_dict: Dict) -> "Poll": + return Poll(int(poll_dict["message_id"]), poll_dict["question"], poll_dict["options"], poll_dict["votes"]) \ No newline at end of file diff --git a/pyvis.py b/pyvis.py index bdd852e..c69360d 100644 --- a/pyvis.py +++ b/pyvis.py @@ -1,20 +1,18 @@ +# load environment variables from .env file +from dotenv import load_dotenv +load_dotenv() + import discord import os -from discord.flags import MessageFlags import requests import json import random from cogs import Programming from cogs import Fun from cogs import Finance -from data import poll_cache from discord.ext import commands -from dotenv import load_dotenv from pretty_help import DefaultMenu, PrettyHelp -# load environment variables from .env file -load_dotenv() - DISCORD_BOT_TOKEN = os.getenv('DISCORD_BOT_TOKEN') DISCORD_GUILD_NAME = os.getenv('DISCORD_GUILD_NAME') GIPHY_API_KEY = os.getenv("GIPHY_API_KEY") @@ -36,7 +34,7 @@ async def on_ready(): @bot.event async def on_member_join(member): - """ WELCOME START """ + """ WELCOME START """ # TODO: extract to a separate file guild = discord.utils.find(lambda g: g.name == DISCORD_GUILD_NAME, bot.guilds) # TODO: check if find will throw an error welcome_channel = discord.utils.find(lambda c: c.name == "👋welcome", guild.channels) # TODO: remove welcome channel hardcode @@ -57,7 +55,7 @@ async def on_message(message): if message.author == bot.user: return - # TODO: is channel is a proposal channel check for poll command + # TODO: if channel is a proposal channel check for poll command # repsond with command prefix if bot.user.mentioned_in(message) and not message.mention_everyone: @@ -71,32 +69,14 @@ async def on_raw_reaction_add(event): # avoid bot self trigger if event.user_id == bot.user.id: return - - """ POLL START """ - # check if message is in cache and if reaction is valid - if event.message_id in poll_cache and event.emoji.name in Fun.OPTION_EMOJIS: - # if user has not voted yet, add their vote - if event.user_id not in poll_cache[event.message_id][event.emoji.name]: - poll_cache[event.message_id][event.emoji.name].append(event.user_id) - """ POLL END """ - - print(poll_cache) + Fun.poll_vote(event) @bot.event async def on_raw_reaction_remove(event): # avoid bot self trigger if event.user_id == bot.user.id: return - - """ POLL START """ - # check if message is in cache and if reaction is valid - if event.message_id in poll_cache and event.emoji.name in Fun.OPTION_EMOJIS: - # if user has voted, remove their vote - if event.user_id in poll_cache[event.message_id][event.emoji.name]: - poll_cache[event.message_id][event.emoji.name].remove(event.user_id) - """ POLL END """ - - print(poll_cache) + Fun.poll_vote(event) # override default help command menu = DefaultMenu('◀️', '▶️', '❌') diff --git a/requirements.txt b/requirements.txt index c0d30c5..7551cd1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,9 @@ discord.py==1.7.3 idna==3.2 multidict==5.1.0 python-dotenv==0.19.0 +redis==3.5.3 requests==2.26.0 typing-extensions==3.10.0.0 urllib3==1.26.6 +walrus==0.8.2 yarl==1.6.3