From 60be1eb2bd16fc89827a50ad18377642aae176be Mon Sep 17 00:00:00 2001 From: Kyomotoi <0w0@imki.moe> Date: Thu, 6 Apr 2023 16:07:46 +0800 Subject: =?UTF-8?q?=F0=9F=8E=A8=20=E4=BC=98=E5=8C=96=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E5=BA=93=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ATRI/plugins/thesaurus/data_source.py | 46 +++++++++++++++------------------- ATRI/plugins/thesaurus/db.py | 47 ----------------------------------- 2 files changed, 20 insertions(+), 73 deletions(-) delete mode 100644 ATRI/plugins/thesaurus/db.py (limited to 'ATRI/plugins/thesaurus') diff --git a/ATRI/plugins/thesaurus/data_source.py b/ATRI/plugins/thesaurus/data_source.py index 4f99b2c..281a777 100644 --- a/ATRI/plugins/thesaurus/data_source.py +++ b/ATRI/plugins/thesaurus/data_source.py @@ -2,23 +2,23 @@ from datetime import datetime, timedelta, timezone as tz from ATRI.message import MessageBuilder from ATRI.exceptions import ThesaurusError +from ATRI.database import DatabaseWrapper, ThesaurusStoragor, ThesaurusAuditList -from .db import DBForTS, DBForTAL -from .db import ThesaurusStoragor + +DBForTS = DatabaseWrapper(ThesaurusStoragor) +DBForTAL = DatabaseWrapper(ThesaurusAuditList) class ThesaurusManager: async def __add_item(self, _id: str, group_id: int, is_main: bool = False): if is_main: try: - async with DBForTS() as db: - await db.add_item(_id, group_id) + await DBForTS.add_sub(_id=_id, group_id=group_id) except Exception: raise ThesaurusError(f"添加词库(ts)数据失败 目标词id: {_id}") else: try: - async with DBForTAL() as db: - await db.add_item(_id, group_id) + await DBForTAL.add_sub(_id=_id, group_id=group_id) except Exception: raise ThesaurusError(f"添加词库(tal)数据失败 目标词id: {_id}") @@ -27,56 +27,52 @@ class ThesaurusManager: ): if is_main: try: - async with DBForTS() as db: - await db.update_item(_id, group_id, update_map) + await DBForTS.update_sub( + update_map=update_map, _id=_id, group_id=group_id + ) except Exception: raise ThesaurusError(f"更新词库(ts)数据失败 目标词id: {_id}") else: try: - async with DBForTAL() as db: - await db.update_item(_id, group_id, update_map) + await DBForTAL.update_sub( + update_map=update_map, _id=_id, group_id=group_id + ) except Exception: raise ThesaurusError(f"更新词库(tal)数据失败 目标词id: {_id}") async def __del_item(self, _id: str, group_id: int, is_main: bool = False): if is_main: try: - async with DBForTS() as db: - await db.del_item({"_id": _id, "group_id": group_id}) + await DBForTS.del_sub({"_id": _id, "group_id": group_id}) except Exception: raise ThesaurusError(f"删除词库(ts)数据失败 目标词id: {_id}") else: try: - async with DBForTAL() as db: - await db.del_item({"_id": _id, "group_id": group_id}) + await DBForTAL.del_sub({"_id": _id, "group_id": group_id}) except Exception: raise ThesaurusError(f"删除词库(tal)数据失败 目标词id: {_id}") async def get_item_list(self, query_map: dict, is_main: bool = False) -> list: if is_main: try: - async with DBForTS() as db: - return await db.get_item_list(query_map) + return await DBForTS.get_sub_list(query_map) except Exception: raise ThesaurusError("获取词库(ts)列表数据失败") else: try: - async with DBForTAL() as db: - return await db.get_item_list(query_map) + return await DBForTAL.get_sub_list(query_map) except Exception: raise ThesaurusError("获取词库(tal)列表数据失败") async def get_all_items(self, is_main: bool = False) -> list: if is_main: try: - async with DBForTS() as db: - return await db.get_all_items() + return await DBForTS.get_all_subs() except Exception: raise ThesaurusError("获取全部词库(ts)列表数据失败") else: try: - async with DBForTAL() as db: - return await db.get_all_items() + return await DBForTAL.get_all_subs() except Exception: raise ThesaurusError("获取全部词库(tal)列表数据失败") @@ -165,8 +161,7 @@ class ThesaurusManager: class ThesaurusListener: async def get_item_by_id(self, _id: str) -> ThesaurusStoragor: try: - async with DBForTS() as db: - data = await db.get_item_list({"_id": _id}) + data = await DBForTS.get_sub_list({"_id": _id}) except Exception: raise ThesaurusError(f"获取词库(ts)数据失败 词条ID: {_id}") @@ -174,7 +169,6 @@ class ThesaurusListener: async def get_item_list(self, group_id: int): try: - async with DBForTS() as db: - return await db.get_item_list({"group_id": group_id}) + return await DBForTS.get_sub_list({"group_id": group_id}) except Exception: raise ThesaurusError(f"获取词库(ts)数据失败 目标群号: {group_id}") diff --git a/ATRI/plugins/thesaurus/db.py b/ATRI/plugins/thesaurus/db.py deleted file mode 100644 index b5394fc..0000000 --- a/ATRI/plugins/thesaurus/db.py +++ /dev/null @@ -1,47 +0,0 @@ -from ATRI.database import ThesaurusStoragor, ThesaurusAuditList - - -class DBForTS: - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - pass - - async def add_item(self, _id: str, group_id: int): - await ThesaurusStoragor.create(_id=_id, group_id=group_id) - - async def update_item(self, _id: str, group_id: int, update_map: dict): - await ThesaurusStoragor.filter(_id=_id, group_id=group_id).update(**update_map) - - async def del_item(self, query_map: dict): - await ThesaurusStoragor.filter(**query_map).delete() - - async def get_item_list(self, query_map: dict) -> list: - return await ThesaurusStoragor.filter(**query_map) - - async def get_all_items(self) -> list: - return await ThesaurusStoragor.all() - - -class DBForTAL: - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - pass - - async def add_item(self, _id: str, group_id: int): - await ThesaurusAuditList.create(_id=_id, group_id=group_id) - - async def update_item(self, _id: str, group_id: int, update_map: dict): - await ThesaurusAuditList.filter(_id=_id, group_id=group_id).update(**update_map) - - async def del_item(self, query_map: dict): - await ThesaurusAuditList.filter(**query_map).delete() - - async def get_item_list(self, query_map: dict) -> list: - return await ThesaurusAuditList.filter(**query_map) - - async def get_all_items(self) -> list: - return await ThesaurusAuditList.all() -- cgit v1.2.3