Improve MixinWSApi.on_message

Author: learnforpractice

pysrc/mixin_msg_test.py renamed to examples/mixin_msg_test.py

@@ -4,29 +4,29 @@
 import asyncio
 import base64
 import logging
+from typing import Optional
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(module)s %(lineno)d %(message)s')
 logger = logging.getLogger(__name__)
 
-from pymixin.mixin_ws_api import MixinWSApi
+from pymixin.mixin_ws_api import MixinWSApi, MessageView, Category, MessageStatus
 
 bot_config = None
 
 with open(sys.argv[1]) as f:
     bot_config = f.read()
     bot_config = json.loads(bot_config)
 
+
 class MixinBot(MixinWSApi):
     def __init__(self):
         super().__init__(bot_config, on_message=self.on_message)
 
-    async def on_message(self, message):
-        logger.info(message)
-        if 'error' in message:
-            return
-        action = message["action"]
+    async def on_message(self, id: str, action: str, msg: Optional[MessageView]) -> None:
+        logger.info("on_message: %s %s %s", id, action, msg)
+
         if action not in ["ACKNOWLEDGE_MESSAGE_RECEIPT", "CREATE_MESSAGE", "LIST_PENDING_MESSAGES"]:
-            logger.info("unknow action %s", action)
+            logger.info("unknown action %s", action)
             return
 
         if action == "ACKNOWLEDGE_MESSAGE_RECEIPT":
@@ -35,37 +35,29 @@ async def on_message(self, message):
         if not action == "CREATE_MESSAGE":
             return
 
-        data = message["data"]
-        msgid = data["message_id"]
-        typeindata = data["type"]
-        categoryindata = data["category"]
-        user_id = data["user_id"]
-        conversation_id = data["conversation_id"]
+        if not msg:
+            return
 
-        created_at = data["created_at"]
-        updated_at = data["updated_at"]
+        msgid = msg.message_id
+        created_at = msg.created_at
+        updated_at = msg.updated_at
 
         await self.echoMessage(msgid)
 
-        logger.info('user_id %s', user_id)
+        logger.info('user_id %s', msg.user_id)
         logger.info("created_at %s",created_at)
 
-        if 'error' in message:
+        if not msg.category in ["SYSTEM_ACCOUNT_SNAPSHOT", "PLAIN_TEXT", "SYSTEM_CONVERSATION", "PLAIN_STICKER", "PLAIN_IMAGE", "PLAIN_CONTACT"]:
+            logger.info("unknown category: %s", msg.category)
             return
 
-        if not categoryindata in ["SYSTEM_ACCOUNT_SNAPSHOT", "PLAIN_TEXT", "SYSTEM_CONVERSATION", "PLAIN_STICKER", "PLAIN_IMAGE", "PLAIN_CONTACT"]:
-            logger.info("unknown category: %s", categoryindata)
+        if not msg.category == "PLAIN_TEXT" and msg.type == "message":
             return
 
-        if not categoryindata == "PLAIN_TEXT" and typeindata == "message":
-            return
-
-        data = data["data"]
-        logger.info(data)
-        data = base64.urlsafe_b64decode(data)
+        logger.info(msg.data)
+        data = base64.urlsafe_b64decode(msg.data)
         logger.info(data)
-        data = data.decode()
-        await self.sendUserText(conversation_id, user_id, data)
+        await self.sendUserText(msg.conversation_id, msg.user_id, data.decode())
 
 bot = MixinBot()
 

examples/ws_example.py

@@ -311,7 +311,7 @@ async def send_messages(self, messages):
         '''
         return await self.post('/messages', messages)
 
-    async def send_message(self, conversation_id: str, category: str, data: str):
+    async def send_message(self, conversation_id: str, category: str, data: Union[str, bytes]):
         if isinstance(data, str):
             data = data.encode()
         msg = {
@@ -338,24 +338,21 @@ async def send_sticker_message(self, conversation_id: str, name: str, album_id:
             "album_id": album_id,
             "sticker_id": sticker_id
         }
-        data = json.dumps(data)
-        return await self.send_message(conversation_id, "PLAIN_STICKER", data)
+        return await self.send_message(conversation_id, "PLAIN_STICKER", json.dumps(data))
 
     async def send_contract_message(self, conversation_id: str, user_id: str):
         data = {
             "user_id": user_id
         }
-        data = json.dumps(data)
-        return await self.send_message(conversation_id, "PLAIN_CONTACT", data)
+        return await self.send_message(conversation_id, "PLAIN_CONTACT", json.dumps(data))
 
     async def send_button_group_message(self, conversation_id: str, label: str, color: str, action: str):
         data = {
             "label": label,
             "action": action,
             "color": color
         }
-        data = json.dumps([data])
-        return await self.send_message(conversation_id, "APP_BUTTON_GROUP", data)
+        return await self.send_message(conversation_id, "APP_BUTTON_GROUP", json.dumps([data]))
 
     @staticmethod
     def _convert_object_to_dict(x):

pysrc/mixin_bot_api.py

@@ -8,23 +8,63 @@
 import uuid
 import gzip
 import time
-from typing import Union
+from typing import Union, Callable, Awaitable, Optional, Literal
 from io import BytesIO
 import base64
 import websockets
 
-from pymixin import mixin_config
-from pymixin.mixin_bot_api import MixinBotApi
-from pymixin import log
+from dataclasses import dataclass
+from dataclasses_json import dataclass_json
+
+from . import mixin_config
+from .mixin_bot_api import MixinBotApi
+from . import log
 
 logger = log.get_logger(__name__)
 
-class MixinWSApi:
+Category = Literal[
+        'PLAIN_TEXT',
+        'PLAIN_AUDIO',
+        'PLAIN_POST',
+        'PLAIN_IMAGE',
+        'PLAIN_DATA',
+        'PLAIN_STICKER',
+        'PLAIN_LIVE',
+        'PLAIN_LOCATION',
+        'PLAIN_VIDEO',
+        'PLAIN_CONTACT',
+        'APP_CARD',
+        'APP_BUTTON_GROUP',
+        'MESSAGE_RECALL',
+        'SYSTEM_CONVERSATION',
+        'SYSTEM_ACCOUNT_SNAPSHOT'
+    ]
+
+MessageStatus = Literal['SENT', 'DELIVERED', 'READ']
+
+@dataclass_json
+@dataclass
+class MessageView:
+    type: str
+    representative_id: str
+    quote_message_id: str
+    conversation_id: str
+    user_id: str
+    session_id: str
+    message_id: str
+    category: Category
+    data: str
+    data_base64: str
+    status: MessageStatus
+    source: str
+    created_at: str
+    updated_at: str
 
-    def __init__(self, bot_config, on_message):
+class MixinWSApi:
+    def __init__(self, bot_config, on_message: Callable[[str, str, Optional[MessageView]], Awaitable[None]]):
         self.bot = MixinBotApi(bot_config)
         self.ws = None
-        self.on_message = on_message
+        self._on_message = on_message
 
     async def connect(self):
         if self.ws:
@@ -67,7 +107,11 @@ async def run(self):
                 msg = gzip.GzipFile(mode="rb", fileobj=msg)
                 msg = msg.read()
                 msg = json.loads(msg)
-                await self.on_message(msg)
+                try:
+                    view = MessageView.from_dict(msg['data'])
+                except KeyError:
+                    view = None
+                await self._on_message(msg['id'], msg['action'], view)
 
     """
     =================

pysrc/mixin_ws_api.py

@@ -5,6 +5,7 @@ python-dateutil>=2.7.5
 requests>=2.21.0
 websocket-client>=0.54.0
 websockets
+dataclasses-json
 
 scikit-build
 cython

requirements-dev.txt

@@ -5,4 +5,5 @@ python-dateutil>=2.7.5
 requests>=2.21.0
 websockets>=9.1
 cryptography>=3.4.7
+dataclasses-json