refactor: 优化代码写法

nfe-w · Nov 23, 2024 · d8d4e50 · d8d4e50
1 parent 1b7cc10
commit d8d4e50
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 40 deletions.
diff --git a/push_channel/__init__.py b/push_channel/__init__.py
@@ -14,32 +14,25 @@
 
 push_channel_dict: dict[str, PushChannel] = {}
 
+_channel_type_to_class = {
+    "serverChan_turbo": ServerChanTurbo,
+    "wecom_apps": WeComApps,
+    "wecom_bot": WeComBot,
+    "dingtalk_bot": DingtalkBot,
+    "feishu_apps": FeishuApps,
+    "feishu_bot": FeishuBot,
+    "telegram_bot": TelegramBot,
+    "bark": Bark,
+    "gotify": Gotify,
+    "webhook": Webhook,
+    "email": Email,
+    "demo": Demo,
+}
+
 
 def get_push_channel(config) -> PushChannel:
     channel_type = config.get("type", None)
-    if channel_type == "serverChan_turbo":
-        return ServerChanTurbo(config)
-    if channel_type == "wecom_apps":
-        return WeComApps(config)
-    if channel_type == "wecom_bot":
-        return WeComBot(config)
-    if channel_type == "dingtalk_bot":
-        return DingtalkBot(config)
-    if channel_type == "feishu_apps":
-        return FeishuApps(config)
-    if channel_type == "feishu_bot":
-        return FeishuBot(config)
-    if channel_type == "telegram_bot":
-        return TelegramBot(config)
-    if channel_type == "bark":
-        return Bark(config)
-    if channel_type == "gotify":
-        return Gotify(config)
-    if channel_type == "webhook":
-        return Webhook(config)
-    if channel_type == "email":
-        return Email(config)
-    if channel_type == "demo":
-        return Demo(config)
+    if channel_type is None or channel_type not in _channel_type_to_class:
+        raise ValueError(f"不支持的通道类型: {channel_type}")
 
-    raise ValueError(f"不支持的通道类型: {channel_type}")
+    return _channel_type_to_class[channel_type](config)
diff --git a/query_task/__init__.py b/query_task/__init__.py
@@ -7,22 +7,20 @@
 from .query_weibo import QueryWeibo
 from .query_xhs import QueryXhs
 
+_task_type_to_class = {
+    "bilibili": QueryBilibili,
+    "weibo": QueryWeibo,
+    "xhs": QueryXhs,
+    "douyin": QueryDouyin,
+    "douyu": QueryDouyu,
+    "huya": QueryHuya,
+    "demo": QueryDemo,
+}
+
 
 def get_query_task(config) -> QueryTask:
     task_type = config.get("type", None)
-    if task_type == "bilibili":
-        return QueryBilibili(config)
-    if task_type == "weibo":
-        return QueryWeibo(config)
-    if task_type == "xhs":
-        return QueryXhs(config)
-    if task_type == "douyin":
-        return QueryDouyin(config)
-    if task_type == "douyu":
-        return QueryDouyu(config)
-    if task_type == "huya":
-        return QueryHuya(config)
-    if task_type == "demo":
-        return QueryDemo(config)
+    if task_type is None or task_type not in _task_type_to_class:
+        raise ValueError(f"不支持的查询任务: {task_type}")
 
-    raise ValueError(f"不支持的查询任务: {task_type}")
+    return _task_type_to_class[task_type](config)