Make the Message attribute names align with the original json from received event (#252)

Author: o926428377

linebot/models/base.py

@@ -86,11 +86,15 @@ def as_json_dict(self):
         return data
 
     @classmethod
-    def new_from_json_dict(cls, data):
+    def new_from_json_dict(cls, data, use_raw_message=False):
         """Create a new instance from a dict.
 
         :param data: JSON dict
+        :param bool use_raw_message: Using original Message key as attribute
         """
+        if use_raw_message:
+            return cls(use_raw_message=use_raw_message, **data)
+
         new_data = {utils.to_snake_case(key): value
                     for key, value in data.items()}
 
@@ -117,7 +121,7 @@ def get_or_new_from_json_dict(data, cls):
 
     @staticmethod
     def get_or_new_from_json_dict_with_types(
-            data, cls_map, type_key='type'
+            data, cls_map, type_key='type', use_raw_message=False
     ):
         """Get `cls` object w/ deserialization from json by using type key hint if needed.
 
@@ -129,12 +133,13 @@ def get_or_new_from_json_dict_with_types(
         :param cls_map:
         :param type_key:
         :rtype: object
+        :param bool use_raw_message: Using original Message key as attribute
         """
         if isinstance(data, tuple(cls_map.values())):
             return data
         elif isinstance(data, dict):
             type_val = data[type_key]
             if type_val in cls_map:
-                return cls_map[type_val].new_from_json_dict(data)
+                return cls_map[type_val].new_from_json_dict(data, use_raw_message=use_raw_message)
 
         return None

linebot/models/events.py

@@ -103,7 +103,7 @@ class MessageEvent(Event):
     """
 
     def __init__(self, mode=None, timestamp=None, source=None, reply_token=None, message=None,
-                 **kwargs):
+                 use_raw_message=False, **kwargs):
         """__init__ method.
 
         :param str mode: Channel state
@@ -113,6 +113,7 @@ def __init__(self, mode=None, timestamp=None, source=None, reply_token=None, mes
         :param str reply_token: Reply token
         :param message: Message object
         :type message: T <= :py:class:`linebot.models.messages.Message`
+        :param bool use_raw_message: Using original Message key as attribute
         :param kwargs:
         """
         super(MessageEvent, self).__init__(
@@ -130,7 +131,7 @@ def __init__(self, mode=None, timestamp=None, source=None, reply_token=None, mes
                 'location': LocationMessage,
                 'sticker': StickerMessage,
                 'file': FileMessage
-            }
+            }, use_raw_message=use_raw_message
         )
 
 

linebot/models/messages.py

@@ -35,17 +35,28 @@
 class Message(with_metaclass(ABCMeta, Base)):
     """Abstract Base Class of Message."""
 
-    def __init__(self, id=None, **kwargs):
+    def __init__(self, id=None, use_raw_message=False, **kwargs):
         """__init__ method.
 
         :param str id: Message ID
+        :param bool use_raw_message: Using original Message key as attribute
         :param kwargs:
         """
         super(Message, self).__init__(**kwargs)
 
+        if use_raw_message:
+            self.__dict__.update(kwargs)
+
         self.type = None
         self.id = id
 
+    def __getitem__(self, key):
+        """__getitem__ method.
+
+        :param str key: Message key
+        """
+        return self.__dict__.get(key, None)
+
 
 @deprecated(reason="Use 'from linebot.v3.webhooks import TextMessageContent' instead. See https://github.com/line/line-bot-sdk-python/blob/master/README.rst for more details.", version='3.0.0', category=LineBotSdkDeprecatedIn30)  # noqa: E501
 class TextMessage(Message):

linebot/webhook.py

@@ -138,14 +138,15 @@ def __init__(self, channel_secret):
         """
         self.signature_validator = SignatureValidator(channel_secret)
 
-    def parse(self, body, signature, as_payload=False):
+    def parse(self, body, signature, as_payload=False, use_raw_message=False):
         """Parse webhook request body as text.
 
         :param str body: Webhook request body (as text)
         :param str signature: X-Line-Signature value (as text)
         :param bool as_payload: (optional) True to return WebhookPayload object.
         :rtype: list[T <= :py:class:`linebot.models.events.Event`]
             | :py:class:`linebot.webhook.WebhookPayload`
+        :param bool use_raw_message: Using original Message key as attribute
         :return: Events list, or WebhookPayload instance
         """
         if not self.signature_validator.validate(body, signature):
@@ -157,7 +158,8 @@ def parse(self, body, signature, as_payload=False):
         for event in body_json['events']:
             event_type = event['type']
             if event_type == 'message':
-                events.append(MessageEvent.new_from_json_dict(event))
+                events.append(MessageEvent.new_from_json_dict(event,
+                              use_raw_message=use_raw_message))
             elif event_type == 'follow':
                 events.append(FollowEvent.new_from_json_dict(event))
             elif event_type == 'unfollow':
@@ -244,13 +246,15 @@ def decorator(func):
 
         return decorator
 
-    def handle(self, body, signature):
+    def handle(self, body, signature, use_raw_message=False):
         """Handle webhook.
 
         :param str body: Webhook request body (as text)
         :param str signature: X-Line-Signature value (as text)
+        :param bool use_raw_message: Using original Message key as attribute
         """
-        payload = self.parser.parse(body, signature, as_payload=True)
+        payload = self.parser.parse(body, signature, as_payload=True,
+                                    use_raw_message=use_raw_message)
 
         for event in payload.events:
             func = None

tests/test_webhook.py

@@ -20,7 +20,7 @@
 import inspect
 
 from linebot import (
-    SignatureValidator, WebhookParser, WebhookHandler
+    SignatureValidator, WebhookParser, WebhookHandler, utils
 )
 from linebot.models import (
     MessageEvent, FollowEvent, UnfollowEvent, JoinEvent,
@@ -614,6 +614,9 @@ def test_parse_webhook_req_without_destination(self):
 class TestWebhookHandler(unittest.TestCase):
     def setUp(self):
         self.handler = WebhookHandler('channel_secret')
+        self.use_raw_message = True
+        self.retrieve_attr_name = lambda x: x if self.use_raw_message \
+            else utils.to_snake_case(x)
 
         @self.handler.add(MessageEvent, message=TextMessage)
         def message_text(event, destination):
@@ -635,6 +638,13 @@ def message_sticker(event):
             self.assertEqual('message', event.type)
             self.assertEqual('sticker', event.message.type)
 
+        @self.handler.add(MessageEvent, message=FileMessage)
+        def message_file(event):
+            self.assertEqual('message', event.type)
+            self.assertEqual('file', event.message.type)
+            self.assertNotEqual(event.message[self.retrieve_attr_name("fileName")], None)
+            self.assertNotEqual(event.message[self.retrieve_attr_name("fileSize")], None)
+
         @self.handler.add(MessageEvent)
         def message(event):
             self.assertEqual('message', event.type)
@@ -680,7 +690,7 @@ def test_handler(self):
         # mock
         self.handler.parser.signature_validator.validate = lambda a, b: True
 
-        self.handler.handle(body, 'signature')
+        self.handler.handle(body, 'signature', self.use_raw_message)
 
 
 class TestInvokeWebhookHandler(unittest.TestCase):