♻️ refactor client of scheduler_config

nonebot_bison/utils/__init__.py

@@ -11,14 +11,16 @@ from nonebot_plugin_saa import Text, Image, MessageSegmentFactory
 from .http import http_client
 from .context import ProcessContext
 from ..plugin_config import plugin_config
-from .scheduler_config import SchedulerConfig, scheduler
 from .image import pic_merge, text_to_image, is_pics_mergable, pic_url_to_image
+from .scheduler_config import ClientManager, SchedulerConfig, DefaultClientManager, scheduler
 
 __all__ = [
     "http_client",
     "Singleton",
     "parse_text",
     "ProcessContext",
+    "ClientManager",
+    "DefaultClientManager",
     "html_to_text",
     "SchedulerConfig",
     "scheduler",

nonebot_bison/utils/context.py

@@ -2,19 +2,25 @@ from base64 import b64encode
 
 from httpx import Response, AsyncClient
 
+from nonebot_bison.types import Target
+
+from .scheduler_config import ClientManager
+
 
 class ProcessContext:
     reqs: list[Response]
+    _client_mgr: ClientManager
 
-    def __init__(self) -> None:
+    def __init__(self, client_mgr: ClientManager) -> None:
         self.reqs = []
+        self._client_mgr = client_mgr
 
-    def log_response(self, resp: Response):
+    def _log_response(self, resp: Response):
         self.reqs.append(resp)
 
-    def register_to_client(self, client: AsyncClient):
+    def _register_to_client(self, client: AsyncClient):
         async def _log_to_ctx(r: Response):
-            self.log_response(r)
+            self._log_response(r)
 
         hooks = {
             "response": [_log_to_ctx],
@@ -41,3 +47,16 @@ class ProcessContext:
                 )
             res.append(log_content)
         return res
+
+    async def get_client(self, target: Target | None = None) -> AsyncClient:
+        client = await self._client_mgr.get_client(target)
+        self._register_to_client(client)
+        return client
+
+    async def get_client_for_static(self) -> AsyncClient:
+        client = await self._client_mgr.get_client_for_static()
+        self._register_to_client(client)
+        return client
+
+    async def refresh_client(self):
+        await self._client_mgr.refresh_client()

nonebot_bison/utils/scheduler_config.py

@@ -1,3 +1,4 @@
+from abc import ABC
 from typing import Literal
 
 from httpx import AsyncClient
@@ -6,10 +7,32 @@ from ..types import Target
 from .http import http_client
 
 
+class ClientManager(ABC):
+    async def get_client(self, target: Target | None) -> AsyncClient: ...
+
+    async def get_client_for_static(self) -> AsyncClient: ...
+
+    async def get_query_name_client(self) -> AsyncClient: ...
+
+    async def refresh_client(self): ...
+
+
+class DefaultClientManager(ClientManager):
+    async def get_client(self, target: Target | None) -> AsyncClient:
+        return http_client()
+
+    async def get_client_for_static(self) -> AsyncClient:
+        return http_client()
+
+    async def get_query_name_client(self) -> AsyncClient:
+        return http_client()
+
+
 class SchedulerConfig:
     schedule_type: Literal["date", "interval", "cron"]
     schedule_setting: dict
     name: str
+    client_mgr: type[ClientManager] = DefaultClientManager
     require_browser: bool = False
 
     def __str__(self):
@@ -18,12 +41,6 @@ class SchedulerConfig:
     def __init__(self):
         self.default_http_client = http_client()
 
-    async def get_client(self, target: Target) -> AsyncClient:
-        return self.default_http_client
-
-    async def get_query_name_client(self) -> AsyncClient:
-        return self.default_http_client
-
 
 def scheduler(schedule_type: Literal["date", "interval", "cron"], schedule_setting: dict) -> type[SchedulerConfig]:
     return type(
@@ -32,5 +49,6 @@ def scheduler(schedule_type: Literal["date", "interval", "cron"], schedule_setti
         {
             "schedule_type": schedule_type,
             "schedule_setting": schedule_setting,
+            "client_mgr": ClientManager,
         },
     )