From e7eea2a7d97f4fb72a803503ef2e60898bca1e17 Mon Sep 17 00:00:00 2001
From: HibiKier <45528451+HibiKier@users.noreply.github.com>
Date: Mon, 4 Oct 2021 22:57:59 +0800
Subject: [PATCH] Update __init__.py

---
 plugins/yiqing/__init__.py | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/plugins/yiqing/__init__.py b/plugins/yiqing/__init__.py
index 760babb3..8c8926fe 100644
--- a/plugins/yiqing/__init__.py
+++ b/plugins/yiqing/__init__.py
@@ -1,9 +1,10 @@
 from nonebot import on_command
-from .data_source import get_yiqing_data
+from .data_source import get_yiqing_data, get_city_and_province_list
 from services.log import logger
 from nonebot.adapters.cqhttp import Bot, MessageEvent, GroupMessageEvent
 from nonebot.typing import T_State
 from utils.utils import get_message_text
+from configs.config import NICKNAME
 
 
 __zx_plugin_name__ = "疫情查询"
@@ -13,7 +14,7 @@ usage：
     指令：
         疫情 中国
         疫情 [省份/城市]
-    * 当省份与城市重名时，可在后添加 “市” *
+    * 当省份与城市重名时，可在后添加 "市" 或 "省" *
     示例：疫情 吉林   <- [省]
     示例：疫情 吉林市  <- [市]
 """.strip()
@@ -36,17 +37,21 @@ yiqing = on_command("疫情", aliases={"查询疫情", "疫情查询"}, priority
 @yiqing.handle()
 async def _(bot: Bot, event: MessageEvent, state: T_State):
     msg = get_message_text(event.json())
+    city_and_province_list = get_city_and_province_list()
     if msg:
-        result = await get_yiqing_data(msg)
-        if result:
-            await yiqing.send(result)
-            logger.info(
-                f"(USER {event.user_id}, GROUP "
-                f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'}) 查询疫情: {msg}"
-            )
+        if msg in city_and_province_list or msg[:-1] in city_and_province_list:
+            result = await get_yiqing_data(msg)
+            if result:
+                await yiqing.send(result)
+                logger.info(
+                    f"(USER {event.user_id}, GROUP "
+                    f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'}) 查询疫情: {msg}"
+                )
+            else:
+                await yiqing.send("查询失败!!!!", at_sender=True)
+                logger.info(
+                    f"(USER {event.user_id}, GROUP "
+                    f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'}) 查询疫情失败"
+                )
         else:
-            await yiqing.send("查询失败!!!!", at_sender=True)
-            logger.info(
-                f"(USER {event.user_id}, GROUP "
-                f"{event.group_id if isinstance(event, GroupMessageEvent) else 'private'}) 查询疫情失败"
-            )
+            await yiqing.send(f"{NICKNAME}只支持国内的疫情查询喔...")