Merge branch 'develop' into inline-mode

This commit is contained in:
Dan 2018-11-09 10:23:47 +01:00
commit 755a392fa2
2 changed files with 85 additions and 161 deletions

View File

@ -22,10 +22,9 @@ from collections import OrderedDict
from queue import Queue from queue import Queue
from threading import Thread from threading import Thread
import pyrogram
from pyrogram.api import types from pyrogram.api import types
from ..ext import utils from ..ext import utils
from ..handlers import RawUpdateHandler, CallbackQueryHandler, MessageHandler, DeletedMessagesHandler, UserStatusHandler from ..handlers import CallbackQueryHandler, MessageHandler, DeletedMessagesHandler, UserStatusHandler, RawUpdateHandler
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -46,15 +45,39 @@ class Dispatcher:
types.UpdateDeleteChannelMessages types.UpdateDeleteChannelMessages
) )
CALLBACK_QUERY_UPDATES = (
types.UpdateBotCallbackQuery,
types.UpdateInlineBotCallbackQuery
)
MESSAGE_UPDATES = NEW_MESSAGE_UPDATES + EDIT_MESSAGE_UPDATES MESSAGE_UPDATES = NEW_MESSAGE_UPDATES + EDIT_MESSAGE_UPDATES
def __init__(self, client, workers): UPDATES = None
def __init__(self, client, workers: int):
self.client = client self.client = client
self.workers = workers self.workers = workers
self.workers_list = [] self.workers_list = []
self.updates = Queue() self.updates = Queue()
self.groups = OrderedDict() self.groups = OrderedDict()
Dispatcher.UPDATES = {
Dispatcher.MESSAGE_UPDATES:
lambda upd, usr, cht: (utils.parse_messages(self.client, upd.message, usr, cht), MessageHandler),
Dispatcher.DELETE_MESSAGE_UPDATES:
lambda upd, usr, cht: (utils.parse_deleted_messages(upd), DeletedMessagesHandler),
Dispatcher.CALLBACK_QUERY_UPDATES:
lambda upd, usr, cht: (utils.parse_callback_query(self.client, upd, usr), CallbackQueryHandler),
(types.UpdateUserStatus,):
lambda upd, usr, cht: (utils.parse_user_status(upd.status, upd.user_id), UserStatusHandler)
}
Dispatcher.UPDATES = {key: value for key_tuple, value in Dispatcher.UPDATES.items() for key in key_tuple}
def start(self): def start(self):
for i in range(self.workers): for i in range(self.workers):
self.workers_list.append( self.workers_list.append(
@ -70,8 +93,8 @@ class Dispatcher:
for _ in range(self.workers): for _ in range(self.workers):
self.updates.put(None) self.updates.put(None)
for i in self.workers_list: for worker in self.workers_list:
i.join() worker.join()
self.workers_list.clear() self.workers_list.clear()
@ -84,60 +107,10 @@ class Dispatcher:
def remove_handler(self, handler, group: int): def remove_handler(self, handler, group: int):
if group not in self.groups: if group not in self.groups:
raise ValueError("Group {} does not exist. " raise ValueError("Group {} does not exist. Handler was not removed.".format(group))
"Handler was not removed.".format(group))
self.groups[group].remove(handler) self.groups[group].remove(handler)
def dispatch(self, update, users: dict = None, chats: dict = None, is_raw: bool = False):
for group in self.groups.values():
try:
for handler in group:
if is_raw:
if not isinstance(handler, RawUpdateHandler):
continue
args = (self.client, update, users, chats)
else:
message = (update.message
or update.channel_post
or update.edited_message
or update.edited_channel_post)
deleted_messages = (update.deleted_channel_posts
or update.deleted_messages)
callback_query = update.callback_query
user_status = update.user_status
if message and isinstance(handler, MessageHandler):
if not handler.check(message):
continue
args = (self.client, message)
elif deleted_messages and isinstance(handler, DeletedMessagesHandler):
if not handler.check(deleted_messages):
continue
args = (self.client, deleted_messages)
elif callback_query and isinstance(handler, CallbackQueryHandler):
if not handler.check(callback_query):
continue
args = (self.client, callback_query)
elif user_status and isinstance(handler, UserStatusHandler):
if not handler.check(user_status):
continue
args = (self.client, user_status)
else:
continue
handler.callback(*args)
break
except Exception as e:
log.error(e, exc_info=True)
def update_worker(self): def update_worker(self):
name = threading.current_thread().name name = threading.current_thread().name
log.debug("{} started".format(name)) log.debug("{} started".format(name))
@ -153,79 +126,32 @@ class Dispatcher:
chats = {i.id: i for i in update[2]} chats = {i.id: i for i in update[2]}
update = update[0] update = update[0]
self.dispatch(update, users=users, chats=chats, is_raw=True) parser = Dispatcher.UPDATES.get(type(update), None)
if isinstance(update, Dispatcher.MESSAGE_UPDATES): if parser is None:
if isinstance(update.message, types.MessageEmpty):
continue
message = utils.parse_messages(
self.client,
update.message,
users,
chats
)
is_edited_message = isinstance(update, Dispatcher.EDIT_MESSAGE_UPDATES)
self.dispatch(
pyrogram.Update(
message=((message if message.chat.type != "channel"
else None) if not is_edited_message
else None),
edited_message=((message if message.chat.type != "channel"
else None) if is_edited_message
else None),
channel_post=((message if message.chat.type == "channel"
else None) if not is_edited_message
else None),
edited_channel_post=((message if message.chat.type == "channel"
else None) if is_edited_message
else None)
)
)
elif isinstance(update, Dispatcher.DELETE_MESSAGE_UPDATES):
is_channel = hasattr(update, 'channel_id')
messages = utils.parse_deleted_messages(
update.messages,
(update.channel_id if is_channel else None)
)
self.dispatch(
pyrogram.Update(
deleted_messages=(messages if not is_channel else None),
deleted_channel_posts=(messages if is_channel else None)
)
)
elif isinstance(update, types.UpdateBotCallbackQuery):
self.dispatch(
pyrogram.Update(
callback_query=utils.parse_callback_query(
self.client, update, users
)
)
)
elif isinstance(update, types.UpdateInlineBotCallbackQuery):
self.dispatch(
pyrogram.Update(
callback_query=utils.parse_inline_callback_query(
self.client, update, users
)
)
)
elif isinstance(update, types.UpdateUserStatus):
self.dispatch(
pyrogram.Update(
user_status=utils.parse_user_status(
update.status, update.user_id
)
)
)
else:
continue continue
update, handler_type = parser(update, users, chats)
for group in self.groups.values():
for handler in group:
args = None
if isinstance(handler, RawUpdateHandler):
args = (update, users, chats)
elif isinstance(handler, handler_type):
if handler.check(update):
args = (update,)
if args is None:
continue
try:
handler.callback(self.client, *args)
except Exception as e:
log.error(e, exc_info=True)
finally:
break
except Exception as e: except Exception as e:
log.error(e, exc_info=True) log.error(e, exc_info=True)

View File

@ -760,10 +760,10 @@ def parse_messages(
return parsed_messages if is_list else parsed_messages[0] return parsed_messages if is_list else parsed_messages[0]
def parse_deleted_messages( def parse_deleted_messages(update) -> pyrogram_types.Messages:
messages: list, messages = update.messages
channel_id: int channel_id = getattr(update, "channel_id", None)
) -> pyrogram_types.Messages:
parsed_messages = [] parsed_messages = []
for message in messages: for message in messages:
@ -870,42 +870,40 @@ def parse_profile_photos(photos):
) )
def parse_callback_query(client, callback_query, users): def parse_callback_query(client, update, users):
peer = callback_query.peer message = None
inline_message_id = None
if isinstance(peer, types.PeerUser): if isinstance(update, types.UpdateBotCallbackQuery):
peer_id = peer.user_id peer = update.peer
elif isinstance(peer, types.PeerChat):
peer_id = -peer.chat_id
else:
peer_id = int("-100" + str(peer.channel_id))
return pyrogram_types.CallbackQuery( if isinstance(peer, types.PeerUser):
id=str(callback_query.query_id), peer_id = peer.user_id
from_user=parse_user(users[callback_query.user_id]), elif isinstance(peer, types.PeerChat):
message=client.get_messages(peer_id, callback_query.msg_id), peer_id = -peer.chat_id
chat_instance=str(callback_query.chat_instance), else:
data=callback_query.data.decode(), peer_id = int("-100" + str(peer.channel_id))
game_short_name=callback_query.game_short_name,
client=client
)
message = client.get_messages(peer_id, update.msg_id)
def parse_inline_callback_query(client, callback_query, users): elif isinstance(update, types.UpdateInlineBotCallbackQuery):
return pyrogram_types.CallbackQuery( inline_message_id = b64encode(
id=str(callback_query.query_id),
from_user=parse_user(users[callback_query.user_id]),
chat_instance=str(callback_query.chat_instance),
inline_message_id=b64encode(
pack( pack(
"<iqq", "<iqq",
callback_query.msg_id.dc_id, update.msg_id.dc_id,
callback_query.msg_id.id, update.msg_id.id,
callback_query.msg_id.access_hash update.msg_id.access_hash
), ),
b"-_" b"-_"
).decode().rstrip("="), ).decode().rstrip("=")
game_short_name=callback_query.game_short_name,
return pyrogram_types.CallbackQuery(
id=str(update.query_id),
from_user=parse_user(users[update.user_id]),
message=message,
inline_message_id=inline_message_id,
chat_instance=str(update.chat_instance),
data=update.data.decode(),
game_short_name=update.game_short_name,
client=client client=client
) )