Refactor Update.extract_ methods to Update.effective_ properties (python-telegram-bot#531)

jh0ker · web-flow · commit ff39e2436ef7 · 2017-03-26T14:36:34.000+02:00
* 🔨 Refactor `Update.extract_` methods to `Update.effective_` properties python-telegram-bot#507 * 🔨 handler.py: use effective_ properties
diff --git a/telegram/ext/conversationhandler.py b/telegram/ext/conversationhandler.py
@@ -145,7 +145,9 @@ def __init__(self,
                     raise ValueError("If 'per_chat=True', 'InlineQueryHandler' doesn't work")
 
     def _get_key(self, update):
-        chat, user = update.extract_chat_and_user()
+        chat = update.effective_chat
+        user = update.effective_user
+
         key = list()
 
         if self.per_chat:
diff --git a/telegram/ext/handler.py b/telegram/ext/handler.py
@@ -104,10 +104,11 @@ def collect_optional_args(self, dispatcher, update=None):
         if self.pass_job_queue:
             optional_args['job_queue'] = dispatcher.job_queue
         if self.pass_user_data or self.pass_chat_data:
-            chat, user = update.extract_chat_and_user()
+            chat = update.effective_chat
+            user = update.effective_user
 
             if self.pass_user_data:
-                optional_args['user_data'] = dispatcher.user_data[user.id]
+                optional_args['user_data'] = dispatcher.user_data[user.id if user else None]
 
             if self.pass_chat_data:
                 optional_args['chat_data'] = dispatcher.chat_data[chat.id if chat else None]
diff --git a/telegram/update.py b/telegram/update.py
@@ -73,6 +73,10 @@ def __init__(self,
         self.channel_post = channel_post
         self.edited_channel_post = edited_channel_post
 
+        self._effective_user = None
+        self._effective_chat = None
+        self._effective_message = None
+
     @staticmethod
     def de_json(data, bot):
         """
@@ -99,25 +103,23 @@ def de_json(data, bot):
 
         return Update(**data)
 
-    def extract_chat_and_user(self):
+    @property
+    def effective_user(self):
         """
-        Helper method to get the sender's chat and user objects from an arbitrary update.
-        Depending on the type of update, one of the available attributes ``message``,
-        ``edited_message`` or ``callback_query`` is used to determine the result.
-
-        Returns:
-            tuple: of (chat, user), with None-values if no object could not be found.
+        A property that contains the ``User`` that sent this update, no matter what kind of update
+        this is. Will be ``None`` for channel posts.
         """
+
+        if self._effective_user:
+            return self._effective_user
+
         user = None
-        chat = None
 
         if self.message:
             user = self.message.from_user
-            chat = self.message.chat
 
         elif self.edited_message:
             user = self.edited_message.from_user
-            chat = self.edited_message.chat
 
         elif self.inline_query:
             user = self.inline_query.from_user
@@ -127,51 +129,67 @@ def extract_chat_and_user(self):
 
         elif self.callback_query:
             user = self.callback_query.from_user
-            chat = self.callback_query.message.chat if self.callback_query.message else None
 
-        return chat, user
+        self._effective_user = user
+        return user
 
-    def extract_message_text(self):
+    @property
+    def effective_chat(self):
+        """
+        A property that contains the ``Chat`` that this update was sent in, no matter what kind of
+        update this is. Will be ``None`` for inline queries and chosen inline results.
         """
-        Helper method to get the message text from an arbitrary update.
-        Depending on the type of update, one of the available attributes ``message``,
-        ``edited_message`` or ``callback_query`` is used to determine the result.
 
-        Returns:
-            str: The extracted message text
+        if self._effective_chat:
+            return self._effective_chat
 
-        Raises:
-            ValueError: If no message text was found in the update
+        chat = None
 
-        """
         if self.message:
-            return self.message.text
+            chat = self.message.chat
+
         elif self.edited_message:
-            return self.edited_message.text
-        elif self.callback_query:
-            return self.callback_query.message.text
-        else:
-            raise ValueError("Update contains no message text.")
+            chat = self.edited_message.chat
 
-    def extract_entities(self):
-        """
-        Helper method to get parsed entities from an arbitrary update.
-        Depending on the type of update, one of the available attributes ``message``,
-        ``edited_message`` or ``callback_query`` is used to determine the result.
+        elif self.callback_query and self.callback_query.message:
+            chat = self.callback_query.message.chat
 
-        Returns:
-            dict[:class:`telegram.MessageEntity`, ``str``]: A dictionary of entities mapped to the
-                text that belongs to them, calculated based on UTF-16 codepoints.
+        elif self.channel_post:
+            chat = self.channel_post.chat
 
-        Raises:
-            ValueError: If no entities were found in the update
+        elif self.edited_channel_post:
+            chat = self.edited_channel_post.chat
 
+        self._effective_chat = chat
+        return chat
+
+    @property
+    def effective_message(self):
+        """
+        A property that contains the ``Message`` included in this update, no matter what kind
+        of update this is. Will be ``None`` for inline queries, chosen inline results and callback
+        queries from inline messages.
         """
+
+        if self._effective_message:
+            return self._effective_message
+
+        message = None
+
         if self.message:
-            return self.message.parse_entities()
+            message = self.message
+
         elif self.edited_message:
-            return self.edited_message.parse_entities()
+            message = self.edited_message
+
         elif self.callback_query:
-            return self.callback_query.message.parse_entities()
-        else:
-            raise ValueError("No message object found in self, therefore no entities available.")
+            message = self.callback_query.message
+
+        elif self.channel_post:
+            message = self.channel_post
+
+        elif self.edited_channel_post:
+            message = self.edited_channel_post
+
+        self._effective_message = message
+        return message
diff --git a/tests/test_update.py b/tests/test_update.py
@@ -76,16 +76,20 @@ def test_update_to_dict(self):
         self.assertEqual(update['update_id'], self.update_id)
         self.assertTrue(isinstance(update['message'], telegram.Message))
 
-    def test_extract_chat_and_user(self):
+    def test_effective_chat(self):
         update = telegram.Update.de_json(self.json_dict, self._bot)
-        chat, user = update.extract_chat_and_user()
+        chat = update.effective_chat
         self.assertEqual(update.message.chat, chat)
+
+    def test_effective_user(self):
+        update = telegram.Update.de_json(self.json_dict, self._bot)
+        user = update.effective_user
         self.assertEqual(update.message.from_user, user)
 
-    def test_extract_message_text(self):
+    def test_effective_message(self):
         update = telegram.Update.de_json(self.json_dict, self._bot)
-        text = update.extract_message_text()
-        self.assertEqual(update.message.text, text)
+        message = update.effective_message
+        self.assertEqual(update.message.text, message.text)
 
 
 if __name__ == '__main__':
