diff --git a/flanabot/bots/flana_bot.py b/flanabot/bots/flana_bot.py index 9dc7d6a..de48efc 100644 --- a/flanabot/bots/flana_bot.py +++ b/flanabot/bots/flana_bot.py @@ -6,7 +6,7 @@ import math import random import re from abc import ABC -from typing import Iterable, Sequence, Type +from typing import Iterable, Sequence import flanaapis.geolocation.functions import flanaapis.weather.functions @@ -45,6 +45,9 @@ class FlanaBot(MultiBot, ABC): self.register(self._on_choose, constants.KEYWORDS['random'], priority=2) self.register(self._on_choose, (constants.KEYWORDS['choose'], constants.KEYWORDS['random']), priority=2) + self.register(self._on_delete_votes, (multibot_constants.KEYWORDS['deactivate'], constants.KEYWORDS['vote'])) + self.register(self._on_delete_votes, (multibot_constants.KEYWORDS['delete'], constants.KEYWORDS['vote'])) + self.register(self._on_dice, constants.KEYWORDS['dice']) self.register(self._on_hello, multibot_constants.KEYWORDS['hello']) @@ -90,8 +93,12 @@ class FlanaBot(MultiBot, ABC): self.register(self._on_unpunish, (multibot_constants.KEYWORDS['deactivate'], constants.KEYWORDS['punish'])) self.register(self._on_unpunish, (multibot_constants.KEYWORDS['activate'], multibot_constants.KEYWORDS['permission'])) - self.register(self._on_weather_chart, constants.KEYWORDS['weather_chart']) - self.register(self._on_weather_chart, (multibot_constants.KEYWORDS['show'], constants.KEYWORDS['weather_chart'])) + self.register(self._on_voting_ban, (multibot_constants.KEYWORDS['deactivate'], multibot_constants.KEYWORDS['permission'], constants.KEYWORDS['vote'])) + + self.register(self._on_voting_unban, (multibot_constants.KEYWORDS['activate'], multibot_constants.KEYWORDS['permission'], constants.KEYWORDS['vote'])) + + self.register(self._on_weather, constants.KEYWORDS['weather_chart']) + self.register(self._on_weather, (multibot_constants.KEYWORDS['show'], constants.KEYWORDS['weather_chart'])) self.register_button(self._on_poll_button_press, ButtonsGroup.POLL) self.register_button(self._on_roles_button_press, ButtonsGroup.ROLES) @@ -149,6 +156,22 @@ class FlanaBot(MultiBot, ABC): else: return text.split() + async def _get_poll_message(self, message: Message) -> Message | None: + if poll_message := message.replied_message: + if poll_message.contents.get('poll') is None: + return + return poll_message + elif ( + (message.chat.is_private or self.is_bot_mentioned(message)) + and + flanautils.cartesian_product_string_matching(message.text, constants.KEYWORDS['poll'], min_ratio=multibot_constants.PARSE_CALLBACKS_MIN_RATIO_DEFAULT) + and + (poll_message := Message.find_one({'contents.poll.is_active': True}, sort_keys=(('date', pymongo.DESCENDING),))) + ): + return await self.get_message(poll_message.chat.id, poll_message.id) + else: + return + @return_if_first_empty(exclude_self_types='FlanaBot', globals_=globals()) async def _manage_exceptions(self, exceptions: BaseException | Iterable[BaseException], context: Chat | Message): if not isinstance(exceptions, Iterable): @@ -298,6 +321,23 @@ class FlanaBot(MultiBot, ABC): async def _unpunish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): pass + async def _update_poll_buttons(self, message: Message): + if message.contents['poll']['is_multiple_answer']: + total_votes = len({option_vote[0] for option_votes in message.contents['poll']['votes'].values() if option_votes for option_vote in option_votes}) + else: + total_votes = sum(len(option_votes) for option_votes in message.contents['poll']['votes'].values()) + + if total_votes: + buttons = [] + for option, option_votes in message.contents['poll']['votes'].items(): + ratio = f'{len(option_votes)}/{total_votes}' + names = f"({', '.join(option_vote[1] for option_vote in option_votes)})" if option_votes else '' + buttons.append(f'{option} ➜ {ratio} {names}') + else: + buttons = list(message.contents['poll']['votes'].keys()) + + await self.edit(self._distribute_buttons(buttons), message) + # ---------------------------------------------- # # HANDLERS # # ---------------------------------------------- # @@ -339,6 +379,19 @@ class FlanaBot(MultiBot, ABC): else: await self.send(random.choice(('¿Que elija el qué?', '¿Y las opciones?', '?', '🤔')), message) + @admin + async def _on_delete_votes(self, message: Message): + if message.chat.is_group and not self.is_bot_mentioned(message) or not (poll_message := await self._get_poll_message(message)): + return + + await self.delete_message(message) + + for user in await self._find_users_to_punish(message): + for option_name, option_votes in poll_message.contents['poll']['votes'].items(): + poll_message.contents['poll']['votes'][option_name] = [option_vote for option_vote in option_votes if option_vote[0] != user.id] + + await self._update_poll_buttons(poll_message) + async def _on_dice(self, message: Message): if message.chat.is_group and not self.is_bot_mentioned(message): return @@ -414,23 +467,39 @@ class FlanaBot(MultiBot, ABC): self._distribute_buttons(final_options), message, buttons_key=ButtonsGroup.POLL, - contents={'poll': {'is_active': True, 'is_multiple_answer': is_multiple_answer, 'votes': {option: [] for option in final_options}}} + contents={'poll': { + 'is_active': True, + 'is_multiple_answer': is_multiple_answer, + 'votes': {option: [] for option in final_options}, + 'banned_users_tries': {} + }} ) else: await self.send(random.choice(('¿Y las opciones?', '?', '🤔')), message) - async def _on_poll_multi(self, message: Message): - await self._on_poll(message, is_multiple_answer=True) - async def _on_poll_button_press(self, message: Message): await self.accept_button_event(message) if not message.contents['poll']['is_active']: return - option_name = results[0] if (results := re.findall('(.*?) ➜.+', message.buttons_info.pressed_text)) else message.buttons_info.pressed_text - selected_option_votes = message.contents['poll']['votes'][option_name] presser_id = message.buttons_info.presser_user.id presser_name = message.buttons_info.presser_user.name.split('#')[0] + if (presser_id_str := str(presser_id)) in message.contents['poll']['banned_users_tries']: + message.contents['poll']['banned_users_tries'][presser_id_str] += 1 + message.save() + if message.contents['poll']['banned_users_tries'][presser_id_str] == 3: + await self.send(random.choice(( + f'Deja de dar por culo {presser_name}, que no puedes votar aqui.', + f'No es pesao {presser_name}, que no tienes permitido votar aqui', + f'Deja de pulsar botones que no puedes votar aqui {presser_name}', + f'{presser_name} deja de intentar votar aqui que no puedes', + f'Te han prohibido votar aqui {presser_name}', + f'No puedes votar aqui {presser_name}' + )), reply_to=message) + return + + option_name = results[0] if (results := re.findall('(.*?) ➜.+', message.buttons_info.pressed_text)) else message.buttons_info.pressed_text + selected_option_votes = message.contents['poll']['votes'][option_name] if [presser_id, presser_name] in selected_option_votes: selected_option_votes.remove([presser_id, presser_name]) @@ -445,21 +514,10 @@ class FlanaBot(MultiBot, ABC): break selected_option_votes.append((presser_id, presser_name)) - if message.contents['poll']['is_multiple_answer']: - total_votes = len({option_vote[0] for option_votes in message.contents['poll']['votes'].values() if option_votes for option_vote in option_votes}) - else: - total_votes = sum(len(option_votes) for option_votes in message.contents['poll']['votes'].values()) + await self._update_poll_buttons(message) - if total_votes: - buttons = [] - for option, option_votes in message.contents['poll']['votes'].items(): - ratio = f'{len(option_votes)}/{total_votes}' - names = f"({', '.join(option_vote[1] for option_vote in option_votes)})" if option_votes else '' - buttons.append(f'{option} ➜ {ratio} {names}') - else: - buttons = list(message.contents['poll']['votes'].keys()) - - await self.edit(self._distribute_buttons(buttons), message) + async def _on_poll_multi(self, message: Message): + await self._on_poll(message, is_multiple_answer=True) @bot_mentioned @group @@ -555,18 +613,7 @@ class FlanaBot(MultiBot, ABC): await self._manage_exceptions(SendError('No hay información musical en ese mensaje.'), message) async def _on_stop_poll(self, message: Message): - if poll_message := message.replied_message: - if poll_message.contents.get('poll') is None: - return - elif ( - (message.chat.is_private or self.is_bot_mentioned(message)) - and - flanautils.cartesian_product_string_matching(message.text, constants.KEYWORDS['poll'], min_ratio=multibot_constants.PARSE_CALLBACKS_MIN_RATIO_DEFAULT) - and - (poll_message := Message.find_one({'contents.poll.is_active': True}, sort_keys=(('date', pymongo.DESCENDING),))) - ): - poll_message = await self.get_message(poll_message.chat.id, poll_message.id) - else: + if not (poll_message := await self._get_poll_message(message)): return winners = [] @@ -603,34 +650,33 @@ class FlanaBot(MultiBot, ABC): for user in await self._find_users_to_punish(message): await self.unpunish(user, message, message) - async def _on_weather_button_press(self, message: Message): - await self.accept_button_event(message) + @admin + async def _on_voting_ban(self, message: Message): + if message.chat.is_group and not self.is_bot_mentioned(message) or not (poll_message := await self._get_poll_message(message)): + return - match message.buttons_info.pressed_text: - case WeatherEmoji.ZOOM_IN.value: - message.weather_chart.zoom_in() - case WeatherEmoji.ZOOM_OUT.value: - message.weather_chart.zoom_out() - case WeatherEmoji.LEFT.value: - message.weather_chart.move_left() - case WeatherEmoji.RIGHT.value: - message.weather_chart.move_right() - case WeatherEmoji.PRECIPITATION_VOLUME.value: - message.weather_chart.trace_metadatas['rain_volume'].show = not message.weather_chart.trace_metadatas['rain_volume'].show - message.weather_chart.trace_metadatas['snow_volume'].show = not message.weather_chart.trace_metadatas['snow_volume'].show - case emoji if emoji in WeatherEmoji.values: - trace_metadata_name = WeatherEmoji(emoji).name.lower() - message.weather_chart.trace_metadatas[trace_metadata_name].show = not message.weather_chart.trace_metadatas[trace_metadata_name].show - case _: - return + await self.delete_message(message) - message.weather_chart.apply_zoom() - message.weather_chart.draw() + for user in await self._find_users_to_punish(message): + if str(user.id) not in poll_message.contents['poll']['banned_users_tries']: + poll_message.contents['poll']['banned_users_tries'][str(user.id)] = 0 + poll_message.save() - image_bytes = message.weather_chart.to_image() - await self.edit(Media(image_bytes, MediaType.IMAGE, 'jpg'), message) + @admin + async def _on_voting_unban(self, message: Message): + if message.chat.is_group and not self.is_bot_mentioned(message) or not (poll_message := await self._get_poll_message(message)): + return - async def _on_weather_chart(self, message: Message): + await self.delete_message(message) + + for user in await self._find_users_to_punish(message): + try: + del poll_message.contents['poll']['banned_users_tries'][str(user.id)] + except KeyError: + pass + poll_message.save() + + async def _on_weather(self, message: Message): bot_state_message: Message | None = None if message.is_inline: show_progress_state = False @@ -775,6 +821,33 @@ class FlanaBot(MultiBot, ABC): # noinspection PyTypeChecker BotAction(Action.AUTO_WEATHER_CHART, message, affected_objects=[bot_message]).save() + async def _on_weather_button_press(self, message: Message): + await self.accept_button_event(message) + + match message.buttons_info.pressed_text: + case WeatherEmoji.ZOOM_IN.value: + message.weather_chart.zoom_in() + case WeatherEmoji.ZOOM_OUT.value: + message.weather_chart.zoom_out() + case WeatherEmoji.LEFT.value: + message.weather_chart.move_left() + case WeatherEmoji.RIGHT.value: + message.weather_chart.move_right() + case WeatherEmoji.PRECIPITATION_VOLUME.value: + message.weather_chart.trace_metadatas['rain_volume'].show = not message.weather_chart.trace_metadatas['rain_volume'].show + message.weather_chart.trace_metadatas['snow_volume'].show = not message.weather_chart.trace_metadatas['snow_volume'].show + case emoji if emoji in WeatherEmoji.values: + trace_metadata_name = WeatherEmoji(emoji).name.lower() + message.weather_chart.trace_metadatas[trace_metadata_name].show = not message.weather_chart.trace_metadatas[trace_metadata_name].show + case _: + return + + message.weather_chart.apply_zoom() + message.weather_chart.draw() + + image_bytes = message.weather_chart.to_image() + await self.edit(Media(image_bytes, MediaType.IMAGE, 'jpg'), message) + # -------------------------------------------------------- # # -------------------- PUBLIC METHODS -------------------- # # -------------------------------------------------------- # diff --git a/flanabot/constants.py b/flanabot/constants.py index 742dd76..4ff9bdb 100644 --- a/flanabot/constants.py +++ b/flanabot/constants.py @@ -105,6 +105,7 @@ KEYWORDS = { 'nombre', 'sonaba', 'sonando', 'song', 'sono', 'sound', 'suena', 'title', 'titulo', 'video'), 'unpunish': ('absolve', 'forgive', 'innocent', 'inocente', 'perdona', 'spare'), + 'vote': ('votacion', 'votar', 'vote', 'voting', 'voto'), 'weather_chart': ('atmosfera', 'atmosferico', 'calle', 'calor', 'caloret', 'clima', 'climatologia', 'cloud', 'cloudless', 'cloudy', 'cold', 'congelar', 'congelado', 'denbora', 'despejado', 'diluvio', 'frio', 'frost', 'hielo', 'humedad', 'llover', 'llueva', 'llueve', 'lluvia', 'nevada', 'nieva', 'nieve', diff --git a/tests/unit/test_parse_callbacks.py b/tests/unit/test_parse_callbacks.py index b28e675..2f89bbc 100644 --- a/tests/unit/test_parse_callbacks.py +++ b/tests/unit/test_parse_callbacks.py @@ -182,4 +182,4 @@ class TestParseCallbacks(unittest.TestCase): 'hara mucho calor en egipto este fin de semana?', 'pfff no ve que frio ahi en oviedo este finde' ] - self._test_no_always_callbacks(phrases, self.flana_tele_bot._on_weather_chart) + self._test_no_always_callbacks(phrases, self.flana_tele_bot._on_weather)