From 823d030c2cd812dcec770afaef2eb44b1cc5b21b Mon Sep 17 00:00:00 2001 From: Bibo-Joshi <22366557+Bibo-Joshi@users.noreply.github.com> Date: Sat, 30 Apr 2022 22:18:11 +0200 Subject: [PATCH] Use enums for dynamic types & rename two attributes in ChatMember (#2817) --- examples/chatmemberbot.py | 4 +-- telegram/_chat.py | 3 ++- telegram/_chatmember.py | 34 +++++++++++++------------ telegram/_messageentity.py | 3 ++- telegram/_poll.py | 3 ++- telegram/_utils/enum.py | 44 ++++++++++++++++++++++----------- telegram/constants.py | 8 +++--- tests/test_chat.py | 8 +++++- tests/test_chatmemberhandler.py | 4 +-- tests/test_chatmemberupdated.py | 8 +++--- tests/test_messageentity.py | 7 ++++++ tests/test_poll.py | 25 +++++++++++++++++++ 12 files changed, 105 insertions(+), 46 deletions(-) diff --git a/examples/chatmemberbot.py b/examples/chatmemberbot.py index 30f7138bc..0c8c78ef2 100644 --- a/examples/chatmemberbot.py +++ b/examples/chatmemberbot.py @@ -50,7 +50,7 @@ def extract_status_change( old_status in [ ChatMember.MEMBER, - ChatMember.CREATOR, + ChatMember.OWNER, ChatMember.ADMINISTRATOR, ] or (old_status == ChatMember.RESTRICTED and old_is_member is True) @@ -59,7 +59,7 @@ def extract_status_change( new_status in [ ChatMember.MEMBER, - ChatMember.CREATOR, + ChatMember.OWNER, ChatMember.ADMINISTRATOR, ] or (new_status == ChatMember.RESTRICTED and new_is_member is True) diff --git a/telegram/_chat.py b/telegram/_chat.py index ea9469721..64a563bd4 100644 --- a/telegram/_chat.py +++ b/telegram/_chat.py @@ -22,6 +22,7 @@ from datetime import datetime from typing import TYPE_CHECKING, List, Optional, ClassVar, Union, Tuple, Any from telegram import ChatPhoto, TelegramObject, constants +from telegram._utils import enum from telegram._utils.types import JSONDict, FileInput, ODVInput, DVInput, ReplyMarkup from telegram._utils.defaultvalue import DEFAULT_NONE @@ -228,7 +229,7 @@ class Chat(TelegramObject): ): # Required self.id = int(id) # pylint: disable=invalid-name - self.type = type + self.type = enum.get_member(constants.ChatType, type, type) # Optionals self.title = title self.username = username diff --git a/telegram/_chatmember.py b/telegram/_chatmember.py index a97a124d5..5e21f46c6 100644 --- a/telegram/_chatmember.py +++ b/telegram/_chatmember.py @@ -43,16 +43,18 @@ class ChatMember(TelegramObject): considered equal, if their :attr:`user` and :attr:`status` are equal. .. versionchanged:: 14.0 - As of Bot API 5.3, :class:`ChatMember` is nothing but the base class for the subclasses - listed above and is no longer returned directly by :meth:`~telegram.Bot.get_chat`. - Therefore, most of the arguments and attributes were removed and you should no longer - use :class:`ChatMember` directly. + * As of Bot API 5.3, :class:`ChatMember` is nothing but the base class for the subclasses + listed above and is no longer returned directly by :meth:`~telegram.Bot.get_chat`. + Therefore, most of the arguments and attributes were removed and you should no longer + use :class:`ChatMember` directly. + * The constant ``ChatMember.CREATOR`` was replaced by :attr:`~telegram.ChatMember.OWNER` + * The constant ``ChatMember.KICKED`` was replaced by :attr:`~telegram.ChatMember.BANNED` Args: user (:class:`telegram.User`): Information about the user. status (:obj:`str`): The member's status in the chat. Can be - :attr:`~telegram.ChatMember.ADMINISTRATOR`, :attr:`~telegram.ChatMember.CREATOR`, - :attr:`~telegram.ChatMember.KICKED`, :attr:`~telegram.ChatMember.LEFT`, + :attr:`~telegram.ChatMember.ADMINISTRATOR`, :attr:`~telegram.ChatMember.OWNER`, + :attr:`~telegram.ChatMember.BANNED`, :attr:`~telegram.ChatMember.LEFT`, :attr:`~telegram.ChatMember.MEMBER` or :attr:`~telegram.ChatMember.RESTRICTED`. Attributes: @@ -65,10 +67,10 @@ class ChatMember(TelegramObject): ADMINISTRATOR: ClassVar[str] = constants.ChatMemberStatus.ADMINISTRATOR """:const:`telegram.constants.ChatMemberStatus.ADMINISTRATOR`""" - CREATOR: ClassVar[str] = constants.ChatMemberStatus.CREATOR - """:const:`telegram.constants.ChatMemberStatus.CREATOR`""" - KICKED: ClassVar[str] = constants.ChatMemberStatus.KICKED - """:const:`telegram.constants.ChatMemberStatus.KICKED`""" + OWNER: ClassVar[str] = constants.ChatMemberStatus.OWNER + """:const:`telegram.constants.ChatMemberStatus.OWNER`""" + BANNED: ClassVar[str] = constants.ChatMemberStatus.BANNED + """:const:`telegram.constants.ChatMemberStatus.BANNED`""" LEFT: ClassVar[str] = constants.ChatMemberStatus.LEFT """:const:`telegram.constants.ChatMemberStatus.LEFT`""" MEMBER: ClassVar[str] = constants.ChatMemberStatus.MEMBER @@ -95,12 +97,12 @@ class ChatMember(TelegramObject): data['until_date'] = from_timestamp(data.get('until_date', None)) _class_mapping: Dict[str, Type['ChatMember']] = { - cls.CREATOR: ChatMemberOwner, + cls.OWNER: ChatMemberOwner, cls.ADMINISTRATOR: ChatMemberAdministrator, cls.MEMBER: ChatMemberMember, cls.RESTRICTED: ChatMemberRestricted, cls.LEFT: ChatMemberLeft, - cls.KICKED: ChatMemberBanned, + cls.BANNED: ChatMemberBanned, } if cls is ChatMember: @@ -132,7 +134,7 @@ class ChatMemberOwner(ChatMember): Attributes: status (:obj:`str`): The member's status in the chat, - always :tg-const:`telegram.ChatMember.CREATOR`. + always :tg-const:`telegram.ChatMember.OWNER`. user (:class:`telegram.User`): Information about the user. is_anonymous (:obj:`bool`): :obj:`True`, if the user's presence in the chat is hidden. @@ -149,7 +151,7 @@ class ChatMemberOwner(ChatMember): custom_title: str = None, **_kwargs: object, ): - super().__init__(status=ChatMember.CREATOR, user=user) + super().__init__(status=ChatMember.OWNER, user=user) self.is_anonymous = is_anonymous self.custom_title = custom_title @@ -436,7 +438,7 @@ class ChatMemberBanned(ChatMember): Attributes: status (:obj:`str`): The member's status in the chat, - always :tg-const:`telegram.ChatMember.KICKED`. + always :tg-const:`telegram.ChatMember.BANNED`. user (:class:`telegram.User`): Information about the user. until_date (:class:`datetime.datetime`): Date when restrictions will be lifted for this user. @@ -446,5 +448,5 @@ class ChatMemberBanned(ChatMember): __slots__ = ('until_date',) def __init__(self, user: User, until_date: datetime.datetime, **_kwargs: object): - super().__init__(status=ChatMember.KICKED, user=user) + super().__init__(status=ChatMember.BANNED, user=user) self.until_date = until_date diff --git a/telegram/_messageentity.py b/telegram/_messageentity.py index 70437077f..c316a0f91 100644 --- a/telegram/_messageentity.py +++ b/telegram/_messageentity.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, ClassVar from telegram import TelegramObject, User, constants +from telegram._utils import enum from telegram._utils.types import JSONDict if TYPE_CHECKING: @@ -74,7 +75,7 @@ class MessageEntity(TelegramObject): **_kwargs: Any, ): # Required - self.type = type + self.type = enum.get_member(constants.MessageEntityType, type, type) self.offset = offset self.length = length # Optionals diff --git a/telegram/_poll.py b/telegram/_poll.py index da9e362ea..3c80ef9fb 100644 --- a/telegram/_poll.py +++ b/telegram/_poll.py @@ -24,6 +24,7 @@ import sys from typing import TYPE_CHECKING, Any, Dict, List, Optional, ClassVar from telegram import MessageEntity, TelegramObject, User, constants +from telegram._utils import enum from telegram._utils.datetime import from_timestamp, to_timestamp from telegram._utils.types import JSONDict @@ -189,7 +190,7 @@ class Poll(TelegramObject): self.total_voter_count = total_voter_count self.is_closed = is_closed self.is_anonymous = is_anonymous - self.type = type + self.type = enum.get_member(constants.PollType, type, type) self.allows_multiple_answers = allows_multiple_answers self.correct_option_id = correct_option_id self.explanation = explanation diff --git a/telegram/_utils/enum.py b/telegram/_utils/enum.py index 63ff62645..6a969a4dc 100644 --- a/telegram/_utils/enum.py +++ b/telegram/_utils/enum.py @@ -1,21 +1,22 @@ +#!/usr/bin/env python # -# A library that provides a Python interface to the Telegram Bot API -# Copyright (C) 2015-2022 -# Leandro Toledo de Souza +# A library that provides a Python interface to the Telegram Bot API +# Copyright (C) 2015-2022 +# Leandro Toledo de Souza # -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. # -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Lesser Public License for more details. +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser Public License for more details. # -# You should have received a copy of the GNU Lesser Public License -# along with this program. If not, see [http://www.gnu.org/licenses/]. -"""This module contains a helper class for Enums that should be subclasses of `str`. +# You should have received a copy of the GNU Lesser Public License +# along with this program. If not, see [http://www.gnu.org/licenses/]. +"""This module contains helper functions related to enums. Warning: Contents of this module are intended to be used internally by the library and *not* by the @@ -23,6 +24,21 @@ Warning: the changelog. """ from enum import Enum +from typing import TypeVar, Union, Type + +_A = TypeVar('_A') +_B = TypeVar('_B') +_Enum = TypeVar('_Enum', bound=Enum) + + +def get_member(enum: Type[_Enum], value: _A, default: _B) -> Union[_Enum, _A, _B]: + """Tries to call ``enum(value)`` to convert the value into an enumeration member. + If that fails, the ``default`` is returned. + """ + try: + return enum(value) + except ValueError: + return default class StringEnum(str, Enum): diff --git a/telegram/constants.py b/telegram/constants.py index 36d71889f..09ffa77f0 100644 --- a/telegram/constants.py +++ b/telegram/constants.py @@ -211,10 +211,10 @@ class ChatMemberStatus(StringEnum): ADMINISTRATOR = 'administrator' """:obj:`str`: A :class:`telegram.ChatMember` who is administrator of the chat.""" - CREATOR = 'creator' - """:obj:`str`: A :class:`telegram.ChatMember` who is the creator of the chat.""" - KICKED = 'kicked' - """:obj:`str`: A :class:`telegram.ChatMember` who was kicked from the chat.""" + OWNER = 'creator' + """:obj:`str`: A :class:`telegram.ChatMember` who is the owner of the chat.""" + BANNED = 'kicked' + """:obj:`str`: A :class:`telegram.ChatMember` who was banned in the chat.""" LEFT = 'left' """:obj:`str`: A :class:`telegram.ChatMember` who has left the chat.""" MEMBER = 'member' diff --git a/tests/test_chat.py b/tests/test_chat.py index eaa5cf4de..f2307dacf 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -20,7 +20,7 @@ import pytest from telegram import Chat, ChatPermissions, ChatLocation, Location, Bot, User -from telegram.constants import ChatAction +from telegram.constants import ChatAction, ChatType from tests.conftest import check_shortcut_signature, check_shortcut_call, check_defaults_handling @@ -122,6 +122,12 @@ class TestChat: assert chat_dict['linked_chat_id'] == chat.linked_chat_id assert chat_dict['location'] == chat.location.to_dict() + def test_enum_init(self): + chat = Chat(id=1, type='foo') + assert chat.type == 'foo' + chat = Chat(id=1, type='private') + assert chat.type is ChatType.PRIVATE + def test_link(self, chat): assert chat.link == f'https://t.me/{chat.username}' chat.username = None diff --git a/tests/test_chatmemberhandler.py b/tests/test_chatmemberhandler.py index 8db1733c7..dab1b0f88 100644 --- a/tests/test_chatmemberhandler.py +++ b/tests/test_chatmemberhandler.py @@ -75,8 +75,8 @@ def chat_member_updated(): Chat(1, 'chat'), User(1, '', False), from_timestamp(int(time.time())), - ChatMember(User(1, '', False), ChatMember.CREATOR), - ChatMember(User(1, '', False), ChatMember.CREATOR), + ChatMember(User(1, '', False), ChatMember.OWNER), + ChatMember(User(1, '', False), ChatMember.OWNER), ) diff --git a/tests/test_chatmemberupdated.py b/tests/test_chatmemberupdated.py index 964805366..dd0c212df 100644 --- a/tests/test_chatmemberupdated.py +++ b/tests/test_chatmemberupdated.py @@ -182,7 +182,7 @@ class TestChatMemberUpdated: Chat(1, 'chat'), User(1, '', False), time, - ChatMember(User(1, '', False), ChatMember.CREATOR), + ChatMember(User(1, '', False), ChatMember.OWNER), new_chat_member, ) # wrong new_chat_member @@ -191,10 +191,10 @@ class TestChatMemberUpdated: User(1, '', False), time, old_chat_member, - ChatMember(User(1, '', False), ChatMember.CREATOR), + ChatMember(User(1, '', False), ChatMember.OWNER), ) # wrong type - g = ChatMember(User(1, '', False), ChatMember.CREATOR) + g = ChatMember(User(1, '', False), ChatMember.OWNER) assert a == b assert hash(a) == hash(b) @@ -255,5 +255,5 @@ class TestChatMemberUpdated: diff = chat_member_updated.difference() assert diff.pop('is_anonymous') == (False, None) assert diff.pop('until_date') == (None, datetime.datetime(2021, 1, 1)) - assert diff.pop('status') == (ChatMember.CREATOR, ChatMember.KICKED) + assert diff.pop('status') == (ChatMember.OWNER, ChatMember.BANNED) assert diff == {} diff --git a/tests/test_messageentity.py b/tests/test_messageentity.py index ef3add809..e142c60af 100644 --- a/tests/test_messageentity.py +++ b/tests/test_messageentity.py @@ -19,6 +19,7 @@ import pytest from telegram import MessageEntity, User +from telegram.constants import MessageEntityType @pytest.fixture(scope="class", params=MessageEntity.ALL_TYPES) @@ -70,6 +71,12 @@ class TestMessageEntity: if message_entity.language: assert entity_dict['language'] == message_entity.language + def test_enum_init(self): + entity = MessageEntity(type='foo', offset=0, length=1) + assert entity.type == 'foo' + entity = MessageEntity(type='url', offset=0, length=1) + assert entity.type is MessageEntityType.URL + def test_equality(self): a = MessageEntity(MessageEntity.BOLD, 2, 3) b = MessageEntity(MessageEntity.BOLD, 2, 3) diff --git a/tests/test_poll.py b/tests/test_poll.py index 2f8b04282..4781a6d1f 100644 --- a/tests/test_poll.py +++ b/tests/test_poll.py @@ -22,6 +22,7 @@ from datetime import datetime from telegram import Poll, PollOption, PollAnswer, User, MessageEntity from telegram._utils.datetime import to_timestamp +from telegram.constants import PollType @pytest.fixture(scope="class") @@ -211,6 +212,30 @@ class TestPoll: assert poll_dict['open_period'] == poll.open_period assert poll_dict['close_date'] == to_timestamp(poll.close_date) + def test_enum_init(self): + poll = Poll( + type='foo', + id='id', + question='question', + options=[], + total_voter_count=0, + is_closed=False, + is_anonymous=False, + allows_multiple_answers=False, + ) + assert poll.type == 'foo' + poll = Poll( + type=PollType.QUIZ, + id='id', + question='question', + options=[], + total_voter_count=0, + is_closed=False, + is_anonymous=False, + allows_multiple_answers=False, + ) + assert poll.type is PollType.QUIZ + def test_parse_entity(self, poll): entity = MessageEntity(type=MessageEntity.URL, offset=13, length=17) poll.explanation_entities = [entity]