使用闭包实现client hook

This commit is contained in:
suyiiyii 2024-09-03 10:52:20 +08:00
parent 4ce6b85f79
commit 055ed6e02a
2 changed files with 28 additions and 6 deletions

View File

@ -22,8 +22,9 @@ class ProcessContext:
async def _log_to_ctx(r: Response):
self._log_response(r)
existing_hooks = client.event_hooks["response"]
hooks = {
"response": [_log_to_ctx],
"response": [*existing_hooks, _log_to_ctx],
}
client.event_hooks = hooks

View File

@ -4,10 +4,12 @@ from abc import ABC, abstractmethod
import httpx
from httpx import AsyncClient
from nonebot.log import logger
from ..types import Target
from ..config import config
from .http import http_client
from ..config.db_model import Cookie
class ClientManager(ABC):
@ -41,22 +43,41 @@ class DefaultClientManager(ClientManager):
class CookieClientManager(ClientManager):
_platform_name: str
async def _choose_cookie(self, target: Target) -> dict[str, str]:
def _generate_hook(self, cookie: Cookie):
async def _response_hook(resp: httpx.Response):
if not cookie:
logger.debug(f"未携带bison cookie: {resp.request.url}")
if resp.status_code == 200:
logger.debug(f"请求成功: {cookie.id} {resp.request.url}")
else:
logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}")
return _response_hook
async def _choose_cookie(self, target: Target) -> Cookie:
if not target:
return {}
return Cookie(content="{}")
cookies = await config.get_cookie_by_target(target, self._platform_name)
if not cookies:
return {}
return Cookie(content="{}")
cookie = sorted(cookies, key=lambda x: x.last_usage, reverse=True)[0]
return json.loads(cookie.content)
return cookie
async def get_client(self, target: Target | None) -> AsyncClient:
client = http_client()
cookie = await self._choose_cookie(target)
if cookie.content != "{}":
logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}")
else:
logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie")
cookies = httpx.Cookies()
if cookie:
cookies.update(cookie)
cookies.update(json.loads(cookie.content))
client.cookies = cookies
client._bison_cookie = cookie
client.event_hooks = {"response": [self._generate_hook(cookie)]}
return client
async def get_client_for_static(self) -> AsyncClient: