8000 Add Filters.via_bot (#2009) · Konano/python-telegram-bot@0189442 · GitHub
[go: up one dir, main page]
Skip to content
8000

Commit 0189442

Browse files
PoolitzerBibo-Joshi
Poolitzer
and
Bibo-Joshi
authored
Add Filters.via_bot (python-telegram-bot#2009)
* feat: via_bot filter also fixing a small mistake in the empty parameter of the user filter and improve docs slightly * fix: forgot to set via_bot to None * fix: redoing subclassing to copy paste solution * Cosmetic changes Co-authored-by: Hinrich Mahler <hinrich.mahler@freenet.de>
1 parent fd0325f commit 0189442

File tree

2 files changed

+292
-4
lines changed

2 files changed

+292
-4
lines changed

‎telegram/ext/filters.py

Lines changed: 166 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -887,7 +887,8 @@ def filter(self, message):
887887
"""Messages sent in a group chat."""
888888

889889
class user(BaseFilter):
890-
"""Filters messages to allow only those which are from specified user ID.
890+
"""Filters messages to allow only those which are from specified user ID(s) or
891+
username(s).
891892
892893
Examples:
893894
``MessageHandler(Filters.user(1234), callback_method)``
@@ -919,7 +920,6 @@ class user(BaseFilter):
919920
RuntimeError: If user_id and username are both present.
920921
921922
"""
922-
923923
def __init__(self, user_id=None, username=None, allow_empty=False):
924924
self.allow_empty = allow_empty
925925
self.__lock = Lock()
@@ -1053,8 +1053,171 @@ def filter(self, message):
10531053
return self.allow_empty
10541054
return False
10551055

1056+
class via_bot(BaseFilter):
1057+
"""Filters messages to allow only those which are from specified via_bot ID(s) or
1058+
username(s).
1059+
1060+
Examples:
1061+
``MessageHandler(Filters.via_bot(1234), callback_method)``
1062+
1063+
Warning:
1064+
:attr:`bot_ids` will give a *copy* of the saved bot ids as :class:`frozenset`. This
1065+
is to ensure thread safety. To add/remove a bot, you should use :meth:`add_usernames`,
1066+
:meth:`add_bot_ids`, :meth:`remove_usernames` and :meth:`remove_bot_ids`. Only update
1067+
the entire set by ``filter.bot_ids/usernames = new_set``, if you are entirely sure
1068+
that it is not causing race conditions, as this will complete replace the current set
1069+
of allowed bots.
1070+
1071+
Attributes:
1072+
bot_ids(set(:obj:`int`), optional): Which bot ID(s) to allow through.
1073+
usernames(set(:obj:`str`), optional): Which username(s) (without leading '@') to allow
1074+
through.
1075+
allow_empty(:obj:`bool`, optional): Whether updates should be processed, if no bot
1076+
is specified in :attr:`bot_ids` and :attr:`usernames`.
1077+
1078+
Args:
1079+
bot_id(:obj:`int` | List[:obj:`int`], optional): Which bot ID(s) to allow
1080+
through.
1081+
username(:obj:`str` | List[:obj:`str`], optional): Which username(s) to allow
1082+
through. Leading '@'s in usernames will be discarded.
1083+
allow_empty(:obj:`bool`, optional): Whether updates should be processed, if no user
1084+
is specified in :attr:`bot_ids` and :attr:`usernames`. Defaults to :obj:`False`
1085+
1086+
Raises:
1087+
RuntimeError: If bot_id and username are both present.
1088+
"""
1089+
1090+
def __init__(self, bot_id=None, username=None, allow_empty=False):
1091+
self.allow_empty = allow_empty
1092+
self.__lock = Lock()
1093+
1094+
self._bot_ids = set()
1095+
self._usernames = set()
1096+
1097+
self._set_bot_ids(bot_id)
1098+
self._set_usernames(username)
1099+
1100+
@staticmethod
1101+
def _parse_bot_id(bot_id):
1102+
if bot_id is None:
1103+
return set()
1104+
if isinstance(bot_id, int):
1105+
return {bot_id}
1106+
return set(bot_id)
1107+
1108+
@staticmethod
1109+
def _parse_username(username):
1110+
if username is None:
1111+
return set()
1112+
if isinstance(username, str):
1113+
return {username[1:] if username.startswith('@') else username}
1114+
return {bot[1:] if bot.startswith('@') else bot for bot in username}
1115+
1116+
def _set_bot_ids(self, bot_id):
1117+
with self.__lock:
1118+
if bot_id and self._usernames:
1119+
raise RuntimeError("Can't set bot_id in conjunction with (already set) "
1120+
"usernames.")
1121+
self._bot_ids = self._parse_bot_id(bot_id)
1122+
1123+
def _set_usernames(self, username):
1124+
with self.__lock:
1125+
if username and self._bot_ids:
1126+
raise RuntimeError("Can't set username in conjunction with (already set) "
1127+
"bot_ids.")
1128+
self._usernames = self._parse_username(username)
1129+
1130+
@property
1131+
def bot_ids(self):
1132+
with self.__lock:
1133+
return frozenset(self._bot_ids)
1134+
1135+
@bot_ids.setter
1136+
def bot_ids(self, bot_id):
1137+
self._set_bot_ids(bot_id)
1138+
1139+
@property
1140+
def usernames(self):
1141+
with self.__lock:
1142+
return frozenset(self._usernames)
1143+
1144+
@usernames.setter
1145+
def usernames(self, username):
1146+
self._set_usernames(username)
1147+
1148+
def add_usernames(self, username):
1149+
"""
1150+
Add one or more users to the allowed usernames.
1151+
Args:
1152+
username(:obj:`str` | List[:obj:`str`], optional): Which username(s) to allow
1153+
through. Leading '@'s in usernames will be discarded.
1154+
"""
1155+
with self.__lock:
1156+
if self._bot_ids:
1157+
raise RuntimeError("Can't set username in conjunction with (already set) "
1158+
"bot_ids.")
1159+
1160+
username = self._parse_username(username)
1161+
self._usernames |= username
1162+
1163+
def add_bot_ids(self, bot_id):
1164+
"""
1165+
Add one or more users to the allowed user ids.
1166+
Args:
1167+
bot_id(:obj:`int` | List[:obj:`int`], optional): Which bot ID(s) to allow
1168+
through.
1169+
"""
1170+
with self.__lock:
1171+
if self._usernames:
1172+
raise RuntimeError("Can't set bot_id in conjunction with (already set) "
1173+
"usernames.")
1174+
1175+
bot_id = self._parse_bot_id(bot_id)
1176+
1177+
self._bot_ids |= bot_id
1178+
1179+
def remove_usernames(self, username):
1180+
"""
1181+
Remove one or more users from allowed usernames.
1182+
Args:
1183+
username(:obj:`str` | List[:obj:`str`], optional): Which username(s) to disallow
1184+
through. Leading '@'s in usernames will be discarded.
1185+
"""
1186+
with self.__lock:
1187+
if self._bot_ids:
1188+
raise RuntimeError("Can't set username in conjunction with (already set) "
1189+
"bot_ids.")
1190+
1191+
username = self._parse_username(username)
1192+
self._usernames -= username
1193+
1194+
def remove_bot_ids(self, bot_id):
1195+
"""
1196+
Remove one or more users from allowed user ids.
1197+
Args:
1198+
bot_id(:obj:`int` | List[:obj:`int`], optional): Which bot ID(s) to disallow
1199+
through.
1200+
"""
1201+
with self.__lock:
1202+
if self._usernames:
1203+
raise RuntimeError("Can't set bot_id in conjunction with (already set) "
1204+
"usernames.")
1205+
bot_id = self._parse_bot_id(bot_id)
1206+
self._bot_ids -= bot_id
1207+
1208+
def filter(self, message):
1209+
"""""" # remove method from docs
1210+
if message.via_bot:
1211+
if self.bot_ids:
1212+
return message.via_bot.id in self.bot_ids
1213+
if self.usernames:
1214+
return (message.via_bot.username
1215+
and message.via_bot.username in self.usernames)
1216+
return self.allow_empty
1217+
return False
1218+
10561219
class chat(BaseFilter):
1057-
"""Filters messages to allow only those which are from specified chat ID.
1220+
"""Filters messages to allow only those which are from a specified chat ID or username.
10581221
10591222
Examples:
10601223
``MessageHandler(Filters.chat(-1234), callback_method)``

‎tests/test_filters.py

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
@pytest.fixture(scope='function')
2929
def update():
3030
return Update(0, Message(0, User(0, 'Testuser', False), datetime.datetime.utcnow(),
31-
Chat(0, 'private')))
31+
Chat(0, 'private'), via_bot=User(0, "Testbot", True)))
3232

3333

3434
@pytest.fixture(scope='function',
@@ -1093,3 +1093,128 @@ def filter(self, _):
10931093
update.message.text = 'test'
10941094
result = (Filters.command | DataFilter('blah'))(update)
10951095
assert result['test'] == ['blah']
1096+
1097+
def test_filters_via_bot_init(self):
1098+
with pytest.raises(RuntimeError, match='in conjunction with'):
1099+
Filters.via_bot(bot_id=1, username='bot')
1100+
1101+
def test_filters_via_bot_allow_empty(self, update):
1102+
assert not Filters.via_bot()(update)
1103+
assert Filters.via_bot(allow_empty=True)(update)
1104+
1105+
def test_filters_via_bot_id(self, update):
1106+
assert not Filters.via_bot(bot_id=1)(update)
1107+
update.message.via_bot.id = 1
1108+
assert Filters.via_bot(bot_id=1)(update)
1109+
update.message.via_bot.id = 2
1110+
assert Filters.via_bot(bot_id=[1, 2])(update)
1111+
assert not Filters.via_bot(bot_id=[3, 4])(update)
1112+
update.message.via_bot = None
1113+
assert not Filters.via_bot(bot_id=[3, 4])(update)
1114+
1115+
def test_filters_via_bot_username(self, update):
1116+
assert not Filters.via_bot(username='bot')(update)
1117+
assert not Filters.via_bot(username='Testbot')(update)
1118+
update.message.via_bot.username = 'bot@'
1119+
assert Filters.via_bot(username='@bot@')(update)
1120+
assert Filters.via_bot(username='bot@')(update)
1121+
assert Filters.via_bot(username=['bot1', 'bot@', 'bot2'])(update)
1122+
assert not Filters.via_bot(username=['@username', '@bot_2'])(update)
1123+
update.message.via_bot = None
1124+
assert not Filters.user(username=['@username', '@bot_2'])(update)
1125+
1126+
def test_filters_via_bot_change_id(self, update):
1127+
f = Filters.via_bot(bot_id=3)
1128+
update.message.via_bot.id = 3
1129+
assert f(update)
1130+
update.message.via_bot.id = 2
1131+
assert not f(update)
1132+
f.bot_ids = 2
1133+
assert f(update)
1134+
1135+
with pytest.raises(RuntimeError, match='username in conjunction'):
1136+
f.usernames = 'user'
1137+
1138+
def test_filters_via_bot_change_username(self, update):
1139+
f = Filters.via_bot(username='bot')
1140+
update.message.via_bot.username = 'bot'
1141+
assert f(update)
1142+
update.message.via_bot.username = 'Bot'
1143+
assert not f(update)
1144+
f.usernames = 'Bot'
1145+
assert f(update)
1146+
1147+
with pytest.raises(RuntimeError, match='bot_id in conjunction'):
1148+
f.bot_ids = 1
1149+
1150+
def test_filters_via_bot_add_user_by_name(self, update):
1151+
users = ['bot_a', 'bot_b', 'bot_c']
1152+
f = Filters.via_bot()
1153+
1154+
for user in users:
1155+
update.message.via_bot.username = user
1156+
assert not f(update)
1157+
1158+
f.add_usernames('bot_a')
1159+
f.add_usernames(['bot_b', 'bot_c'])
1160+
1161+
for user in users:
1162+
update.message.via_bot.username = user
1163+
assert f(update)
1164+
1165+
with pytest.raises(RuntimeError, match='bot_id in conjunction'):
1166+
f.add_bot_ids(1)
1167+
1168+
def test_filters_via_bot_add_user_by_id(self, update):
1169+
users = [1, 2, 3]
1170+
f = Filters.via_bot()
1171+
1172+
for user in users:
1173+
update.message.via_bot.id = user
1174+
assert not f(update)
1175+
1176+
f.add_bot_ids(1)
1177+
f.add_bot_ids([2, 3])
1178+
1179+
for user in users:
1180+
update.message.via_bot.username = user
1181+
assert f(update)
1182+
1183+
with pytest.raises(RuntimeError, match='username in conjunction'):
1184+
f.add_usernames('bot')
1185+
1186+
def test_filters_via_bot_remove_user_by_name(self, update):
1187+
users = ['bot_a', 'bot_b', 'bot_c']
1188+
f = Filters.via_bot(username=users)
1189+
1190+
with pytest.raises(RuntimeError, match='bot_id in conjunction'):
1191+
f.remove_bot_ids(1)
1192+
1193+
for user in users:
1194+
update.message.via_bot.username = user
1195+
assert f(update)
1196+
1197+
f.remove_usernames('bot_a')
1198+
f.remove_usernames(['bot_b', 'bot_c'])
1199+
1200+
for user in users:
1201+
update.message.via_bot.username = user
1202+
assert not f(update)
1203+
1204+
def test_filters_via_bot_remove_user_by_id(self, update):
1205+
users = [1, 2, 3]
1206+
f = Filters.via_bot(bot_id=users)
1207+
1208+
with pytest.raises(RuntimeError, match='username in conjunction'):
1209+
f.remove_usernames('bot')
1210+
1211+
for user in users:
1212+
update.message.via_bot.id = user
1213+
assert f(update)
1214+
1215+
f.remove_bot_ids(1)
1216+
f.remove_bot_ids([2, 3])
1217+
1218+
for user in users:
1219+
update.message.via_bot.username = user
1220+
assert not f(update)

0 commit comments

Comments
 (0)
0