diff --git a/telethon/client/dialogs.py b/telethon/client/dialogs.py index 67cd52d0..1191a724 100644 --- a/telethon/client/dialogs.py +++ b/telethon/client/dialogs.py @@ -160,7 +160,7 @@ class DialogMethods(UserMethods): def conversation( self, entity, *, timeout=None, total_timeout=60, max_messages=100, - replies_are_responses=True): + exclusive=True, replies_are_responses=True): """ Creates a `Conversation ` with the given entity so you can easily send messages and await for @@ -187,6 +187,16 @@ class DialogMethods(UserMethods): specified chat, subsequent actions will result in ``ValueError``. + exclusive (`bool`, optional): + By default, conversations are exclusive within a single + chat. That means that while a conversation is open in a + chat, you can't open another one in the same chat, unless + you disable this flag. + + If you try opening an exclusive conversation for + a chat where it's already open, it will raise + ``AlreadyInConversationError``. + replies_are_responses (`bool`, optional): Whether replies should be treated as responses or not. @@ -229,6 +239,7 @@ class DialogMethods(UserMethods): timeout=timeout, total_timeout=total_timeout, max_messages=max_messages, + exclusive=exclusive, replies_are_responses=replies_are_responses ) diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py index 60f6269f..90f42cf1 100644 --- a/telethon/client/telegrambaseclient.py +++ b/telethon/client/telegrambaseclient.py @@ -272,6 +272,7 @@ class TelegramBaseClient(abc.ABC): # Some further state for subclasses self._event_builders = [] self._conversations = {} + self._ids_in_conversations = {} # chat_id: count # Default parse mode self._parse_mode = markdown diff --git a/telethon/client/updates.py b/telethon/client/updates.py index 1c212d21..84c07eb4 100644 --- a/telethon/client/updates.py +++ b/telethon/client/updates.py @@ -283,6 +283,10 @@ class UpdateMethods(UserMethods): try: await callback(event) + except errors.AlreadyInConversationError: + name = getattr(callback, '__name__', repr(callback)) + __log__.debug('Event handler "%s" already has an open ' + 'conversation, ignoring new one', name) except events.StopPropagation: name = getattr(callback, '__name__', repr(callback)) __log__.debug( diff --git a/telethon/errors/__init__.py b/telethon/errors/__init__.py index 03125dc5..faf88e8d 100644 --- a/telethon/errors/__init__.py +++ b/telethon/errors/__init__.py @@ -8,7 +8,8 @@ from threading import Thread from .common import ( ReadCancelledError, TypeNotFoundError, InvalidChecksumError, - InvalidBufferError, SecurityError, CdnFileTamperedError, MultiError + InvalidBufferError, SecurityError, CdnFileTamperedError, + AlreadyInConversationError, MultiError ) # This imports the base errors too, as they're imported there diff --git a/telethon/errors/common.py b/telethon/errors/common.py index 65cd363c..c9f22b41 100644 --- a/telethon/errors/common.py +++ b/telethon/errors/common.py @@ -79,6 +79,17 @@ class CdnFileTamperedError(SecurityError): ) +class AlreadyInConversationError(Exception): + """ + Occurs when another exclusive conversation is opened in the same chat. + """ + def __init__(self): + super().__init__( + 'Cannot open exclusive conversation in a ' + 'chat that already has one open conversation' + ) + + class MultiError(Exception): """Exception container for multiple `TLRequest`'s.""" diff --git a/telethon/tl/custom/conversation.py b/telethon/tl/custom/conversation.py index 66ec79f8..97b6816a 100644 --- a/telethon/tl/custom/conversation.py +++ b/telethon/tl/custom/conversation.py @@ -3,7 +3,7 @@ import itertools import time from .chatgetter import ChatGetter -from ... import utils +from ... import utils, errors class Conversation(ChatGetter): @@ -23,7 +23,7 @@ class Conversation(ChatGetter): def __init__(self, client, input_chat, *, timeout, total_timeout, max_messages, - replies_are_responses): + exclusive, replies_are_responses): self._id = Conversation._id_counter Conversation._id_counter += 1 @@ -50,6 +50,8 @@ class Conversation(ChatGetter): self._pending_edits = {} self._pending_reads = {} + self._exclusive = exclusive + # The user is able to expect two responses for the same message. # {desired message ID: next incoming index} self._response_indices = {} @@ -380,11 +382,20 @@ class Conversation(ChatGetter): return self._client.loop.run_until_complete(self.__aenter__()) async def __aenter__(self): - self._client._conversations[self._id] = self self._input_chat = \ await self._client.get_input_entity(self._input_chat) self._chat_peer = utils.get_peer(self._input_chat) + + # Make sure we're the only conversation in this chat if it's exclusive + chat_id = utils.get_peer_id(self._chat_peer) + count = self._client._ids_in_conversations.get(chat_id, 0) + if self._exclusive and count: + raise errors.AlreadyInConversationError() + + self._client._ids_in_conversations[chat_id] = count + 1 + self._client._conversations[self._id] = self + self._last_outgoing = 0 self._last_incoming = 0 for d in ( @@ -405,5 +416,11 @@ class Conversation(ChatGetter): return self._client.loop.run_until_complete(self.__aexit__(*args)) async def __aexit__(self, *args): + chat_id = utils.get_peer_id(self._chat_peer) + if self._client._ids_in_conversations[chat_id] == 1: + del self._client._ids_in_conversations[chat_id] + else: + self._client._ids_in_conversations[chat_id] -= 1 + del self._client._conversations[self._id] self._cancel_all()