Dry run for rndinfo command

.gitignore

@@ -79,6 +79,9 @@ target/
 # Jupyter Notebook
 .ipynb_checkpoints
 
+# Visual Studio Code
+.vscode/
+
 # IPython
 profile_default/
 ipython_config.py

locales/evgfilim1-userbot.pot

@@ -12,7 +12,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: evgfilim1/userbot 0.6.x\n"
 "Report-Msgid-Bugs-To: https://github.com/evgfilim1/userbot/issues\n"
-"POT-Creation-Date: 2023-04-20 00:46+0500\n"
+"POT-Creation-Date: 2023-05-08 19:59+0300\n"
 "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language-Team: LANGUAGE <LL@li.org>\n"
@@ -282,17 +282,17 @@ msgstr ""
 msgid "{icon} <code>{value}</code> has been invited to the chat"
 msgstr ""
 
-#: userbot/commands/chat_info.py:95
+#: userbot/commands/chat_info.py:118
 #, python-brace-format
 msgid "{icon} <b>New chat avatar was set!</b> <a href='{msg_link}'>Source</a>\n"
 msgstr ""
 
-#: userbot/commands/chat_info.py:101
+#: userbot/commands/chat_info.py:127
 #, python-brace-format
 msgid "{icon} <b>New chat title was set!</b> <a href='{msg_link}'>Source</a>"
 msgstr ""
 
-#: userbot/commands/chat_info.py:113
+#: userbot/commands/chat_info.py:139
 msgid "<a href='{msg.link}'>Random message (#{msg.id})</a>"
 msgstr ""
 

userbot/commands/chat_info.py

@@ -20,23 +20,34 @@
 
 _log = logging.getLogger(__name__)
 
+ChatTitle = str
+
 
 async def _get_random_message(
     chat_id: int,
     message_filter: MessagesFilter,
     client: Client,
 ) -> Message:
+    """Gets a random message"""
     total = await client.search_messages_count(chat_id, filter=message_filter)
     offset = randint(0, total - 1)
     gen = client.search_messages(chat_id, filter=message_filter, offset=offset, limit=1)
     return await gen.__anext__()
 
 
+async def _get_random_photo_message(
+    chat_id: int,
+    client: Client,
+) -> Message:
+    """Gets a random photo message"""
+    return await _get_random_message(chat_id, MessagesFilter.PHOTO, client)
+
+
 async def _set_random_chat_photo(chat_id: int, client: Client) -> Message:
     """Sets a random chat photo."""
     retries = 20
     while retries >= 0:
-        message = await _get_random_message(chat_id, MessagesFilter.PHOTO, client)
+        message = await _get_random_photo_message(chat_id, client)
         try:
             await client.set_chat_photo(chat_id, photo=message.photo.file_id)
         except PhotoCropSizeSmall:
@@ -49,34 +60,37 @@ async def _set_random_chat_photo(chat_id: int, client: Client) -> Message:
         raise RuntimeError("Retries exceeded")
 
 
-async def _set_random_chat_title(chat_id: int, client: Client) -> Message:
-    """Sets a random chat title."""
+async def _get_random_chat_title(chat_id: int, client: Client) -> tuple[Message, ChatTitle]:
+    """Gets a random chat title."""
     retries = 50
     chat = await client.get_chat(chat_id)
     title_prefix = chat.title.split(" — ")[0]
     while retries >= 0:
         message = await _get_random_message(chat_id, MessagesFilter.EMPTY, client)
-        try:
-            if message.text is None:
-                continue
+        if message.text is not None:
             stripped_text = message.text[: 128 - len(title_prefix) - 3]
-            await client.set_chat_title(chat_id, f"{title_prefix} — {stripped_text}")
-        except BadRequest as e:
-            _log.warning(
-                "An error occurred while setting the chat title in %d",
-                chat_id,
-                exc_info=e,
-            )
-            continue
-        else:
-            return message
-        finally:
-            retries -= 1
-    else:
-        raise RuntimeError("Retries exceeded")
+            chat_title = f"{title_prefix} — {stripped_text}"
+
+            return message, chat_title
+        retries -= 1
+    raise RuntimeError("Retries exceeded")
+
+
+async def _set_random_chat_title(chat_id: int, client: Client) -> Message:
+    """Sets a random chat title."""
+    message, chat_title = await _get_random_chat_title(chat_id, client)
+    try:
+        await client.set_chat_title(chat_id, chat_title)
+    except BadRequest as e:
+        _log.warning(
+            "An error occurred while setting the chat title in %d",
+            chat_id,
+            exc_info=e,
+        )
+    return message
 
 
-@commands.add("rndinfo", usage="['photo'|'title']")
+@commands.add("rndinfo", usage="['photo'|'title'|'all'] ['dry-run']")
 async def random_chat_info(
     client: Client,
     message: Message,
@@ -86,18 +100,30 @@ async def random_chat_info(
     """Sets random chat photo and/or title.
 
     Sets both if no argument is given.
+    Specify `dry-run` as the second argument to get the message without actually setting it.
     """
     _ = tr.gettext
     text = ""
     what = command.args[0]
-    if what == "photo" or what is None:
-        msg = await _set_random_chat_photo(message.chat.id, client)
+    if what is None:
+        what = "all"
+        is_dry_run = False
+    else:
+        is_dry_run = command.args[1] == "dry-run"
+    if what in {"photo", "all"}:
+        if is_dry_run:
+            msg = await _get_random_photo_message(message.chat.id, client)
+        else:
+            msg = await _set_random_chat_photo(message.chat.id, client)
         text += _(
             "{icon} <b>New chat avatar was set!</b> <a href='{msg_link}'>Source</a>\n"
         ).format(icon=Icons.PICTURE, msg_link=msg.link)
         await sleep(0.1)
-    if what == "title" or what is None:
-        msg = await _set_random_chat_title(message.chat.id, client)
+    if what in {"title", "all"}:
+        if is_dry_run:
+            msg, _chat_title = await _get_random_chat_title(message.chat.id, client)
+        else:
+            msg = await _set_random_chat_title(message.chat.id, client)
         text += _("{icon} <b>New chat title was set!</b> <a href='{msg_link}'>Source</a>").format(
             icon=Icons.PENCIL, msg_link=msg.link
         )