refactor code

plugins/yiqing/__init__.py

@@ -1,27 +1,27 @@
 from nonebot import on_command
 from .data_source import get_yiqing_data, clear_data
 from services.log import logger
-from nonebot.adapters.cqhttp import Bot, Event
+from nonebot.adapters.cqhttp import Bot, MessageEvent, GroupMessageEvent
 from nonebot.typing import T_State
 from .config import city_list
 from utils.utils import scheduler
 
-__plugin_name__ = '疫情查询'
-__plugin_usage__ = '查询疫情帮助:\n\t对我说 查询疫情 省份/城市，我会回复疫情的实时数据\n\t示例: 查询疫情 温州'
+__plugin_name__ = "疫情查询"
+__plugin_usage__ = "查询疫情帮助:\n\t对我说 查询疫情 省份/城市，我会回复疫情的实时数据\n\t示例: 查询疫情 温州"
 
 
 yiqing = on_command("疫情", aliases={"查询疫情", "疫情查询"}, priority=5, block=True)
 
 
 @yiqing.handle()
-async def _(bot: Bot, event: Event, state: T_State):
-    if str(event.get_message()).strip() in ['帮助']:
-        await yiqing.finish(__plugin_usage__)
+async def _(bot: Bot, event: MessageEvent, state: T_State):
     msg = str(event.get_message()).strip()
+    if not msg or msg in ["帮助"]:
+        await yiqing.finish(__plugin_usage__)
     if msg:
         if msg in city_list.keys():
             province = msg
-            city = ''
+            city = ""
         else:
             for key in city_list.keys():
                 if msg in city_list.get(key):
@@ -35,17 +35,22 @@ async def _(bot: Bot, event: Event, state: T_State):
         if result:
             await yiqing.send(result)
             logger.info(
-                f"(USER {event.user_id}, GROUP {event.group_id if event.message_type != 'private' else 'private'}) 查询疫情:" + result)
+                f"(USER {event.user_id}, GROUP "
+                f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'}) 查询疫情:"
+                + result
+            )
         else:
             await yiqing.send("查询失败!!!!", at_sender=True)
             logger.info(
-                f"(USER {event.user_id}, GROUP {event.group_id if event.message_type != 'private' else 'private'}) 查询疫情失败")
+                f"(USER {event.user_id}, GROUP "
+                f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'}) 查询疫情失败"
+            )
     except UnboundLocalError:
-        await yiqing.finish('参数正确吗？只要一个参数啊', at_sender=True)
+        await yiqing.finish("参数正确吗？只要一个参数啊", at_sender=True)
 
 
 @scheduler.scheduled_job(
-    'cron',
+    "cron",
     hour=0,
     minute=1,
 )

plugins/yiqing/data_source.py

@@ -10,32 +10,50 @@ from utils.utils import get_local_proxy
 url = "https://api.yimian.xyz/coro/"
 
 
-async def get_yiqing_data(province, city_=''):
+async def get_yiqing_data(province, city_=""):
     if not os.path.exists(TXT_PATH + "yiqing/"):
         os.mkdir(TXT_PATH + "yiqing/")
     if not os.path.exists(TXT_PATH + "yiqing/" + str(datetime.now().date()) + ".json"):
         async with aiohttp.ClientSession(headers=get_user_agent()) as session:
             async with session.get(url, proxy=get_local_proxy(), timeout=7) as response:
                 datalist = await response.json()
-        with open(TXT_PATH + "yiqing/" + str(datetime.now().date()) + ".json", 'w') as f:
+        with open(
+            TXT_PATH + "yiqing/" + str(datetime.now().date()) + ".json", "w"
+        ) as f:
             json.dump(datalist, f)
-    datalist = json.load(open(TXT_PATH + "yiqing/" + str(datetime.now().date()) + ".json", 'r'))
-    result = ''
+    datalist = json.load(
+        open(TXT_PATH + "yiqing/" + str(datetime.now().date()) + ".json", "r")
+    )
+    result = ""
     for data in datalist:
-        if data['provinceShortName'] == province:
-            if city_ == '':
-                result = province + "疫情数据:\n现存确诊: " + \
-                         str(data['currentConfirmedCount']) + "\n累计确诊: " + \
-                         str(data['confirmedCount']) + "\n治愈: " + \
-                         str(data['curedCount']) + "\n死亡: " + \
-                         str(data['deadCount'])
+        if data["provinceShortName"] == province:
+            if city_ == "":
+                result = (
+                    province
+                    + "疫情数据:\n现存确诊: "
+                    + str(data["currentConfirmedCount"])
+                    + "\n累计确诊: "
+                    + str(data["confirmedCount"])
+                    + "\n治愈: "
+                    + str(data["curedCount"])
+                    + "\n死亡: "
+                    + str(data["deadCount"])
+                )
                 break
             else:
-                for city in data['cities']:
-                    if city['cityName'] == city_:
-                        result = city_ + "疫情数据:\n现存确诊: " + \
-                                 str(city['currentConfirmedCount']) + "\n累计确诊: " + str(city['confirmedCount']) +\
-                                 "\n治愈: " + str(city['curedCount']) + "\n死亡: " + str(city['deadCount'])
+                for city in data["cities"]:
+                    if city["cityName"] == city_:
+                        result = (
+                            city_
+                            + "疫情数据:\n现存确诊: "
+                            + str(city["currentConfirmedCount"])
+                            + "\n累计确诊: "
+                            + str(city["confirmedCount"])
+                            + "\n治愈: "
+                            + str(city["curedCount"])
+                            + "\n死亡: "
+                            + str(city["deadCount"])
+                        )
                         break
     return result
 
@@ -45,6 +63,5 @@ def clear_data():
         os.remove(TXT_PATH + "yiqing/" + file)
 
 
-if __name__ == '__main__':
-    print(get_yiqing_data("浙江", city_=''))
-
+if __name__ == "__main__":
+    print(get_yiqing_data("浙江", city_=""))