format_platform

noneplugin · Sep 7, 2024 · 180d80c · 180d80c
1 parent 66600cc
commit 180d80c
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 19 deletions.
diff --git a/nonebot_plugin_chatrecorder/adapters/onebot_v12.py b/nonebot_plugin_chatrecorder/adapters/onebot_v12.py
@@ -9,7 +9,7 @@
 from typing_extensions import override
 
 from ..config import plugin_config
-from ..consts import SupportedAdapter, SupportedPlatform
+from ..consts import SupportedAdapter
 from ..message import (
     MessageDeserializer,
     MessageSerializer,
@@ -18,6 +18,7 @@
     serialize_message,
 )
 from ..model import MessageRecord
+from ..utils import format_platform
 
 try:
     from nonebot.adapters.onebot.v12 import Bot, Message, MessageEvent
@@ -67,16 +68,10 @@ async def record_send_msg(
             elif detail_type == "private":
                 level = SessionLevel.LEVEL1
 
-            platform = bot.platform
-            if platform in ("kook",):
-                platform = SupportedPlatform.kaiheila
-            elif platform not in list(SupportedPlatform):
-                platform = SupportedPlatform.unknown
-
             session = Session(
                 bot_id=bot.self_id,
                 bot_type=bot.type,
-                platform=bot.platform,
+                platform=format_platform(bot.platform),
                 level=level,
                 id1=data.get("user_id"),
                 id2=data.get("group_id") or data.get("channel_id"),

diff --git a/nonebot_plugin_chatrecorder/adapters/satori.py b/nonebot_plugin_chatrecorder/adapters/satori.py
@@ -9,7 +9,7 @@
 from typing_extensions import override
 
 from ..config import plugin_config
-from ..consts import SupportedAdapter, SupportedPlatform
+from ..consts import SupportedAdapter
 from ..message import (
     MessageDeserializer,
     MessageSerializer,
@@ -18,7 +18,7 @@
     serialize_message,
 )
 from ..model import MessageRecord
-from ..utils import remove_timezone
+from ..utils import format_platform, remove_timezone
 
 try:
     from nonebot.adapters.satori import Bot, Message, MessageEvent
@@ -76,18 +76,10 @@ async def record_send_msg(
             id2 = result_message.channel.id if result_message.channel else None
             id3 = result_message.guild.id if result_message.guild else None
 
-            platform = bot.platform
-            if platform in ("onebot", "red", "chronocat"):
-                platform = SupportedPlatform.qq
-            elif platform in ("kook",):
-                platform = SupportedPlatform.kaiheila
-            elif platform not in list(SupportedPlatform):
-                platform = SupportedPlatform.unknown
-
             session = Session(
                 bot_id=bot.self_id,
                 bot_type=bot.type,
-                platform=platform,
+                platform=format_platform(bot.platform),
                 level=level,
                 id1=id1,
                 id2=id2,

diff --git a/nonebot_plugin_chatrecorder/utils.py b/nonebot_plugin_chatrecorder/utils.py
@@ -1,5 +1,7 @@
 from datetime import datetime, timezone
 
+from .consts import SupportedPlatform
+
 
 def remove_timezone(dt: datetime) -> datetime:
     """移除时区"""
@@ -8,3 +10,13 @@ def remove_timezone(dt: datetime) -> datetime:
     # 先转至 UTC 时间，再移除时区
     dt = dt.astimezone(timezone.utc)
     return dt.replace(tzinfo=None)
+
+
+def format_platform(platform: str) -> str:
+    if platform in ("onebot", "red", "chronocat"):
+        return SupportedPlatform.qq
+    elif platform in ("kook",):
+        return SupportedPlatform.kaiheila
+    elif platform not in list(SupportedPlatform):
+        return SupportedPlatform.unknown
+    return platform