Add poll functionality (#24)

Author: klercke

protobot.py

@@ -12,20 +12,22 @@
 from dotenv import load_dotenv      #
 import random                       #
 from discord.ext import commands    #
+from discord.ext import tasks       #
 import logging                      #
 import time                         #
 import asyncio                      #
 import activity                     #
 import string                       #
 import schedule                     #
+import threading                    #
                                     #
 #####################################
 
 
 #####################################
                                     #
 COMMAND_PREFIX = '!'                #
-VERSION = "v0.4.1-alpha"            #
+VERSION = "v0.5.0-alpha"            #
 ACTIVITY = discord.Game("!help")    #
 LOG_LEVEL = logging.INFO            #
                                     #
@@ -74,6 +76,7 @@ async def on_ready():
     for guild in bot.guilds:
         add_all_users_from_guild_to_database(guild)
 
+    # Start scheduled events
     logging.info("Scheduling events...")
     logging.info("Daily events")
     schedule.every().day.at("00:00").do(run_once_every_day)
@@ -82,6 +85,8 @@ async def on_ready():
     logging.info("Minutely events")
     schedule.every().minute.do(run_once_every_minute)
     logging.info("Scheduling complete!")
+    t = threading.Thread(target=schedule_worker)
+    t.start()
 
 
 @bot.event
@@ -293,6 +298,104 @@ async def about(ctx):
     await ctx.message.channel.send(f"ProtoBot {VERSION}. Source code and bug tracker: https://github.com/klercke/ProtoBot")
 
 
+@bot.command(name="poll", help="Creates a poll. Usage: !poll \"PROMPT\" TIME OPTION1 OPTION2 [OPTION3...]")
+async def poll(ctx):
+    """
+    Allows users to create timed polls
+    """
+
+    user_input = ctx.message.content
+    user = ctx.message.author
+    
+    # Get the user's question
+    prompt = ""
+    writing = False
+    count = 0
+    for character in user_input:
+        count += 1
+        if not writing and character == "\"":
+            writing = True
+        elif writing:
+            if character == "\"":
+                writing = False
+                break
+            else:
+                prompt += character
+        elif count == len(user_input) - 1:
+            await ctx.channel.send("Sorry, I couldn't understand your command. Please make sure the poll question is in quotes.")
+            return
+    
+    # Get the TTL for the poll
+    user_input = user_input[count + 1:]
+    if not user_input:
+        await ctx.channel.send("You didn't provide a time. Assuming 5 minutes.")
+        poll_time = "5m"
+    else:
+        poll_time = ""
+        count = 0
+        for character in user_input:
+            count += 1
+            if character != " ":
+                poll_time += character
+            else:
+                break
+
+    unit_long = ""
+    unit = poll_time[-1].lower()
+    poll_time = int(poll_time[:-1])
+    poll_time_in_sec = 0
+    if unit == 's':
+        unit_long = "seconds"
+        poll_time_in_sec = poll_time
+    elif unit == 'm':
+        unit_long = "minutes"
+        poll_time_in_sec = poll_time * 60
+    elif unit == 'h':
+        unit_long = "hours"
+        poll_time_in_sec = poll_time * 3600
+    elif unit == 'd':
+        unit_long = "days"
+        poll_time_in_sec = poll_time * 86400
+
+
+    # Get emoji options
+    user_input = user_input[count:]
+    user_input = user_input.split()
+    options = []
+    message_sent = await ctx.message.channel.send(f"<@{user.id}> has started a poll:\n{prompt}\nVoting will last {poll_time} {unit_long}.")
+    for emoji in user_input:
+        options += emoji
+        await message_sent.add_reaction(emoji)
+
+    async def count_poll_results(message_sent, poll_time_in_sec):
+        # Wait for voting to finish
+        await asyncio.sleep(poll_time_in_sec)
+
+        # Create dictionary for results
+        results = {}
+        for option in options:
+            results[option] = 0
+
+        # Get original message
+        message_sent = await message_sent.channel.fetch_message(message_sent.id)
+
+        # Count results 
+        total_votes = 0
+        for reaction in message_sent.reactions:
+            results[reaction.emoji] = reaction.count - 1
+            total_votes += reaction.count - 1
+
+        if total_votes == 0:
+            await message_sent.channel.send(f"Voting for \"{prompt}\" complete. Nobody voted!")
+        else:
+            # Display results
+            winner = max(results, key = results.get)
+            winner_percentage = round((results[winner] / total_votes) * 100, 2)
+            await message_sent.channel.send(f"Voting for \"{prompt}\" complete. {winner} is the winner with {results[winner]} ({winner_percentage}%) votes!")
+
+    await count_poll_results(message_sent, poll_time_in_sec)
+    
+
 def run_once_every_day():
     """
     Runs a block of code every day.
@@ -353,6 +456,13 @@ def add_user_to_database(uuid, name, score=0, allowmoderator=True, rankexempt=Fa
     activity.write_database()
 
 
+def schedule_worker():
+    logging.info("Schedule worker thread initialized.")
+    while True:
+        schedule.run_pending()
+        time.sleep(1)
+
+
 def main():
     # Load bot token from .env
     load_dotenv()