diff options
| author | Kyomotoi <0w0@imki.moe> | 2023-04-06 16:07:46 +0800 |
|---|---|---|
| committer | Kyomotoi <0w0@imki.moe> | 2023-04-06 16:07:46 +0800 |
| commit | 60be1eb2bd16fc89827a50ad18377642aae176be (patch) | |
| tree | 657a08ad802b2b67c2f5d941d93b67c67160b315 /ATRI/plugins/bilibili_dynamic | |
| parent | 170211bc72311fc4524e7c731303645eac2aa041 (diff) | |
| download | ATRI-60be1eb2bd16fc89827a50ad18377642aae176be.tar.gz ATRI-60be1eb2bd16fc89827a50ad18377642aae176be.tar.bz2 ATRI-60be1eb2bd16fc89827a50ad18377642aae176be.zip | |
🎨 优化数据库代码
Diffstat (limited to 'ATRI/plugins/bilibili_dynamic')
| -rw-r--r-- | ATRI/plugins/bilibili_dynamic/data_source.py | 18 | ||||
| -rw-r--r-- | ATRI/plugins/bilibili_dynamic/db.py | 26 |
2 files changed, 7 insertions, 37 deletions
diff --git a/ATRI/plugins/bilibili_dynamic/data_source.py b/ATRI/plugins/bilibili_dynamic/data_source.py index d629454..15babe5 100644 --- a/ATRI/plugins/bilibili_dynamic/data_source.py +++ b/ATRI/plugins/bilibili_dynamic/data_source.py @@ -5,8 +5,8 @@ from operator import itemgetter from ATRI.message import MessageBuilder from ATRI.utils import TimeDealer from ATRI.exceptions import BilibiliDynamicError +from ATRI.database import DatabaseWrapper, BilibiliSubscription -from .db import DB from .api import API @@ -17,27 +17,25 @@ _OUTPUT_FORMAT = ( .text("链接: {up_dy_link}") .done() ) +DB = DatabaseWrapper(BilibiliSubscription) class BilibiliDynamicSubscriptor: async def __add_sub(self, uid: int, group_id: int): try: - async with DB() as db: - await db.add_sub(uid, group_id) + await DB.add_sub(uid=uid, group_id=group_id) except Exception: raise BilibiliDynamicError("添加订阅失败") async def update_sub(self, uid: int, group_id: int, update_map: dict): try: - async with DB() as db: - await db.update_sub(uid, group_id, update_map) + await DB.update_sub(update_map=update_map, uid=uid, group_id=group_id) except Exception: raise BilibiliDynamicError("更新订阅失败") async def __del_sub(self, uid: int, group_id: int): try: - async with DB() as db: - await db.del_sub({"uid": uid, "group_id": group_id}) + await DB.del_sub({"uid": uid, "group_id": group_id}) except Exception: raise BilibiliDynamicError("删除订阅失败") @@ -48,15 +46,13 @@ class BilibiliDynamicSubscriptor: query_map = {"uid": uid, "group_id": group_id} try: - async with DB() as db: - return await db.get_sub_list(query_map) + return await DB.get_sub_list(query_map) except Exception: raise BilibiliDynamicError("获取订阅列表失败") async def get_all_subs(self) -> list: try: - async with DB() as db: - return await db.get_all_subs() + return await DB.get_all_subs() except Exception: raise BilibiliDynamicError("获取全部订阅列表失败") diff --git a/ATRI/plugins/bilibili_dynamic/db.py b/ATRI/plugins/bilibili_dynamic/db.py deleted file mode 100644 index e6bb8bc..0000000 --- a/ATRI/plugins/bilibili_dynamic/db.py +++ /dev/null @@ -1,26 +0,0 @@ -from ATRI.database import BilibiliSubscription - - -class DB: - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - pass - - async def add_sub(self, uid: int, group_id: int): - await BilibiliSubscription.create(uid=uid, group_id=group_id) - - async def update_sub(self, uid: int, group_id: int, update_map: dict): - await BilibiliSubscription.filter(uid=uid, group_id=group_id).update( - **update_map - ) - - async def del_sub(self, query_map: dict): - await BilibiliSubscription.filter(**query_map).delete() - - async def get_sub_list(self, query_map: dict) -> list: - return await BilibiliSubscription.filter(**query_map) - - async def get_all_subs(self) -> list: - return await BilibiliSubscription.all() |