PyroFork: Add listen and ask

Co-authored-by: wulan17 <wulan17@nusantararom.org>
Signed-off-by: wulan17 <wulan17@nusantararom.org>

compiler/docs/compiler.py

@@ -144,6 +144,8 @@ def pyrogram_api():
             stop_transmission
             export_session_string
             set_parse_mode
+            ask
+            listen
         """,
         messages="""
         Messages

pyrogram/client.py

@@ -275,6 +275,8 @@ class Client(Methods):
 
         self.loop = asyncio.get_event_loop()
 
+        self.listening = {}
+
     def __enter__(self):
         return self.start()
 

pyrogram/handlers/message_handler.py

@@ -46,4 +46,27 @@ class MessageHandler(Handler):
     """
 
     def __init__(self, callback: Callable, filters=None):
-        super().__init__(callback, filters)
+        #super().__init__(callback, filters)
+        self.user_callback = callback
+        super().__init__(self.resolve_listener, filters)
+
+    async def resolve_listener(self, client, message, *args):
+        listener = client.listening.get(message.chat.id)
+        if listener and not listener['future'].done():
+            listener['future'].set_result(message)
+        else:
+            if listener and listener['future'].done():
+                client.clear_listener(message.chat.id, listener['future'])
+            await self.user_callback(client, message, *args)
+
+    async def check(self, client, update):
+        listener = client.listening.get(update.chat.id)
+
+        if listener and not listener['future'].done():
+            return await listener['filters'](client, update) if callable(listener['filters']) else True
+
+        return (
+            await self.filters(client, update)
+            if callable(self.filters)
+            else True
+        )

pyrogram/methods/utilities/__init__.py

@@ -17,7 +17,9 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from .add_handler import AddHandler
+from .ask import Ask
 from .export_session_string import ExportSessionString
+from .listen import Listen
 from .remove_handler import RemoveHandler
 from .restart import Restart
 from .run import Run
@@ -28,7 +30,9 @@ from .stop_transmission import StopTransmission
 
 class Utilities(
     AddHandler,
+    Ask,
     ExportSessionString,
+    Listen,
     RemoveHandler,
     Restart,
     Run,

pyrogram/methods/utilities/ask.py

@@ -0,0 +1,68 @@
+#  Pyrogram - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-present Dan <https://github.com/delivrance>
+#
+#  This file is part of Pyrogram.
+#
+#  Pyrogram is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU Lesser General Public License as published
+#  by the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  Pyrogram is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public License
+#  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+import logging
+
+import pyrogram
+
+from typing import Union
+
+log = logging.getLogger(__name__)
+
+
+class Ask:
+    async def ask(
+        self: "pyrogram.Client",
+        chat_id: Union[str, int],
+        text: str,
+        filters=None,
+        timeout: int = None,
+        *args,
+        **kwargs
+    ):
+        """Send message then awaits for a new message in the specified chat.
+
+        Parameters:
+            chat_id (``int`` | ``str``):
+                Unique identifier (int) or username (str) of the target chat.
+                For your personal cloud (Saved Messages) you can simply use "me" or "self".
+                For a contact that exists in your Telegram address book you can use his phone number (str).
+
+            text (``str``):
+                Text of the message to be sent.
+
+            filters (:obj:`~pyrogram.filters`, *optional*):
+                Pass one or more filters to allow only a subset of messages to be passed
+                in your function.
+
+            timeout (``int``, *optional*):
+                The waiting timeout.
+
+        Returns:
+            :obj:`~pyrogram.types.Message`: On success, text message is returned.
+
+        Example:
+            .. code-block:: python
+
+                answer = await Client.listen(chat_id, "Your name:")
+                name = answer.text
+        """
+        request = await self.send_message(chat_id, text, *args, **kwargs)
+        response = await self.listen(chat_id, filters, timeout)
+        response.request = request
+        return response

pyrogram/methods/utilities/listen.py

@@ -0,0 +1,92 @@
+#  Pyrogram - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-present Dan <https://github.com/delivrance>
+#
+#  This file is part of Pyrogram.
+#
+#  Pyrogram is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU Lesser General Public License as published
+#  by the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  Pyrogram is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public License
+#  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+import asyncio
+import functools
+import logging
+
+import pyrogram
+
+from typing import Union
+
+log = logging.getLogger(__name__)
+
+
+class Listen:
+    async def listen(
+        self: "pyrogram.Client",
+        chat_id: Union[str, int],
+        filters=None,
+        timeout: int = None
+    ):
+        """Awaits for a new message in the specified chat.
+
+        Parameters:
+            chat_id (``int`` | ``str``):
+                Unique identifier (int) or username (str) of the target chat.
+                For your personal cloud (Saved Messages) you can simply use "me" or "self".
+                For a contact that exists in your Telegram address book you can use his phone number (str).
+
+            filters (:obj:`~pyrogram.filters`, *optional*):
+                Pass one or more filters to allow only a subset of messages to be passed
+                in your function.
+
+            timeout (``int``, *optional*):
+                The waiting timeout.
+
+        Returns:
+            :obj:`~pyrogram.types.Message`: On success, text message is returned.
+
+        Example:
+            .. code-block:: python
+
+                await Client.send_message(chat_id, "Your name:")
+                answer = await Client.listen(chat_id)
+                name = answer.text
+        """
+        if type(chat_id) != int:
+            chat = await self.get_chat(chat_id)
+            chat_id = chat.id
+
+        future = self.loop.create_future()
+        future.add_done_callback(
+            functools.partial(self.clear_listener, chat_id)
+        )
+        self.listening.update({
+            chat_id: {"future": future, "filters": filters}
+        })
+        return await asyncio.wait_for(future, timeout)
+
+    def clear_listener(
+        self: "pyrogram.Client",
+        chat_id: Union[str, int],
+        future
+    ):
+        if chat_id in self.listening and future == self.listening[chat_id]["future"]:
+            self.listening.pop(chat_id, None)
+
+    def cancel_listener(
+        self: "pyrogram.Client",
+        chat_id: Union[str, int]
+    ):
+        listener = self.listening.get(chat_id)
+        if not listener or listener['future'].done():
+            return
+
+        listener['future'].set_exception(ListenerCanceled())
+        self.clear_listener(chat_id, listener['future'])

pyrogram/types/user_and_chats/chat.py

@@ -967,3 +967,12 @@ class Chat(Object):
         """
 
         return await self._client.unpin_all_chat_messages(self.id)
+
+    def listen(self, *args, **kwargs):
+        return self._client.listen(self.id, *args, **kwargs)
+
+    def ask(self, *args, **kwargs):
+        return self._client.ask(self.id, *args, **kwargs)
+
+    def cancel_listener(self):
+        return self._client.cancel_listener(self.id)

pyrogram/types/user_and_chats/user.py

@@ -397,3 +397,12 @@ class User(Object, Update):
         """
 
         return self._client.get_common_chats(self.id)
+
+    def listen(self, *args, **kwargs):
+        return self._client.listen(self.id, *args, **kwargs)
+
+    def ask(self, *args, **kwargs):
+        return self._client.ask(self.id, *args, **kwargs)
+
+    def cancel_listener(self):
+        return self._client.cancel_listener(self.id)