Filter for admins; randomly pick for ties

icon-poller.py

@@ -7,12 +7,27 @@ import logging
 from PIL import Image
 import random
 
+from mwt import MWT
+
 submissions = {}
 names = {}
 polls = {}
 poll_options = {}
 status = {}
 
+@MWT(timeout=60*60)
+def get_admin_ids(bot, chat_id):
+    """Returns a list of admin IDs for a given chat. Results are cached for 1 hour."""
+    return [admin.user.id for admin in bot.get_chat_administrators(chat_id)]
+
+def ensure_admin(func):
+    def wrapper(update, context, *args, **kwargs):
+        chat_id = update.effective_chat.id
+        user_id = update.effective_user.id
+        if user_id in get_admin_ids(context.bot, chat_id):
+            return func(update, context, *args, **kwargs)
+    return wrapper
+
 def ensure_chat(func):
     def wrapper(update, *args, **kwargs):
         chat_id = update.effective_chat.id
@@ -85,13 +100,16 @@ def accept_photo(update: Update, context: CallbackContext):
     print(file_id)
 
 @ensure_chat
+@ensure_admin
 @ensure_closed
 def open_submissions(update: Update, context: CallbackContext):
     chat_id = update.effective_chat.id
+    submissions[chat_id] = {}
     context.bot.send_message(chat_id=chat_id, text="Send your images as a reply to this post.")
     status[chat_id] = "open"
     
 @ensure_chat
+@ensure_admin
 @ensure_open
 def start_poll(update: Update, context: CallbackContext):
     chat_id = update.effective_chat.id
@@ -102,30 +120,31 @@ def start_poll(update: Update, context: CallbackContext):
 
     print("Users with submissions: ", users)
 
-    chosen_users = random.choices(users, k=min(len(users), 10))
+    chosen_users = random.sample(users, k=min(len(users), 10))
 
     options = []
     chosen = []
 
-    for user_id in chosen_users:
-        options.append(names[chat_id][user_id])
+    for index, user_id in enumerate(chosen_users):
+        options.append(str(index+1) + ": " + names[chat_id][user_id])
         chosen.append(submissions[chat_id][user_id])
         poll_options[chat_id].append(user_id)
 
     if len(chosen) == 0:
         context.bot.send_message(chat_id=chat_id, text="There were no submissions.")
         status[chat_id] = "closed"
-    if len(chosen) == 1:
+    elif len(chosen) == 1:
         context.bot.send_photo(chat_id=chat_id, photo=chosen[0])
         context.bot.send_message(chat_id=chat_id, text="There was only one submission. " + options[0] + " is the winner.")
         status[chat_id] = "closed"
     else:
-        input_media = list(map(lambda file_id: InputMediaPhoto(file_id), chosen))
+        input_media = list(map(lambda pair: InputMediaPhoto(pair[1], caption=pair[0]), zip(options, chosen)))
         context.bot.send_media_group(chat_id=chat_id, media=input_media)    
 
         polls[chat_id] = context.bot.send_poll(chat_id=chat_id, options = options, question="Pick an icon").message_id
 
 @ensure_chat
+@ensure_admin
 @ensure_polling
 def decide_poll(update: Update, context: CallbackContext):
     chat_id = update.effective_chat.id
@@ -134,13 +153,16 @@ def decide_poll(update: Update, context: CallbackContext):
     result = context.bot.stop_poll(chat_id=chat_id, message_id=polls[chat_id])
 
     votes = -1
-    winner = 0
+    winners = []
 
     for index, option in enumerate(result.options):
         if option.voter_count > votes:
-            winner = index
+            winners = []
+        if option.voter_count >= votes:
+            winners.append(index)
             votes = option.voter_count
 
+    winner = random.choice(winners)
     winner_id = poll_options[chat_id][winner]
     winner_name = names[chat_id][winner_id]
     winner_photo = submissions[chat_id][winner_id]
@@ -149,6 +171,7 @@ def decide_poll(update: Update, context: CallbackContext):
     context.bot.send_photo(chat_id=chat_id, photo=winner_photo)
 
 @ensure_chat
+@ensure_admin
 def cancel(update: Update, context: CallbackContext):
     chat_id = update.effective_chat.id
     status[chat_id] = "closed"

mwt.py

@@ -0,0 +1,38 @@
+import time
+
+class MWT(object):
+    """Memoize With Timeout"""
+    _caches = {}
+    _timeouts = {}
+
+    def __init__(self,timeout=2):
+        self.timeout = timeout
+
+    def collect(self):
+        """Clear cache of results which have timed out"""
+        for func in self._caches:
+            cache = {}
+            for key in self._caches[func]:
+                if (time.time() - self._caches[func][key][1]) < self._timeouts[func]:
+                    cache[key] = self._caches[func][key]
+            self._caches[func] = cache
+
+    def __call__(self, f):
+        self.cache = self._caches[f] = {}
+        self._timeouts[f] = self.timeout
+
+        def func(*args, **kwargs):
+            kw = sorted(kwargs.items())
+            key = (args, tuple(kw))
+            try:
+                v = self.cache[key]
+                print("cache")
+                if (time.time() - v[1]) > self.timeout:
+                    raise KeyError
+            except KeyError:
+                print("new")
+                v = self.cache[key] = f(*args,**kwargs),time.time()
+            return v[0]
+        func.func_name = f.__name__
+
+        return func