NEW_MESSAGE_UPDATES = (

UpdateNewMessage,

UpdateNewChannelMessage,

UpdateNewScheduledMessage,

)

EDIT_MESSAGE_UPDATES = (UpdateEditMessage, UpdateEditChannelMessage)

DELETE_MESSAGES_UPDATES = (UpdateDeleteMessages, UpdateDeleteChannelMessages)

CALLBACK_QUERY_UPDATES = (UpdateBotCallbackQuery, UpdateInlineBotCallbackQuery)

CHAT_MEMBER_UPDATES = (UpdateChatParticipant, UpdateChannelParticipant)

USER_STATUS_UPDATES = (UpdateUserStatus,)

BOT_INLINE_QUERY_UPDATES = (UpdateBotInlineQuery,)

POLL_UPDATES = (UpdateMessagePoll,)

CHOSEN_INLINE_RESULT_UPDATES = (UpdateBotInlineSend,)

CHAT_JOIN_REQUEST_UPDATES = (UpdateBotChatInviteRequester,)

def __init__(self, client: "pyrogram.Client"):

self.client = client

self.loop = asyncio.get_event_loop()

self.handler_worker_tasks = []

self.locks_list = []

self.updates_queue = asyncio.Queue()

self.groups = OrderedDict()

async def message_parser(update, users, chats):

return (

await pyrogram.types.Message._parse(

self.client,

update.message,

users,

chats,

isinstance(update, UpdateNewScheduledMessage),

),

MessageHandler,

)

async def edited_message_parser(update, users, chats):

# Edited messages are parsed the same way as new messages, but the handler is different

parsed, _ = await message_parser(update, users, chats)

return (parsed, EditedMessageHandler)

async def deleted_messages_parser(update, users, chats):

return (

utils.parse_deleted_messages(self.client, update),

DeletedMessagesHandler,

)

async def callback_query_parser(update, users, chats):

return (

await pyrogram.types.CallbackQuery._parse(self.client, update, users),

CallbackQueryHandler,

)

async def user_status_parser(update, users, chats):

return (

pyrogram.types.User._parse_user_status(self.client, update),

UserStatusHandler,

)

async def inline_query_parser(update, users, chats):

return (

pyrogram.types.InlineQuery._parse(self.client, update, users),

InlineQueryHandler,

)

async def poll_parser(update, users, chats):

return (pyrogram.types.Poll._parse_update(self.client, update), PollHandler)

async def chosen_inline_result_parser(update, users, chats):

return (

pyrogram.types.ChosenInlineResult._parse(self.client, update, users),

ChosenInlineResultHandler,

)

async def chat_member_updated_parser(update, users, chats):

return (

pyrogram.types.ChatMemberUpdated._parse(

self.client, update, users, chats

),

ChatMemberUpdatedHandler,

)

async def chat_join_request_parser(update, users, chats):

return (

pyrogram.types.ChatJoinRequest._parse(

self.client, update, users, chats

),

ChatJoinRequestHandler,

)

self.update_parsers = {

Dispatcher.NEW_MESSAGE_UPDATES: message_parser,

Dispatcher.EDIT_MESSAGE_UPDATES: edited_message_parser,

Dispatcher.DELETE_MESSAGES_UPDATES: deleted_messages_parser,

Dispatcher.CALLBACK_QUERY_UPDATES: callback_query_parser,

Dispatcher.USER_STATUS_UPDATES: user_status_parser,

Dispatcher.BOT_INLINE_QUERY_UPDATES: inline_query_parser,

Dispatcher.POLL_UPDATES: poll_parser,

Dispatcher.CHOSEN_INLINE_RESULT_UPDATES: chosen_inline_result_parser,

Dispatcher.CHAT_MEMBER_UPDATES: chat_member_updated_parser,

Dispatcher.CHAT_JOIN_REQUEST_UPDATES: chat_join_request_parser,

}

self.update_parsers = {

key: value

for key_tuple, value in self.update_parsers.items()

for key in key_tuple

}

async def start(self):

if not self.client.no_updates:

for i in range(self.client.workers):

self.locks_list.append(asyncio.Lock())

self.handler_worker_tasks.append(

self.loop.create_task(self.handler_worker(self.locks_list[-1]))

)

log.info("Started %s HandlerTasks", self.client.workers)

async def stop(self):

if not self.client.no_updates:

for i in range(self.client.workers):

self.updates_queue.put_nowait(None)

for i in self.handler_worker_tasks:

await i

self.handler_worker_tasks.clear()

self.groups.clear()

log.info("Stopped %s HandlerTasks", self.client.workers)

def add_handler(self, handler, group: int):

async def fn():

for lock in self.locks_list:

await lock.acquire()

try:

if group not in self.groups:

self.groups[group] = []

self.groups = OrderedDict(sorted(self.groups.items()))

self.groups[group].append(handler)

finally:

for lock in self.locks_list:

lock.release()

self.loop.create_task(fn())

def remove_handler(self, handler, group: int):

async def fn():

for lock in self.locks_list:

await lock.acquire()

try:

if group not in self.groups:

raise ValueError(

f"Group {group} does not exist. Handler was not removed."

)

self.groups[group].remove(handler)

finally:

for lock in self.locks_list:

lock.release()

self.loop.create_task(fn())

async def handler_worker(self, lock):

while True:

packet = await self.updates_queue.get()

if packet is None:

break

try:

update, users, chats = packet

parser = self.update_parsers.get(type(update), None)

parsed_update, handler_type = (

await parser(update, users, chats)

if parser is not None

else (None, type(None))

)

async with lock:

for group in self.groups.values():

for handler in group:

args = None

if isinstance(handler, handler_type):

try:

if await handler.check(self.client, parsed_update):

args = (parsed_update,)

except Exception as e:

log.exception(e)

continue

elif isinstance(handler, RawUpdateHandler):

args = (update, users, chats)

if args is None:

continue

try:

if inspect.iscoroutinefunction(handler.callback):

await handler.callback(self.client, *args)

else:

await self.loop.run_in_executor(

self.client.executor,

handler.callback,

self.client,

*args,

)

except pyrogram.StopPropagation:

raise

except pyrogram.ContinuePropagation:

continue

except Exception as e:

log.exception(e)

break

except pyrogram.StopPropagation:

pass

except Exception as e:

log.exception(e)

async def _process_diff(self, state: tuple) -> tuple[int, int] | None:

id, local_pts, local_qts, local_date, local_seq = state

try:

diff = await self.client.invoke(

raw.functions.updates.GetChannelDifference(

channel=await self.client.resolve_peer(id),

filter=raw.types.ChannelMessagesFilterEmpty(),

pts=local_pts,

limit=10_000,

)

if id != 0

else raw.functions.updates.GetDifference(

pts=local_pts, date=local_date, qts=0

)

)

except (errors.ChannelPrivate, errors.ChannelInvalid):

await self.client.storage.delete_state(id)

return

if isinstance(diff, raw.types.updates.DifferenceEmpty):

return

elif isinstance(diff, raw.types.updates.DifferenceTooLong):

return

elif isinstance(diff, raw.types.updates.ChannelDifferenceEmpty):

return

elif isinstance(diff, raw.types.updates.ChannelDifferenceTooLong):

return

elif isinstance(diff, raw.types.updates.Difference):

local_pts = diff.state.pts

elif isinstance(diff, raw.types.updates.ChannelDifference):

local_qts = diff.pts

elif isinstance(diff, raw.types.updates.DifferenceSlice):

if diff.intermediate_state.pts == local_pts:

return

local_pts = diff.intermediate_state.pts

local_date = diff.intermediate_state.date

await self._process_diff((id, local_pts, local_qts, local_date, local_seq))

return

users = {i.id: i for i in diff.users}

chats = {i.id: i for i in diff.chats}

for message in diff.new_messages:

self.updates_queue.put_nowait(

(

raw.types.UpdateNewMessage(

message=message, pts=local_pts, pts_count=-1

)

if id == 0

else raw.types.UpdateNewChannelMessage(

message=message, pts=local_pts, pts_count=-1

),

users,

chats,

)

)

for update in diff.other_updates:

self.updates_queue.put_nowait((update, users, chats))

await self.client.storage.update_state(

(id, local_pts, local_qts, local_date, local_seq)

)

return len(diff.new_messages), len(diff.other_updates)

async def fill_difference_gap(self):

states = await self.client.storage.get_states()

messages_count = updates_count = 0

for state in states:

result = await self._process_diff(state)

if result is not None:

m, u = result

messages_count += m

updates_count += u