add support for User.language_code (python-telegram-bot#624)

jh0ker · web-flow · commit 1c4c228cf147 · 2017-05-21T14:00:53.000+02:00
* add support for User.language_code

* Add language filter

Useful is you wanna do something like restrict your shop to a single or a few locales or something like that.
diff --git a/telegram/ext/filters.py b/telegram/ext/filters.py
@@ -19,6 +19,11 @@
 """ This module contains the Filters for use with the MessageHandler class """
 from telegram import Chat
 
+try:
+    str_type = base_string
+except NameError:
+    str_type = str
+
 
 class BaseFilter(object):
     """Base class for all Message Filters
@@ -263,3 +268,24 @@ def filter(self, message):
             return message.chat.type in [Chat.GROUP, Chat.SUPERGROUP]
 
     group = _Group()
+
+    class language(BaseFilter):
+        """
+        Filters messages to only allow those which are from users with a certain language code.
+        Note that according to telegrams documentation, every single user does not have the
+        language_code attribute.
+
+        Args:
+            lang (str|list): Which language code(s) to allow through. This will be matched using
+                .startswith meaning that 'en' will match both 'en_US' and 'en_GB'
+        """
+
+        def __init__(self, lang):
+            if isinstance(lang, str_type):
+                self.lang = [lang]
+            else:
+                self.lang = lang
+
+        def filter(self, message):
+            return message.from_user.language_code and any(
+                [message.from_user.language_code.startswith(x) for x in self.lang])
diff --git a/telegram/user.py b/telegram/user.py
@@ -26,21 +26,23 @@ class User(TelegramObject):
     """This object represents a Telegram User.
 
     Attributes:
-        id (int):
-        first_name (str):
-        last_name (str):
-        username (str):
-        type (str):
+        id (int): Unique identifier for this user or bot
+        first_name (str): User's or bot's first name
+        last_name (str): User's or bot's last name
+        username (str): User's or bot's username
+        language_code (str): IETF language tag of the user's language
+        type (str): Deprecated
 
     Args:
-        id (int):
-        first_name (str):
+        id (int): Unique identifier for this user or bot
+        first_name (str): User's or bot's first name
         **kwargs: Arbitrary keyword arguments.
 
     Keyword Args:
-        type (Optional[str]):
-        last_name (Optional[str]):
-        username (Optional[str]):
+        type (Optional[str]): Deprecated
+        last_name (Optional[str]): User's or bot's last name
+        username (Optional[str]): User's or bot's username
+        language_code (Optional[str]): IETF language tag of the user's language
         bot (Optional[Bot]): The Bot to use for instance methods
     """
 
@@ -50,6 +52,7 @@ def __init__(self,
                  type=None,
                  last_name=None,
                  username=None,
+                 language_code=None,
                  bot=None,
                  **kwargs):
         # Required
@@ -59,6 +62,7 @@ def __init__(self,
         self.type = type
         self.last_name = last_name
         self.username = username
+        self.language_code = language_code
 
         self.bot = bot
 
diff --git a/tests/test_filters.py b/tests/test_filters.py
@@ -274,6 +274,32 @@ class _CustomFilter(BaseFilter):
         with self.assertRaises(NotImplementedError):
             (custom & Filters.text)(self.message)
 
+    def test_language_filter_single(self):
+        self.message.from_user.language_code = 'en_US'
+        self.assertTrue((Filters.language('en_US'))(self.message))
+        self.assertTrue((Filters.language('en'))(self.message))
+        self.assertFalse((Filters.language('en_GB'))(self.message))
+        self.assertFalse((Filters.language('da'))(self.message))
+        self.message.from_user.language_code = 'en_GB'
+        self.assertFalse((Filters.language('en_US'))(self.message))
+        self.assertTrue((Filters.language('en'))(self.message))
+        self.assertTrue((Filters.language('en_GB'))(self.message))
+        self.assertFalse((Filters.language('da'))(self.message))
+        self.message.from_user.language_code = 'da'
+        self.assertFalse((Filters.language('en_US'))(self.message))
+        self.assertFalse((Filters.language('en'))(self.message))
+        self.assertFalse((Filters.language('en_GB'))(self.message))
+        self.assertTrue((Filters.language('da'))(self.message))
+
+    def test_language_filter_multiple(self):
+        f = Filters.language(['en_US', 'da'])
+        self.message.from_user.language_code = 'en_US'
+        self.assertTrue(f(self.message))
+        self.message.from_user.language_code = 'en_GB'
+        self.assertFalse(f(self.message))
+        self.message.from_user.language_code = 'da'
+        self.assertTrue(f(self.message))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_user.py b/tests/test_user.py
@@ -37,13 +37,15 @@ def setUp(self):
         self.first_name = "Leandro"
         self.last_name = "S."
         self.username = "leandrotoledo"
+        self.language_code = "pt-BR"
         self.type = "private"
 
         self.json_dict = {
             'id': self._id,
             'first_name': self.first_name,
             'last_name': self.last_name,
             'username': self.username,
+            'language_code': self.language_code,
             'type': self.type
         }
 
@@ -54,6 +56,7 @@ def test_user_de_json(self):
         self.assertEqual(user.first_name, self.first_name)
         self.assertEqual(user.last_name, self.last_name)
         self.assertEqual(user.username, self.username)
+        self.assertEqual(user.language_code, self.language_code)
         self.assertEqual(user.type, self.type)
 
         self.assertEqual(user.name, '@leandrotoledo')
@@ -98,6 +101,7 @@ def test_user_to_dict(self):
         self.assertEqual(user['first_name'], self.first_name)
         self.assertEqual(user['last_name'], self.last_name)
         self.assertEqual(user['username'], self.username)
+        self.assertEqual(user['language_code'], self.language_code)
         self.assertEqual(user['type'], self.type)
 
     @flaky(3, 1)
