8000 Fixes for `Roles` (#90) · python-telegram-bot/ptbcontrib@222ae22 · GitHub
[go: up one dir, main page]
Skip to content

Commit 222ae22

Browse files
Bibo-JoshiBibo-Joshi
authored
Fixes for Roles (#90)
1 parent 29189d8 commit 222ae22

File tree

4 files changed

+72
-23
lines changed

4 files changed

+72
-23
lines changed

‎ptbcontrib/roles/README.md

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ Provides classes and methods for granular, hierarchical user access management.
44

55
```python
66
from telegram import Update
7-
from telegram.ext import TypeHandler, ApplicationBuilder, MessageHandler, SomeHandler
7+
from telegram.ext import TypeHandler, ApplicationBuilder, MessageHandler, SomeHandler, filters
88
from ptbcontrib.roles import setup_roles, RolesHandler, Role
99

1010

@@ -23,7 +23,8 @@ async def post_init(application):
2323
# and before the polling starts.
2424
# This ensures that the roles are initialized *after* the
2525
# persistence has been loaded, if persistence is used.
26-
# See also the wiki page at https://github.com/python-telegram-bot/python-telegram-bot/wiki/Making-your-bot-persistent
26+
# See also the wiki page at
27+
# https://github.com/python-telegram-bot/python-telegram-bot/wiki/Making-your-bot-persistent
2728
roles = setup_roles(application)
2829

2930
if 'my_role_1' not in roles:
@@ -33,15 +34,17 @@ async def post_init(application):
3334
if 'my_role_2' not in roles:
3435
roles.add_role(name='my_role_2')
3536

36-
roles.add_admin('authors_user_id')
37+
roles.add_admin(123)
3738
my_role_2 = roles['my_role_2']
3839
my_role_1.add_child_role(my_role_2)
3940

4041
# Anyone can add themself to my_role_2
41-
application.add_handler(TypeHandler(Update, add_to_my_role_2))
42+
application.add_handler(RolesHandler(TypeHandler(Update, add_to_my_role_2), roles=None))
4243

4344
# Only the admin can add users to my_role_1
44-
application.add_handler(RolesHandler(MessageHandler(Filters.text, add_to_my_role_1), roles=roles.admins))
45+
application.add_handler(
46+
RolesHandler(MessageHandler(filters.TEXT, add_to_my_role_1), roles=roles.admins)
47+
)
4548

4649
# This will be accessible by my_role_2, my_role_1 and the admin
4750
application.add_handler(RolesHandler(SomeHandler(...), roles=my_role_2))
@@ -51,20 +54,13 @@ async def post_init(application):
5154

5255
# This will be accessible by anyone except my_role_1 and my_role_2
5356
application.add_handler(RolesHandler(SomeHandler(...), roles=~my_role_1))
54-
55-
# This will be accessible only by the admins of the group the update was sent in
56-
application.add_handler(RolesHandler(SomeHandler(...), roles=roles.chat_admins))
57-
58-
# This will be accessible only by the creator of the group the update was sent in
59-
application.add_handler(RolesHandler(SomeHandler(...), roles=roles.chat_creator))
6057

6158
# You can compare the roles regarding hierarchy:
62-
roles.ADMINS >= roles['my_role_1'] # True
63-
roles.ADMINS >= roles['my_role_2'] # True
59+
roles['my_role_1'] >= roles['my_role_2'] # True
6460
roles['my_role_1'] < roles['my_role_2'] # False
65-
roles.ADMINS >= Role(...) # False, since neither of those is a parent of the other
61+
roles.admins >= Role(...) # False, since neither of those is a parent of the other
6662

67-
application = ApplicationBuilder.token('TOKEN').post_init(post_init).build()
63+
application = ApplicationBuilder().token('TOKEN').post_init(post_init).build()
6864
```
6965

7066
Please see the docstrings for more details.

‎ptbcontrib/roles/roles.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
_REPLACED_LOCK: str = "ptbcontrib_roles_replaced_lock"
2828

2929

30+
# We only inherit from UpdateFilter to get the nice syntax of the bitwise operators.
31+
# If not for that, we could do without inheritance.
3032
class Role(UpdateFilter):
3133
"""This class represents a security level used by :class:`telegram.ext.Roles`. Roles have a
3234
hierarchy, i.e. a role can do everthing, its child roles can do. To compare two roles you may
@@ -133,6 +135,13 @@ def __init__(
133135
self._admin_event.wait()
134136
self._admin.add_child_role(self)
135137

138+
def check_update(self, update: Update) -> bool:
139+
"""Check if the update is allowed by this role. This is just an alias for
140+
:meth:`filter`.
141+
"""
142+
# Override UpdateFilter.check_update to handle also updates that are not messages
143+
return self.filter(update)
144+
136145
@staticmethod
137146
def _parse_chat_id(chat_id: Union[int, List[int], Tuple[int, ...], None]) -> Set[int]:
138147
if chat_id is None:
@@ -421,8 +430,7 @@ def __init__(self, bot: Bot) -> None:
421430

422431
def set_bot(self, bot: Bot) -> None:
423432
"""If for some reason you can't pass the bot on initialization, you can set it with this
424-
method. Make sure to set the bot before the first call of :attr:`chat_admins` or
425-
:attr:`chat_creator`.
433+
method.
426434
427435
Args:
428436
bot (:class:`telegram.Bot`): The bot to set.

‎ptbcontrib/roles/roleshandler.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -103,20 +103,25 @@ class RolesHandler(BaseHandler[Update, _CCT]):
103103
104104
Args:
105105
handler (:class:`telegram.ext.BaseHandler`): The handler to wrap.
106-
roles (:class:`ptbcontrib.roles.Roles`): The roles to apply to the handler. Can be combined
107-
with bitwise operations.
106+
roles (:class:`ptbcontrib.roles.Roles`| :obj:`None`): The roles to apply to the handler.
107+
Can be combined with bitwise operations. :obj:`None` is accepted as input, in which
108+
case no roles restrictions will be applied. This can be useful if you want to have
109+
``context.roles`` available without restricting access.
110+
108111
"""
109112

110-
def __init__(self, handler: BaseHandler[Update, _CCT], roles: Role) -> None:
113+
def __init__(
114+
self, handler: BaseHandler[Update, _CCT], roles: Union[Role, InvertedRole, None]
115+
) -> None:
111116
self.handler = handler
112-
self.roles: Union[Role, InvertedRole] = roles
117+
self.roles: Union[Role, InvertedRole, None] = roles
113118
super().__init__(self.handler.callback)
114119

115120
def check_update(self, update: object) -> Any:
116121
"""Checks if the update should be handled."""
117122
if not isinstance( F438 update, Update):
118123
return False
119-
if self.roles.check_update(update):
124+
if self.roles is None or self.roles.check_update(update):
120125
return self.handler.check_update(update)
121126
return False
122127

‎tests/test_roles.py

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
from typing import Optional
2525

2626
import pytest
27-
from telegram import Chat, Message, Update, User
27+
from telegram import CallbackQuery, Chat, Message, Update, User
2828
from telegram.ext import CallbackContext, MessageHandler, PicklePersistence, filters
2929

3030
from ptbcontrib.roles import BOT_DATA_KEY, Role, Roles, RolesBotData, RolesHandler, setup_roles
@@ -293,6 +293,20 @@ def test_always_allow_admin(self, update, role):
293293
finally:
294294
role._admin.kick_member(0)
295295

296+
def test_non_message_update(self, update, role):
297+
update.message = None
298+
assert not role.check_update(update)
299+
300+
update.callback_query = CallbackQuery(
301+
id="id",
302+
from_user=User(id=0, is_bot=False, first_name="first_name"),
303+
chat_instance="chat_instance",
304+
)
305+
assert not role.check_update(update)
306+
307+
role.add_member(0)
308+
assert role.check_update(update)
309+
296310
def test_pickle(self, role, parent_role):
297311
role.add_member([0, 1, 3])
298312
parent_role.add_member([4, 5, 6])
@@ -526,6 +540,32 @@ def callback(_, context: CallbackContext):
526540
await app.process_update(update)
527541
assert self.test_flag
528542

543+
@pytest.mark.parametrize("roles_bot_data", [True, False])
544+
async def test_callback_and_context_no_roles(self, app, update, roles_bot_data):
545+
if roles_bot_data:
546+
app.bot_data = RolesData()
547+
self.roles = setup_roles(app)
548+
self.roles.admins.add_member(42)
549+
self.roles.add_role(name="role", chat_ids=[1])
550+
self.test_flag = False
551+
552+
def callback(_, context: CallbackContext):
553+
self.test_flag = context.roles is self.roles
554+
555+
handler = MessageHandler(filters.ALL, callback=callback)
556+
roles_handler = RolesHandler(handler, roles=None)
557+
558+
assert roles_handler.check_update(update)
559+
update.message.from_user.id = 1
560+
assert roles_handler.check_update(update)
561+
update.message.from_user.id = 42
562+
assert roles_handler.check_update(update)
563+
564+
app.add_handler(roles_handler)
565+
async with app:
566+
await app.process_update(update)
567+
assert self.test_flag
568+
529569
@pytest.mark.parametrize("roles_bot_data", [True, False])
530570
async def test_handler_error_message(self, app, update, roles_bot_data):
531571
if roles_bot_data:

0 commit comments

Comments
 (0)
0