diff --git a/icon-poller.py b/icon-poller.py index fd27539..198db1d 100644 --- a/icon-poller.py +++ b/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" diff --git a/mwt.py b/mwt.py new file mode 100644 index 0000000..de30d28 --- /dev/null +++ b/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 \ No newline at end of file