From 27323fb78ec1cf402a9a9119c2415aed61502a6a Mon Sep 17 00:00:00 2001
From: Anton <m@antonz.org>
Date: Sat, 13 Apr 2024 07:21:57 +0500
Subject: [PATCH] fix: current user is in update, not in message

---
 bot/bot.py              | 16 ++++++++++------
 bot/commands/imagine.py |  4 +++-
 bot/commands/message.py |  2 +-
 bot/commands/retry.py   |  4 +++-
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/bot/bot.py b/bot/bot.py
index a8aa1d6..afb2d8d 100644
--- a/bot/bot.py
+++ b/bot/bot.py
@@ -5,7 +5,7 @@
 import textwrap
 import time
 
-from telegram import Chat, Message
+from telegram import Chat, Message, Update
 from telegram.ext import (
     Application,
     ApplicationBuilder,
@@ -112,8 +112,10 @@ async def post_shutdown(application: Application) -> None:
 def with_message_limit(func):
     """Refuses to reply if the user has exceeded the message limit."""
 
-    async def wrapper(message: Message, context: CallbackContext, question: str) -> None:
-        username = message.from_user.username
+    async def wrapper(
+        update: Update, message: Message, context: CallbackContext, question: str
+    ) -> None:
+        username = update.effective_user.username
         user = UserData(context.user_data)
 
         # check if the message counter exceeds the message limit
@@ -129,17 +131,19 @@ async def wrapper(message: Message, context: CallbackContext, question: str) ->
 
         # this is a known user or they have not exceeded the message limit,
         # so proceed to the actual message handler
-        await func(message, context, question)
+        await func(update=update, message=message, context=context, question=question)
 
         # increment the message counter
         message_count = user.message_counter.increment()
-        logger.debug(f"user={message.from_user.username}, n_messages={message_count}")
+        logger.debug(f"user={username}, n_messages={message_count}")
 
     return wrapper
 
 
 @with_message_limit
-async def reply_to(message: Message, context: CallbackContext, question: str) -> None:
+async def reply_to(
+    update: Update, message: Message, context: CallbackContext, question: str
+) -> None:
     """Replies to a specific question."""
     await message.chat.send_action(action="typing", message_thread_id=message.message_thread_id)
 
diff --git a/bot/commands/imagine.py b/bot/commands/imagine.py
index 9f917ad..a7d2f43 100644
--- a/bot/commands/imagine.py
+++ b/bot/commands/imagine.py
@@ -38,4 +38,6 @@ async def __call__(self, update: Update, context: CallbackContext) -> None:
                 parse_mode=ParseMode.HTML,
             )
             return
-        await self.reply_func(update.message, context, question=message.text)
+        await self.reply_func(
+            update=update, message=update.message, context=context, question=message.text
+        )
diff --git a/bot/commands/message.py b/bot/commands/message.py
index 8da82b6..251b5db 100644
--- a/bot/commands/message.py
+++ b/bot/commands/message.py
@@ -30,4 +30,4 @@ async def __call__(self, update: Update, context: CallbackContext) -> None:
             # this is not a question to the bot, so ignore it
             return
 
-        await self.reply_func(message, context, question=question)
+        await self.reply_func(update=update, message=message, context=context, question=question)
diff --git a/bot/commands/retry.py b/bot/commands/retry.py
index 1e26004..04c1493 100644
--- a/bot/commands/retry.py
+++ b/bot/commands/retry.py
@@ -18,4 +18,6 @@ async def __call__(self, update: Update, context: CallbackContext) -> None:
         if not last_message:
             await update.message.reply_text("No message to retry 🤷‍♂️")
             return
-        await self.reply_func(update.message, context, question=last_message.question)
+        await self.reply_func(
+            update=update, message=update.message, context=context, question=last_message.question
+        )