8000 Add `filters.Mention` (#3941) · guillemap/python-telegram-bot@300ec92 · GitHub
[go: up one dir, main page]
Skip to content

Commit 300ec92

Browse files
Bibo-JoshiElmurodovJavohir
Bibo-Joshi
and
ElmurodovJavohir
authored
Add filters.Mention (python-telegram-bot#3941)
Co-authored-by: Javohir Elmurodov <elmurodovjavohir@gmail.com>
1 parent 075f517 commit 300ec92

File tree

2 files changed

+131
-1
lines changed

2 files changed

+131
-1
lines changed

‎telegram/ext/filters.py

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@
6666
"LOCATION",
6767
"Language",
6868
"MessageFilter",
69+
"Mention",
6970
"PASSPORT_DATA",
7071
"PHOTO",
7172
"POLL",
< 10000 /code>
@@ -91,14 +92,14 @@
9192
"VOICE",
9293
"ViaBot",
9394
)
94-
9595
import mimetypes
9696
import re
9797
from abc import ABC, abstractmethod
9898
from typing import (
9999
Collection,
100100
Dict,
101101
FrozenSet,
102+
Iterable,
102103
List,
103104
Match,
104105
NoReturn,
@@ -1521,6 +1522,73 @@ def filter(self, message: Message) -> bool:
15211522
"""Messages that contain :attr:`telegram.Message.location`."""
15221523

15231524

1525+
class Mention(MessageFilter):
1526+
"""Messages containing mentions of specified users or chats.
1527+
1528+
Examples:
1529+
.. code-block:: python
1530+
1531+
MessageHandler(filters.Mention("username"), callback)
1532+
MessageHandler(filters.Mention(["@username", 123456]), callback)
1533+
1534+
.. versionadded:: NEXT.VERSION
1535+
1536+
Args:
1537+
mentions (:obj:`int` | :obj:`str` | :class:`telegram.User` | Collection[:obj:`int` | \
1538+
:obj:`str` | :class:`telegram.User`]):
1539+
Specifies the users and chats to filter for. Messages that do not mention at least one
1540+
of the specified users or chats will not be handled. Leading ``'@'`` s in usernames
1541+
will be discarded.
1542+
"""
1543+
1544+
__slots__ = ("_mentions",)
1545+
1546+
def __init__(self, mentions: SCT[Union[int, str, TGUser]]):
1547+
super().__init__(name=f"filters.Mention({mentions})")
1548+
if isinstance(mentions, Iterable) and not isinstance(mentions, str):
1549+
self._mentions = {self._fix_mention_username(mention) for mention in mentions}
1550+
else:
1551+
self._mentions = {self._fix_mention_username(mentions)}
1552+
1553+
@staticmethod
1554+
def _fix_mention_username(mention: Union[int, str, TGUser]) -> Union[int, str, TGUser]:
1555+
if not isinstance(mention, str):
1556+
return mention
1557+
return mention.lstrip("@")
1558+
1559+
@classmethod
1560+
def _check_mention(cls, message: Message, mention: Union[int, str, TGUser]) -> bool:
1561+
if not message.entities:
1562+
return False
1563+
1564+
entity_texts = message.parse_entities(
1565+
types=[MessageEntity.MENTION, MessageEntity.TEXT_MENTION]
1566+
)
1567+
1568+
if isinstance(mention, TGUser):
1569+
return any(
1570+
mention.id == entity.user.id
1571+
or mention.username == entity.user.username
1572+
or mention.username == cls._fix_mention_username(entity_texts[entity])
1573+
for entity in message.entities
1574+
if entity.user
1575+
) or any(
1576+
mention.username == cls._fix_mention_username(entity_text)
1577+
for entity_text in entity_texts.values()
1578+
)
1579+
if isinstance(mention, int):
1580+
return bool(
1581+
any(mention == entity.user.id for entity in message.entities if entity.user)
1582+
)
1583+
return any(
1584+
mention == cls._fix_mention_username(entity_text)
1585+
for entity_text in entity_texts.values()
1586+
)
1587+
1588+
def filter(self, message: Message) -> bool:
1589+
return any(self._check_mention(message, mention) for mention in self._mentions)
1590+
1591+
15241592
class _PassportData(MessageFilter):
15251593
__slots__ = ()
15261594

‎tests/ext/test_filters.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2423,3 +2423,65 @@ def test_filters_attachment(self, update):
24232423
),
24242424
)
24252425
assert filters.ATTACHMENT.check_update(up)
2426+
2427+
def test_filters_mention_no_entities(self, update):
2428+
update.message.text = "test"
2429+
assert not filters.Mention("@test").check_update(update)
2430+
assert not filters.Mention(123456).check_update(update)
2431+
assert not filters.Mention("123456").check_update(update)
2432+
assert not filters.Mention(User(1, "first_name", False)).check_update(update)
2433+
assert not filters.Mention(
2434+
["@test", 123456, "123456", User(1, "first_name", False)]
2435+
).check_update(update)
2436+
2437+
def test_filters_mention_type_mention(self, update):
2438+
update.message.text = "@test1 @test2 user"
2439+
update.message.entities = [
2440+
MessageEntity(MessageEntity.MENTION, 0, 6),
2441+
MessageEntity(MessageEntity.MENTION, 7, 6),
2442+
]
2443+
2444+
user_no_username = User(123456, "first_name", False)
2445+
user_wrong_username = User(123456, "first_name", False, username="wrong")
2446+
user_1 = User(111, "first_name", False, username="test1")
2447+
user_2 = User(222, "first_name", False, username="test2")
2448+
2449+
for username in ("@test1", "@test2"):
2450+
assert filters.Mention(username).check_update(update)
2451+
assert filters.Mention({username}).check_update(update)
2452+
2453+
for user in (user_1, user_2):
2454+
assert filters.Mention(user).check_update(update)
2455+
assert filters.Mention({user}).check_update(update)
2456+
2457+
assert not filters.Mention(
2458+
["@test3", 123, user_no_username, user_wrong_username]
2459+
).check_update(update)
2460+
2461+
def test_filters_mention_type_text_mention(self, update):
2462+
user_1 = User(111, "first_name", False, username="test1")
2463+
user_2 = User(222, "first_name", False, username="test2")
2464+
user_no_username = User(123456, "first_name", False)
2465+
user_wrong_username = User(123456, "first_name", False, username="wrong")
2466+
2467+
update.message.text = "test1 test2 user"
2468+
update.message.entities = [
2469+
MessageEntity(MessageEntity.TEXT_MENTION, 0, 5, user=user_1),
2470+
MessageEntity(MessageEntity.TEXT_MENTION, 6, 5, user=user_2),
2471+
]
2472+
2473+
for username in ("@test1", "@test2"):
2474+
assert filters.Mention(username).check_update(update)
2475+
assert filters.Mention({username}).check_update(update)
2476+
2477+
for user in (user_1, user_2):
2478+
assert filters.Mention(user).check_update(update)
2479+
assert filters.Mention({user}).check_update(update)
2480+
2481+
for user_id in (111, 222):
2482+
assert filters.Mention(user_id).check_update(update)
2483+
assert filters.Mention({user_id}).check_update(update)
2484+
2485+
assert not filters.Mention(
2486+
["@test3", 123, user_no_username, user_wrong_username]
2487+
).check_update(update)

0 commit comments

Comments
 (0)
0