Add new Restriction object and make User and Chat objects use it

delivrance · delivrance · commit 5ce62bd79c35 · 2019-09-07T13:28:05.000+02:00
diff --git a/compiler/docs/compiler.py b/compiler/docs/compiler.py
@@ -309,6 +309,7 @@ def get_title_list(s: str) -> list:
             ChatMember
             ChatPermissions
             Dialog
+            Restriction
         """,
         messages_media="""
         Messages & Media
diff --git a/pyrogram/client/types/user_and_chats/chat.py b/pyrogram/client/types/user_and_chats/chat.py
@@ -16,12 +16,13 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import Union
+from typing import Union, List
 
 import pyrogram
 from pyrogram.api import types
 from .chat_permissions import ChatPermissions
 from .chat_photo import ChatPhoto
+from .restriction import Restriction
 from ..object import Object
 from ...ext import utils
 
@@ -87,8 +88,8 @@ class Chat(Object):
         members_count (``int``, *optional*):
             Chat members count, for groups, supergroups and channels only.
 
-        restriction_reason (``str``, *optional*):
-            The reason why this chat might be unavailable to some users.
+        restrictions (List of :obj:`Restriction`, *optional*):
+            The list of reasons why this chat might be unavailable to some users.
             This field is available only in case *is_restricted* is True.
 
         permissions (:obj:`ChatPermissions` *optional*):
@@ -120,7 +121,7 @@ def __init__(
         sticker_set_name: str = None,
         can_set_sticker_set: bool = None,
         members_count: int = None,
-        restriction_reason: str = None,
+        restrictions: List[Restriction] = None,
         permissions: "pyrogram.ChatPermissions" = None,
         distance: int = None
     ):
@@ -143,7 +144,7 @@ def __init__(
         self.sticker_set_name = sticker_set_name
         self.can_set_sticker_set = can_set_sticker_set
         self.members_count = members_count
-        self.restriction_reason = restriction_reason
+        self.restrictions = restrictions
         self.permissions = permissions
         self.distance = distance
 
@@ -162,7 +163,7 @@ def _parse_user_chat(client, user: types.User) -> "Chat":
             first_name=user.first_name,
             last_name=user.last_name,
             photo=ChatPhoto._parse(client, user.photo, peer_id),
-            restriction_reason=user.restriction_reason,
+            restrictions=pyrogram.List([Restriction._parse(r) for r in user.restriction_reason]) or None,
             client=client
         )
 
@@ -183,6 +184,7 @@ def _parse_chat_chat(client, chat: types.Chat) -> "Chat":
     @staticmethod
     def _parse_channel_chat(client, channel: types.Channel) -> "Chat":
         peer_id = utils.get_channel_id(channel.id)
+        restriction_reason = getattr(channel, "restriction_reason", [])
 
         return Chat(
             id=peer_id,
@@ -193,7 +195,7 @@ def _parse_channel_chat(client, channel: types.Channel) -> "Chat":
             title=channel.title,
             username=getattr(channel, "username", None),
             photo=ChatPhoto._parse(client, getattr(channel, "photo", None), peer_id),
-            restriction_reason=getattr(channel, "restriction_reason", None),
+            restrictions=pyrogram.List([Restriction._parse(r) for r in restriction_reason]) or None,
             permissions=ChatPermissions._parse(getattr(channel, "default_banned_rights", None)),
             members_count=getattr(channel, "participants_count", None),
             client=client
diff --git a/pyrogram/client/types/user_and_chats/restriction.py b/pyrogram/client/types/user_and_chats/restriction.py
@@ -0,0 +1,50 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram 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 General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+from pyrogram.api import types
+from ..object import Object
+
+
+class Restriction(Object):
+    """A restriction applied to bots or chats.
+
+    Parameters:
+        platform (``str``):
+            The platform the restriction is applied to, e.g. "ios", "android"
+
+        reason (``str``):
+            The restriction reason, e.g. "porn", "copyright".
+
+        text (``str``):
+            The restriction text.
+    """
+
+    def __init__(self, *, platform: str, reason: str, text: str):
+        super().__init__(None)
+
+        self.platform = platform
+        self.reason = reason
+        self.text = text
+
+    @staticmethod
+    def _parse(restriction: types.RestrictionReason) -> "Restriction":
+        return Restriction(
+            platform=restriction.platform,
+            reason=restriction.reason,
+            text=restriction.text
+        )
diff --git a/pyrogram/client/types/user_and_chats/user.py b/pyrogram/client/types/user_and_chats/user.py
@@ -17,10 +17,12 @@
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import html
+from typing import List
 
 import pyrogram
 from pyrogram.api import types
 from .chat_photo import ChatPhoto
+from .restriction import Restriction
 from ..object import Object
 from ..update import Update
 
@@ -101,8 +103,8 @@ class User(Object, Update):
         photo (:obj:`ChatPhoto <pyrogram.ChatPhoto>`, *optional*):
             User's or bot's current profile photo. Suitable for downloads only.
 
-        restriction_reason (``str``, *optional*):
-            The reason why this bot might be unavailable to some users.
+        restrictions (List of :obj:`Restriction`, *optional*):
+            The list of reasons why this bot might be unavailable to some users.
             This field is available only in case *is_restricted* is True.
     """
 
@@ -130,7 +132,7 @@ def __init__(
         dc_id: int = None,
         phone_number: str = None,
         photo: ChatPhoto = None,
-        restriction_reason: str = None
+        restrictions: List[Restriction] = None
     ):
         super().__init__(client)
 
@@ -154,7 +156,7 @@ def __init__(
         self.dc_id = dc_id
         self.phone_number = phone_number
         self.photo = photo
-        self.restriction_reason = restriction_reason
+        self.restrictions = restrictions
 
     def __format__(self, format_spec):
         if format_spec == "mention":
@@ -186,7 +188,7 @@ def _parse(client, user: types.User) -> "User" or None:
             dc_id=getattr(user.photo, "dc_id", None),
             phone_number=user.phone,
             photo=ChatPhoto._parse(client, user.photo, user.id),
-            restriction_reason=user.restriction_reason,
+            restrictions=pyrogram.List([Restriction._parse(r) for r in user.restriction_reason]) or None,
             client=client
         )
 
