From e496bf82e68af08312cc43a7c46f470da008991a Mon Sep 17 00:00:00 2001
From: suyiiyii <suyiiyii@gmail.com>
Date: Thu, 5 Dec 2024 16:08:42 +0800
Subject: [PATCH] =?UTF-8?q?:bug:=20=E4=BF=AE=E5=A4=8D=20cookie=20=E6=A8=A1?=
=?UTF-8?q?=E5=9D=97=20type=20hint=20(#658)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
nonebot_bison/admin_page/api.py | 14 ++++++++++----
nonebot_bison/config/db_config.py | 4 ++++
nonebot_bison/config/subs_io/nbesf_model/v3.py | 7 +++++--
nonebot_bison/config/subs_io/subs_io.py | 2 +-
nonebot_bison/config/subs_io/utils.py | 6 ++++--
nonebot_bison/script/cli.py | 2 +-
nonebot_bison/sub_manager/add_cookie_target.py | 6 ++----
nonebot_bison/sub_manager/utils.py | 8 ++++----
nonebot_bison/utils/site.py | 8 ++++----
tests/subs_io/test_cli.py | 2 +-
tests/subs_io/test_subs_io.py | 6 +++---
11 files changed, 39 insertions(+), 26 deletions(-)
diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py
index 4fbc758..66f431d 100644
--- a/nonebot_bison/admin_page/api.py
+++ b/nonebot_bison/admin_page/api.py
@@ -24,6 +24,7 @@ from ..utils.site import CookieClientManager, site_manager, is_cookie_client_man
from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config
from .types import (
Cookie,
+ Target,
TokenResp,
GlobalConf,
SiteConfig,
@@ -213,7 +214,7 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We
@router.get("/cookie", dependencies=[Depends(check_is_superuser)])
-async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]:
+async def get_cookie(site_name: str | None = None, target: str | None = None) -> list[Cookie]:
cookies_in_db = await config.get_cookie(site_name, is_anonymous=False)
return [
Cookie(
@@ -252,7 +253,12 @@ async def get_cookie_target(
cookie_targets = await config.get_cookie_target()
# TODO: filter in SQL
return [
- x
+ CookieTarget(
+ target=Target(
+ platform_name=x.target.platform_name, target_name=x.target.target_name, target=x.target.target
+ ),
+ cookie_id=x.cookie.id,
+ )
for x in cookie_targets
if (site_name is None or x.cookie.site_name == site_name)
and (target is None or x.target.target == target)
@@ -261,13 +267,13 @@ async def get_cookie_target(
@router.post("/cookie_target", dependencies=[Depends(check_is_superuser)])
-async def add_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp:
+async def add_cookie_target(platform_name: str, target: T_Target, cookie_id: int) -> StatusResp:
await config.add_cookie_target(target, platform_name, cookie_id)
return StatusResp(ok=True, msg="")
@router.delete("/cookie_target", dependencies=[Depends(check_is_superuser)])
-async def del_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp:
+async def del_cookie_target(platform_name: str, target: T_Target, cookie_id: int) -> StatusResp:
await config.delete_cookie_target(target, platform_name, cookie_id)
return StatusResp(ok=True, msg="")
diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py
index 7451292..c76078c 100644
--- a/nonebot_bison/config/db_config.py
+++ b/nonebot_bison/config/db_config.py
@@ -288,6 +288,8 @@ class DBConfig:
async def get_cookie_by_id(self, cookie_id: int) -> Cookie:
async with create_session() as sess:
cookie = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
+ if not cookie:
+ raise NoSuchTargetException(f"cookie {cookie_id} not found")
return cookie
async def add_cookie(self, cookie: Cookie) -> int:
@@ -317,6 +319,8 @@ class DBConfig:
.outerjoin(CookieTarget)
.options(selectinload(Cookie.targets))
)
+ if not cookie:
+ raise NoSuchTargetException(f"cookie {cookie_id} not found")
if len(cookie.targets) > 0:
raise Exception(f"cookie {cookie.id} in use")
await sess.execute(delete(Cookie).where(Cookie.id == cookie_id))
diff --git a/nonebot_bison/config/subs_io/nbesf_model/v3.py b/nonebot_bison/config/subs_io/nbesf_model/v3.py
index a2d5e42..c6f1c03 100644
--- a/nonebot_bison/config/subs_io/nbesf_model/v3.py
+++ b/nonebot_bison/config/subs_io/nbesf_model/v3.py
@@ -8,8 +8,11 @@ from pydantic import BaseModel
from nonebot_plugin_saa.registries import AllSupportedPlatformTarget
from nonebot.compat import PYDANTIC_V2, ConfigDict, model_dump, type_validate_json, type_validate_python
+from nonebot_bison.types import Tag
+from nonebot_bison.types import Category
+from nonebot_bison.types import Target as T_Target
+
from ..utils import NBESFParseErr
-from ....types import Tag, Category
from .base import NBESFBase, SubReceipt
from ...db_model import Cookie as DBCookie
from ...db_config import SubscribeDupException, config
@@ -114,7 +117,7 @@ async def magic_cookie_gen(nbesf_data: SubGroup):
new_cookie = DBCookie(**model_dump(cookie, exclude={"targets"}))
cookie_id = await config.add_cookie(new_cookie)
for target in cookie.targets:
- await config.add_cookie_target(target.target, target.platform_name, cookie_id)
+ await config.add_cookie_target(T_Target(target.target), target.platform_name, cookie_id)
except Exception as e:
logger.error(f"！添加 Cookie 条目 {repr(cookie)} 失败: {repr(e)}")
else:
diff --git a/nonebot_bison/config/subs_io/subs_io.py b/nonebot_bison/config/subs_io/subs_io.py
index 9b9ce48..c871096 100644
--- a/nonebot_bison/config/subs_io/subs_io.py
+++ b/nonebot_bison/config/subs_io/subs_io.py
@@ -65,7 +65,7 @@ async def subscribes_export(selector: Callable[[Select], Select]) -> v3.SubGroup
target_payload = type_validate_python(v3.Target, cookie_target.target)
cookie_target_dict[cookie_target.cookie].append(target_payload)
- def cookie_transform(cookie: Cookie, targets: [Target]) -> v3.Cookie:
+ def cookie_transform(cookie: Cookie, targets: list[v3.Target]) -> v3.Cookie:
cookie_dict = row2dict(cookie)
cookie_dict["tags"] = cookie.tags
cookie_dict["targets"] = targets
diff --git a/nonebot_bison/config/subs_io/utils.py b/nonebot_bison/config/subs_io/utils.py
index 1ba7558..21c3a5a 100644
--- a/nonebot_bison/config/subs_io/utils.py
+++ b/nonebot_bison/config/subs_io/utils.py
@@ -1,4 +1,6 @@
-from ..db_model import Model
+from typing import Any
+
+from sqlalchemy.orm import DeclarativeBase
class NBESFVerMatchErr(Exception): ...
@@ -7,7 +9,7 @@ class NBESFVerMatchErr(Exception): ...
class NBESFParseErr(Exception): ...
-def row2dict(row: Model) -> dict:
+def row2dict(row: DeclarativeBase) -> dict[str, Any]:
d = {}
for column in row.__table__.columns:
d[column.name] = str(getattr(row, column.name))
diff --git a/nonebot_bison/script/cli.py b/nonebot_bison/script/cli.py
index 790a7dd..f17e431 100644
--- a/nonebot_bison/script/cli.py
+++ b/nonebot_bison/script/cli.py
@@ -83,7 +83,7 @@ async def subs_export(path: Path, format: str):
export_file = path / f"bison_subscribes_export_{int(time.time())}.{format}"
logger.info("正在获取订阅信息...")
- export_data: v2.SubGroup = await subscribes_export(lambda x: x)
+ export_data: v3.SubGroup = await subscribes_export(lambda x: x)
with export_file.open("w", encoding="utf-8") as f:
match format:
diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py
index 8fd4780..919d0ca 100644
--- a/nonebot_bison/sub_manager/add_cookie_target.py
+++ b/nonebot_bison/sub_manager/add_cookie_target.py
@@ -26,8 +26,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]):
@add_cookie_target_matcher.got("target_idx", parameterless=[handle_cancel])
async def got_target_idx(state: T_State, target_idx: str = ArgPlainText()):
try:
- target_idx = int(target_idx)
- state["target"] = state["sub_table"][target_idx]
+ state["target"] = state["sub_table"][int(target_idx)]
state["site"] = platform_manager[state["target"]["platform_name"]].site
except Exception:
await add_cookie_target_matcher.reject("序号错误")
@@ -57,8 +56,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]):
@add_cookie_target_matcher.got("cookie_idx", MessageTemplate("{_prompt}"), [handle_cancel])
async def got_cookie_idx(state: T_State, cookie_idx: str = ArgPlainText()):
try:
- cookie_idx = int(cookie_idx)
- state["cookie"] = state["cookies"][cookie_idx - 1]
+ state["cookie"] = state["cookies"][int(cookie_idx) - 1]
except Exception:
await add_cookie_target_matcher.reject("序号错误")
diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py
index 8126b9a..2418cfc 100644
--- a/nonebot_bison/sub_manager/utils.py
+++ b/nonebot_bison/sub_manager/utils.py
@@ -13,6 +13,7 @@ from nonebot_plugin_saa import PlatformTarget, extract_target
from ..config import config
from ..types import Category
+from ..types import Target as T_Target
from ..platform import platform_manager
from ..plugin_config import plugin_config
from ..utils.site import is_cookie_client_manager
@@ -88,7 +89,7 @@ async def generate_sub_list_text(
sub_list = [
sub
for sub in sub_list
- if is_cookie_client_manager(platform_manager.get(sub.target.platform_name).site.client_mgr)
+ if is_cookie_client_manager(platform_manager[sub.target.platform_name].site.client_mgr)
]
if not sub_list:
await matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅")
@@ -109,7 +110,7 @@ async def generate_sub_list_text(
res += " {}".format(", ".join(sub.tags)) + "\n"
if is_show_cookie:
target_cookies = await config.get_cookie(
- target=sub.target.target, site_name=platform.site.name, is_anonymous=False
+ target=T_Target(sub.target.target), site_name=platform.site.name, is_anonymous=False
)
if target_cookies:
res += " 关联的 Cookie：\n"
@@ -126,6 +127,5 @@ async def only_allow_private(
event: Event,
matcher: type[Matcher],
):
- # if not issubclass(PrivateMessageEvent, event.__class__):
- if event.message_type != "private":
+ if not (hasattr(event, "message_type") and getattr(event, "message_type") == "private"):
await matcher.finish("请在私聊中使用此命令")
diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py
index 8dcd371..aecb00f 100644
--- a/nonebot_bison/utils/site.py
+++ b/nonebot_bison/utils/site.py
@@ -123,8 +123,8 @@ class CookieClientManager(ClientManager):
async def _choose_cookie(self, target: Target | None) -> Cookie:
"""选择 cookie 的具体算法"""
cookies = await config.get_cookie(self._site_name, target)
- cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
- cookie = min(cookies, key=lambda x: x.last_usage)
+ avaliable_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
+ cookie = min(avaliable_cookies, key=lambda x: x.last_usage)
return cookie
async def get_client(self, target: Target | None) -> AsyncClient:
@@ -183,8 +183,8 @@ class SiteMeta(type):
cls._key = kwargs.get("key")
elif not kwargs.get("abstract"):
# this is the subclass
- if hasattr(cls, "name"):
- site_manager[cls.name] = cls
+ if "name" in namespace:
+ site_manager[namespace["name"]] = cls
super().__init__(name, bases, namespace, **kwargs)
diff --git a/tests/subs_io/test_cli.py b/tests/subs_io/test_cli.py
index d74bac5..61fe271 100644
--- a/tests/subs_io/test_cli.py
+++ b/tests/subs_io/test_cli.py
@@ -78,7 +78,7 @@ async def test_subs_export(app: App, tmp_path: Path):
cookie_name="test cookie",
)
)
- await config.add_cookie_target("weibo_id", "weibo", cookie_id)
+ await config.add_cookie_target(TTarget("weibo_id"), "weibo", cookie_id)
assert len(await config.list_subs_with_all_info()) == 3
diff --git a/tests/subs_io/test_subs_io.py b/tests/subs_io/test_subs_io.py
index 700becf..5164672 100644
--- a/tests/subs_io/test_subs_io.py
+++ b/tests/subs_io/test_subs_io.py
@@ -16,7 +16,7 @@ async def test_subs_export(app: App, init_scheduler):
await config.add_subscribe(
TargetQQGroup(group_id=1232),
- target=TTarget("weibo_id"),
+ target=TTarget(TTarget("weibo_id")),
target_name="weibo_name",
platform_name="weibo",
cats=[],
@@ -24,7 +24,7 @@ async def test_subs_export(app: App, init_scheduler):
)
await config.add_subscribe(
TargetQQGroup(group_id=2342),
- target=TTarget("weibo_id"),
+ target=TTarget(TTarget("weibo_id")),
target_name="weibo_name",
platform_name="weibo",
cats=[],
@@ -45,7 +45,7 @@ async def test_subs_export(app: App, init_scheduler):
cookie_name="test cookie",
)
)
- await config.add_cookie_target("weibo_id", "weibo", cookie_id)
+ await config.add_cookie_target(TTarget("weibo_id"), "weibo", cookie_id)
data = await config.list_subs_with_all_info()
assert len(data) == 3