diff --git a/flanabot/bots/connect_4_bot.py b/flanabot/bots/connect_4_bot.py index df3b673..41a75c1 100644 --- a/flanabot/bots/connect_4_bot.py +++ b/flanabot/bots/connect_4_bot.py @@ -272,38 +272,29 @@ class Connect4Bot(MultiBot, ABC): ) -> bool: if board[i][j] in self._check_winners(i, j, board): winner, loser = (player_1, player_2) if board[i][j] == player_1.number else (player_2, player_1) - try: - message.data['connect_4']['is_active'] = False - except KeyError: - pass - await self.edit( - Media( - connect_4_frontend.make_image(board, winner=winner, loser=loser, highlight=(i, j), win_position=(i, j)), - MediaType.IMAGE, - 'png', - Source.LOCAL - ), - message, - buttons=[] - ) - return True + edit_kwargs = {'winner': winner, 'loser': loser, 'win_position': (i, j)} + elif turn >= constants.CONNECT_4_N_ROWS * constants.CONNECT_4_N_COLUMNS: + edit_kwargs = {'tie': True} + else: + return False - if turn >= constants.CONNECT_4_N_ROWS * constants.CONNECT_4_N_COLUMNS: - try: - message.data['connect_4']['is_active'] = False - except KeyError: - pass - await self.edit( - Media( - connect_4_frontend.make_image(board, highlight=(i, j), tie=True), - MediaType.IMAGE, - 'png', - Source.LOCAL - ), - message, - buttons=[] - ) - return True + try: + message.buttons_info.data['is_active'] = False + except AttributeError: + pass + + await self.edit( + Media( + connect_4_frontend.make_image(board, highlight=(i, j), **edit_kwargs), + MediaType.IMAGE, + 'png', + Source.LOCAL + ), + message, + buttons=[] + ) + + return True @staticmethod def _check_winner_left(i: int, j: int, board: list[list[int | None]]) -> int | None: @@ -481,26 +472,25 @@ class Connect4Bot(MultiBot, ABC): message=message, buttons=self.distribute_buttons([str(n) for n in range(1, constants.CONNECT_4_N_COLUMNS + 1)]), buttons_key=ButtonsGroup.CONNECT_4, - data={'connect_4': { + buttons_data={ 'is_active': True, 'board': board, 'player_1': player_1.to_dict(), 'player_2': player_2.to_dict(), 'turn': 0 - }} + } ) await self.delete_message(message) async def _on_connect_4_button_press(self, message: Message): await self.accept_button_event(message) - is_active = message.data['connect_4']['is_active'] - board = message.data['connect_4']['board'] - player_1 = Player.from_dict(message.data['connect_4']['player_1']) - player_2 = Player.from_dict(message.data['connect_4']['player_2']) - turn = message.data['connect_4']['turn'] + is_active = message.buttons_info.data['is_active'] + board = message.buttons_info.data['board'] + player_1 = Player.from_dict(message.buttons_info.data['player_1']) + player_2 = Player.from_dict(message.buttons_info.data['player_2']) - if turn % 2 == 0: + if message.buttons_info.data['turn'] % 2 == 0: current_player = player_1 next_player = player_2 else: @@ -511,12 +501,11 @@ class Connect4Bot(MultiBot, ABC): if not is_active or current_player.id != presser_id or board[0][move_column] is not None: return - message.data['connect_4']['is_active'] = False - message.save() + message.buttons_info.data['is_active'] = False i, j = self.insert_piece(move_column, current_player.number, board) - turn += 1 - if await self._check_game_finished(i, j, player_1, player_2, turn, board, message): + message.buttons_info.data['turn'] += 1 + if await self._check_game_finished(i, j, player_1, player_2, message.buttons_info.data['turn'], board, message): return await self.edit( @@ -530,22 +519,20 @@ class Connect4Bot(MultiBot, ABC): ) if player_2.id == self.id: - turn += 1 + message.buttons_info.data['turn'] += 1 if await self._ai_turn( player_1, player_2, next_player, current_player, - turn, + message.buttons_info.data['turn'], constants.CONNECT_4_AI_DELAY_SECONDS, board, message ): return - message.data['connect_4']['turn'] = turn - message.data['connect_4']['is_active'] = True - message.save() + message.buttons_info.data['is_active'] = True async def _on_connect_4_vs_itself(self, message: Message): if message.chat.is_group and not self.is_bot_mentioned(message): diff --git a/flanabot/bots/flana_bot.py b/flanabot/bots/flana_bot.py index c99b5fb..e2853e5 100644 --- a/flanabot/bots/flana_bot.py +++ b/flanabot/bots/flana_bot.py @@ -263,12 +263,12 @@ class FlanaBot(Connect4Bot, PenaltyBot, PollBot, ScraperBot, WeatherBot, MultiBo self.distribute_buttons(options, vertically=True), message, buttons_key=ButtonsGroup.ROLES, - data={'user_id': message.author.id} + buttons_data={'user_id': message.author.id} ) async def _on_roles_button_press(self, message: Message): await self.accept_button_event(message) - if message.buttons_info.presser_user.id != message.data['user_id']: + if message.buttons_info.presser_user.id != message.buttons_info.data['user_id']: return role = await self.find_role(message.buttons_info.pressed_text[1:].strip(), message) diff --git a/flanabot/bots/poll_bot.py b/flanabot/bots/poll_bot.py index cb661e4..3fb5fab 100644 --- a/flanabot/bots/poll_bot.py +++ b/flanabot/bots/poll_bot.py @@ -64,34 +64,25 @@ class PollBot(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.data.get('poll') is None: - return + @staticmethod + def _get_poll_message(message: Message) -> Message | None: + if (poll_message := message.replied_message) and poll_message.buttons_info and poll_message.buttons_info.key == ButtonsGroup.POLL: 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_score=multibot_constants.PARSE_CALLBACKS_MIN_SCORE_DEFAULT) - and - (poll_message := self.Message.find_one({'data.poll.is_active': True}, sort_keys=(('date', pymongo.DESCENDING),))) - ): - return await self.get_message(poll_message.chat.id, poll_message.id) async def _update_poll_buttons(self, message: Message): - if message.data['poll']['is_multiple_answer']: - total_votes = len({option_vote[0] for option_votes in message.data['poll']['votes'].values() if option_votes for option_vote in option_votes}) + if message.buttons_info.data['is_multiple_answer']: + total_votes = len({option_vote[0] for option_votes in message.buttons_info.data['votes'].values() if option_votes for option_vote in option_votes}) else: - total_votes = sum(len(option_votes) for option_votes in message.data['poll']['votes'].values()) + total_votes = sum(len(option_votes) for option_votes in message.buttons_info.data['votes'].values()) if total_votes: buttons = [] - for option, option_votes in message.data['poll']['votes'].items(): + for option, option_votes in message.buttons_info.data['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.data['poll']['votes'].keys()) + buttons = list(message.buttons_info.data['votes'].keys()) await self.edit(self.distribute_buttons(buttons, vertically=True), message) @@ -143,12 +134,12 @@ class PollBot(MultiBot, ABC): await self.delete_message(message) if all_: - for option_name, option_votes in poll_message.data['poll']['votes'].items(): - poll_message.data['poll']['votes'][option_name].clear() + for option_name, option_votes in poll_message.buttons_info.data['votes'].items(): + poll_message.buttons_info.data['votes'][option_name].clear() else: for user in await self._find_users_to_punish(message): - for option_name, option_votes in poll_message.data['poll']['votes'].items(): - poll_message.data['poll']['votes'][option_name] = [option_vote for option_vote in option_votes if option_vote[0] != user.id] + for option_name, option_votes in poll_message.buttons_info.data['votes'].items(): + poll_message.buttons_info.data['votes'][option_name] = [option_vote for option_vote in option_votes if option_vote[0] != user.id] await self._update_poll_buttons(poll_message) @@ -174,27 +165,26 @@ class PollBot(MultiBot, ABC): self.distribute_buttons(final_options, vertically=True), message, buttons_key=ButtonsGroup.POLL, - data={'poll': { + buttons_data={ '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_button_press(self, message: Message): await self.accept_button_event(message) - if not message.data['poll']['is_active']: + if not message.buttons_info.data['is_active']: return 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.data['poll']['banned_users_tries']: - message.data['poll']['banned_users_tries'][presser_id_str] += 1 - message.save() - if message.data['poll']['banned_users_tries'][presser_id_str] == 3: + if (presser_id_str := str(presser_id)) in message.buttons_info.data['banned_users_tries']: + message.buttons_info.data['banned_users_tries'][presser_id_str] += 1 + if message.buttons_info.data['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', @@ -206,13 +196,13 @@ class PollBot(MultiBot, ABC): return option_name = results[0] if (results := re.findall('(.*?) ➜.+', message.buttons_info.pressed_text)) else message.buttons_info.pressed_text - selected_option_votes = message.data['poll']['votes'][option_name] + selected_option_votes = message.buttons_info.data['votes'][option_name] if [presser_id, presser_name] in selected_option_votes: selected_option_votes.remove([presser_id, presser_name]) else: - if not message.data['poll']['is_multiple_answer']: - for option_votes in message.data['poll']['votes'].values(): + if not message.buttons_info.data['is_multiple_answer']: + for option_votes in message.buttons_info.data['votes'].values(): try: option_votes.remove([presser_id, presser_name]) except ValueError: @@ -232,7 +222,7 @@ class PollBot(MultiBot, ABC): winners = [] max_votes = 1 - for option, votes in poll_message.data['poll']['votes'].items(): + for option, votes in poll_message.buttons_info.data['votes'].items(): if len(votes) > max_votes: winners = [option] max_votes = len(votes) @@ -248,7 +238,7 @@ class PollBot(MultiBot, ABC): case _: text = 'Encuesta finalizada.' - poll_message.data['poll']['is_active'] = False + poll_message.buttons_info.data['is_active'] = False await self.edit(text, poll_message) if not message.replied_message: @@ -262,9 +252,8 @@ class PollBot(MultiBot, ABC): await self.delete_message(message) for user in await self._find_users_to_punish(message): - if str(user.id) not in poll_message.data['poll']['banned_users_tries']: - poll_message.data['poll']['banned_users_tries'][str(user.id)] = 0 - poll_message.save() + if str(user.id) not in poll_message.buttons_info.data['banned_users_tries']: + poll_message.buttons_info.data['banned_users_tries'][str(user.id)] = 0 @admin(send_negative=True) async def _on_voting_unban(self, message: Message): @@ -275,10 +264,9 @@ class PollBot(MultiBot, ABC): for user in await self._find_users_to_punish(message): try: - del poll_message.data['poll']['banned_users_tries'][str(user.id)] + del poll_message.buttons_info.data['banned_users_tries'][str(user.id)] except KeyError: pass - poll_message.save() # -------------------------------------------------------- # # -------------------- PUBLIC METHODS -------------------- # diff --git a/flanabot/bots/weather_bot.py b/flanabot/bots/weather_bot.py index af2acfb..7823b78 100644 --- a/flanabot/bots/weather_bot.py +++ b/flanabot/bots/weather_bot.py @@ -152,6 +152,7 @@ class WeatherBot(MultiBot, ABC): ], message, buttons_key=ButtonsGroup.WEATHER, + buttons_data={'weather_chart': weather_chart}, send_as_file=False ) await self.send_inline_results(message) @@ -159,36 +160,35 @@ class WeatherBot(MultiBot, ABC): if bot_state_message: await self.delete_message(bot_state_message) - if bot_message: - bot_message.weather_chart = weather_chart - bot_message.save() - if not self.is_bot_mentioned(message): - # noinspection PyTypeChecker - BotAction(Action.AUTO_WEATHER_CHART, message, affected_objects=[bot_message]).save() + if bot_message and not self.is_bot_mentioned(message): + # 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) + weather_chart = message.buttons_info.data['weather_chart'] + match message.buttons_info.pressed_text: case WeatherEmoji.ZOOM_IN.value: - message.weather_chart.zoom_in() + weather_chart.zoom_in() case WeatherEmoji.ZOOM_OUT.value: - message.weather_chart.zoom_out() + weather_chart.zoom_out() case WeatherEmoji.LEFT.value: - message.weather_chart.move_left() + weather_chart.move_left() case WeatherEmoji.RIGHT.value: - message.weather_chart.move_right() + 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 + weather_chart.trace_metadatas['rain_volume'].show = not weather_chart.trace_metadatas['rain_volume'].show + weather_chart.trace_metadatas['snow_volume'].show = not 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 + weather_chart.trace_metadatas[trace_metadata_name].show = not weather_chart.trace_metadatas[trace_metadata_name].show case _: return - message.weather_chart.apply_zoom() - message.weather_chart.draw() + weather_chart.apply_zoom() + weather_chart.draw() - image_bytes = message.weather_chart.to_image() + image_bytes = weather_chart.to_image() await self.edit(Media(image_bytes, MediaType.IMAGE, 'jpg'), message) diff --git a/flanabot/models/message.py b/flanabot/models/message.py index 74dee8f..0b7ddb0 100644 --- a/flanabot/models/message.py +++ b/flanabot/models/message.py @@ -8,7 +8,6 @@ from flanautils import Media, OrderedSet from multibot.models import Message as MultiBotMessage, User from flanabot.models.chat import Chat -from flanabot.models.weather_chart import WeatherChart @dataclass(eq=False) @@ -17,5 +16,4 @@ class Message(MultiBotMessage): mentions: list[User] = field(default_factory=list) chat: Chat = None replied_message: Message = None - weather_chart: WeatherChart = None song_infos: OrderedSet[Media] = field(default_factory=OrderedSet)