Allow filtering new messages by sender/forwards

This commit is contained in:
Lonami Exo 2018-06-20 11:49:18 +02:00
parent ad1ab64415
commit c893eaaa84

View File

@ -1,6 +1,6 @@
import re
from .common import EventBuilder, EventCommon, name_inner_event
from .common import EventBuilder, EventCommon, name_inner_event, _into_id_set
from ..tl import types, custom
@ -18,14 +18,28 @@ class NewMessage(EventBuilder):
If set to ``True``, only **outgoing** messages will be handled.
Mutually exclusive with ``incoming`` (can only set one of either).
from_users (`entity`, optional):
Unlike `chats`, this parameter filters the *sender* of the message.
That is, only messages *sent by this user* will be handled. Use
`chats` if you want private messages with this/these users.
`from_users` lets you filter by messages sent by one or more users
across the desired chats.
forwards (`bool`, optional):
Whether forwarded messages should be handled or not. By default,
both forwarded and normal messages are included. If it's ``True``
*only* forwards will be handled. If it's ``False`` only messages
that are *not* forwards will be handled.
pattern (`str`, `callable`, `Pattern`, optional):
If set, only messages matching this pattern will be handled.
You can specify a regex-like string which will be matched
against the message, a callable function that returns ``True``
if a message is acceptable, or a compiled regex pattern.
"""
def __init__(self, incoming=None, outgoing=None,
chats=None, blacklist_chats=False, pattern=None):
def __init__(self, chats=None, *, blacklist_chats=False,
incoming=None, outgoing=None,
from_users=None, forwards=None, pattern=None):
if incoming is not None and outgoing is None:
outgoing = not incoming
elif outgoing is not None and incoming is None:
@ -40,6 +54,8 @@ class NewMessage(EventBuilder):
super().__init__(chats=chats, blacklist_chats=blacklist_chats)
self.incoming = incoming
self.outgoing = outgoing
self.from_users = from_users
self.forwards = forwards
if isinstance(pattern, str):
self.pattern = re.compile(pattern).match
elif not pattern or callable(pattern):
@ -49,6 +65,16 @@ class NewMessage(EventBuilder):
else:
raise TypeError('Invalid pattern type given')
# Should we short-circuit? E.g. perform no check at all
self._no_check = all(x is None for x in (
self.chats, self.incoming, self.outgoing,
self.from_users, self.forwards, self.from_users
))
async def resolve(self, client):
await super().resolve(client)
self.from_users = await _into_id_set(client, self.from_users)
def build(self, update):
if isinstance(update,
(types.UpdateNewMessage, types.UpdateNewChannelMessage)):
@ -102,15 +128,20 @@ class NewMessage(EventBuilder):
return self._message_filter_event(event)
def _message_filter_event(self, event):
# Short-circuit if we let pass all events
if all(x is None for x in (self.incoming, self.outgoing, self.chats,
self.pattern)):
if self._no_check:
return event
if self.incoming and event.message.out:
return
if self.outgoing and not event.message.out:
return
if self.forwards is not None:
if bool(self.forwards) != bool(event.message.fwd_from):
return
if self.from_users is not None:
if event.message.from_id not in self.from_users:
return
if self.pattern:
match = self.pattern(event.message.message or '')