🏷️ update type hints

This commit is contained in:
suyiiyii 2024-11-25 14:35:32 +08:00
parent 630845fe0d
commit 2c605fc2dc
Signed by: suyiiyii
GPG Key ID: 044704CB29B8AD85
5 changed files with 23 additions and 11 deletions

View File

@ -22,6 +22,7 @@ from ..utils.site import CookieClientManager, site_manager, is_cookie_client_man
from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config
from .types import ( from .types import (
Cookie, Cookie,
Target,
TokenResp, TokenResp,
GlobalConf, GlobalConf,
SiteConfig, SiteConfig,
@ -211,7 +212,7 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We
@router.get("/cookie", dependencies=[Depends(check_is_superuser)]) @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) cookies_in_db = await config.get_cookie(site_name, is_anonymous=False)
return [ return [
Cookie( Cookie(
@ -250,7 +251,12 @@ async def get_cookie_target(
cookie_targets = await config.get_cookie_target() cookie_targets = await config.get_cookie_target()
# TODO: filter in SQL # TODO: filter in SQL
return [ 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 for x in cookie_targets
if (site_name is None or x.cookie.site_name == site_name) if (site_name is None or x.cookie.site_name == site_name)
and (target is None or x.target.target == target) and (target is None or x.target.target == target)
@ -259,13 +265,13 @@ async def get_cookie_target(
@router.post("/cookie_target", dependencies=[Depends(check_is_superuser)]) @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) await config.add_cookie_target(target, platform_name, cookie_id)
return StatusResp(ok=True, msg="") return StatusResp(ok=True, msg="")
@router.delete("/cookie_target", dependencies=[Depends(check_is_superuser)]) @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) await config.delete_cookie_target(target, platform_name, cookie_id)
return StatusResp(ok=True, msg="") return StatusResp(ok=True, msg="")

View File

@ -288,6 +288,8 @@ class DBConfig:
async def get_cookie_by_id(self, cookie_id: int) -> Cookie: async def get_cookie_by_id(self, cookie_id: int) -> Cookie:
async with create_session() as sess: async with create_session() as sess:
cookie = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) cookie = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
if not cookie:
raise NoSuchTargetException(f"cookie {cookie_id} not found")
return cookie return cookie
async def add_cookie(self, cookie: Cookie) -> int: async def add_cookie(self, cookie: Cookie) -> int:
@ -317,6 +319,8 @@ class DBConfig:
.outerjoin(CookieTarget) .outerjoin(CookieTarget)
.options(selectinload(Cookie.targets)) .options(selectinload(Cookie.targets))
) )
if not cookie:
raise NoSuchTargetException(f"cookie {cookie_id} not found")
if len(cookie.targets) > 0: if len(cookie.targets) > 0:
raise Exception(f"cookie {cookie.id} in use") raise Exception(f"cookie {cookie.id} in use")
await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) await sess.execute(delete(Cookie).where(Cookie.id == cookie_id))

View File

@ -8,9 +8,11 @@ from pydantic import BaseModel
from nonebot_plugin_saa.registries import AllSupportedPlatformTarget from nonebot_plugin_saa.registries import AllSupportedPlatformTarget
from nonebot.compat import PYDANTIC_V2, ConfigDict, model_dump, type_validate_json, type_validate_python from nonebot.compat import PYDANTIC_V2, ConfigDict, model_dump, type_validate_json, type_validate_python
from ....types import Tag
from ....types import Category
from ..utils import NBESFParseErr from ..utils import NBESFParseErr
from ....types import Tag, Category
from .base import NBESFBase, SubReceipt from .base import NBESFBase, SubReceipt
from ....types import Target as T_Target
from ...db_model import Cookie as DBCookie from ...db_model import Cookie as DBCookie
from ...db_config import SubscribeDupException, config from ...db_config import SubscribeDupException, config
@ -114,7 +116,7 @@ async def magic_cookie_gen(nbesf_data: SubGroup):
new_cookie = DBCookie(**model_dump(cookie, exclude={"targets"})) new_cookie = DBCookie(**model_dump(cookie, exclude={"targets"}))
cookie_id = await config.add_cookie(new_cookie) cookie_id = await config.add_cookie(new_cookie)
for target in cookie.targets: 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: except Exception as e:
logger.error(f"!添加 Cookie 条目 {repr(cookie)} 失败: {repr(e)}") logger.error(f"!添加 Cookie 条目 {repr(cookie)} 失败: {repr(e)}")
else: else:

View File

@ -123,8 +123,8 @@ class CookieClientManager(ClientManager):
async def _choose_cookie(self, target: Target | None) -> Cookie: async def _choose_cookie(self, target: Target | None) -> Cookie:
"""选择 cookie 的具体算法""" """选择 cookie 的具体算法"""
cookies = await config.get_cookie(self._site_name, target) cookies = await config.get_cookie(self._site_name, target)
cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now()) avaliable_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
cookie = min(cookies, key=lambda x: x.last_usage) cookie = min(avaliable_cookies, key=lambda x: x.last_usage)
return cookie return cookie
async def get_client(self, target: Target | None) -> AsyncClient: async def get_client(self, target: Target | None) -> AsyncClient:
@ -183,8 +183,8 @@ class SiteMeta(type):
cls._key = kwargs.get("key") cls._key = kwargs.get("key")
elif not kwargs.get("abstract"): elif not kwargs.get("abstract"):
# this is the subclass # this is the subclass
if hasattr(cls, "name"): if "name" in namespace:
site_manager[cls.name] = cls site_manager[namespace["name"]] = cls
super().__init__(name, bases, namespace, **kwargs) super().__init__(name, bases, namespace, **kwargs)

View File

@ -78,7 +78,7 @@ async def test_subs_export(app: App, tmp_path: Path):
cookie_name="test cookie", 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 assert len(await config.list_subs_with_all_info()) == 3