From e6f17c434381f30f5cd146c62bcd33deee8e3eff Mon Sep 17 00:00:00 2001
From: AkashiCoin <l1040186796@gmail.com>
Date: Mon, 16 Sep 2024 19:41:57 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=A6=84=20refactor(http=5Futils):=20?=
 =?UTF-8?q?=E9=83=A8=E5=88=86=E9=87=8D=E6=9E=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 zhenxun/utils/http_utils.py | 96 ++++++++++++++++++++++++-------------
 1 file changed, 63 insertions(+), 33 deletions(-)

diff --git a/zhenxun/utils/http_utils.py b/zhenxun/utils/http_utils.py
index 2de1c43a..1891e265 100644
--- a/zhenxun/utils/http_utils.py
+++ b/zhenxun/utils/http_utils.py
@@ -56,28 +56,61 @@ class AsyncHttpx:
             proxy: 指定代理
             timeout: 超时时间
         """
-        if not isinstance(url, list):
-            url = [url]
+        urls = [url] if isinstance(url, str) else url
+        return await cls._get_first_successful(
+            urls,
+            params=params,
+            headers=headers,
+            cookies=cookies,
+            verify=verify,
+            use_proxy=use_proxy,
+            proxy=proxy,
+            timeout=timeout,
+            **kwargs,
+        )
+
+    @classmethod
+    async def _get_first_successful(
+        cls,
+        urls: list[str],
+        **kwargs,
+    ) -> Response:
+        last_exception = None
+        for url in urls:
+            try:
+                return await cls._get_single(url, **kwargs)
+            except Exception as e:
+                last_exception = e
+                if url != urls[-1]:
+                    logger.warning(f"获取 {url} 失败, 尝试下一个")
+        raise last_exception or Exception("All URLs failed")
+
+    @classmethod
+    async def _get_single(
+        cls,
+        url: str,
+        *,
+        params: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+        cookies: dict[str, str] | None = None,
+        verify: bool = True,
+        use_proxy: bool = True,
+        proxy: dict[str, str] | None = None,
+        timeout: int = 30,
+        **kwargs,
+    ) -> Response:
         if not headers:
             headers = get_user_agent()
-        last_exception = Exception
         _proxy = proxy if proxy else cls.proxy if use_proxy else None
-        for u in url:
-            try:
-                async with httpx.AsyncClient(proxies=_proxy, verify=verify) as client:  # type: ignore
-                    return await client.get(
-                        u,
-                        params=params,
-                        headers=headers,
-                        cookies=cookies,
-                        timeout=timeout,
-                        **kwargs,
-                    )
-            except Exception:
-                last_exception = Exception
-                if u != url[-1]:
-                    logger.warning(f"获取 {u} 失败, 尝试下一个")
-        raise last_exception
+        async with httpx.AsyncClient(proxies=_proxy, verify=verify) as client:  # type: ignore
+            return await client.get(
+                url,
+                params=params,
+                headers=headers,
+                cookies=cookies,
+                timeout=timeout,
+                **kwargs,
+            )
 
     @classmethod
     async def head(
@@ -208,8 +241,8 @@ class AsyncHttpx:
                 if not isinstance(url, list):
                     url = [url]
                 for u in url:
-                    if not stream:
-                        try:
+                    try:
+                        if not stream:
                             response = await cls.get(
                                 u,
                                 params=params,
@@ -225,17 +258,14 @@ class AsyncHttpx:
                             content = response.content
                             async with aiofiles.open(path, "wb") as wf:
                                 await wf.write(content)
-                                logger.info(
-                                    f"下载 {url} 成功.. Path：{path.absolute()}"
-                                )
+                                logger.info(f"下载 {u} 成功.. Path：{path.absolute()}")
                             return True
-                        except (TimeoutError, ConnectTimeout, HTTPStatusError):
-                            pass
-                    else:
-                        if not headers:
-                            headers = get_user_agent()
-                        _proxy = proxy if proxy else cls.proxy if use_proxy else None
-                        try:
+                        else:
+                            if not headers:
+                                headers = get_user_agent()
+                            _proxy = (
+                                proxy if proxy else cls.proxy if use_proxy else None
+                            )
                             async with httpx.AsyncClient(
                                 proxies=_proxy,  # type: ignore
                                 verify=verify,
@@ -279,8 +309,8 @@ class AsyncHttpx:
                                             f"Path：{path.absolute()}"
                                         )
                             return True
-                        except (TimeoutError, ConnectTimeout, HTTPStatusError):
-                            pass
+                    except (TimeoutError, ConnectTimeout, HTTPStatusError):
+                        logger.warning(f"下载 {u} 失败.. 尝试下一个地址..")
             else:
                 logger.error(f"下载 {url} 下载超时.. Path：{path.absolute()}")
         except Exception as e: