From 4291309abac8c91d82dd9a54fe5d7148cab90ae6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B4=9B=E6=B0=B4=E5=B1=85=E5=AE=A4?=
 <luoshuijs@outlook.com>
Date: Mon, 1 May 2023 21:35:51 +0800
Subject: [PATCH] :sparkles: Add timeout configuration for the `BaseClient`

---
 simnet/client/base.py      | 59 +++++++++++++++++---------------------
 simnet/client/starrail.pyi |  3 +-
 simnet/utils/types.py      |  5 ++++
 3 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/simnet/client/base.py b/simnet/client/base.py
index 7ed041f..faaba18 100644
--- a/simnet/client/base.py
+++ b/simnet/client/base.py
@@ -3,7 +3,7 @@ import uuid
 from types import TracebackType
 from typing import AsyncContextManager, Type, Optional, Any
 
-from httpx import AsyncClient, TimeoutException, Response, HTTPError
+from httpx import AsyncClient, TimeoutException, Response, HTTPError, Timeout
 
 from simnet.client.cookies import Cookies
 from simnet.client.headers import Headers
@@ -16,46 +16,32 @@ from simnet.utils.types import (
     CookieTypes,
     RequestData,
     QueryParamTypes,
+    TimeoutTypes,
 )
 
-_LOGGER = logging.getLogger("simnet.BaseClient")
+_LOGGER = logging.getLogger("SIMNet.BaseClient")
 
 
 class BaseClient(AsyncContextManager["BaseClient"]):
     """
     This is the base class for simnet clients. It provides common methods and properties for simnet clients.
 
-    Parameters:
-    -----------
-    cookies: Optional[CookieTypes]
-        The cookies used for the client.
-    headers: Optional[HeaderTypes]
-        The headers used for the client.
-    account_id: Optional[int]
-        The account id used for the client.
-    player_id: Optional[int]
-        The player id used for the client.
-    region: Region
-        The region used for the client.
-    lang: str
-        The language used for the client.
+    Args:
+        cookies (Optional[CookieTypes], optional): The cookies used for the client.
+        headers (Optional[HeaderTypes], optional): The headers used for the client.
+        account_id (Optional[int], optional): The account id used for the client.
+        player_id (Optional[int], optional): The player id used for the client.
+        region (Region, optional): The region used for the client.
+        lang (str, optional): The language used for the client.
+        timeout (Optional[TimeoutTypes], optional): Timeout configuration for the client.
 
     Attributes:
-    -----------
-    cookies: Cookies
-        The cookies used for the client.
-    headers: Headers
-        The headers used for the client.
-    player_id: Optional[int]
-        The player id used for the client.
-    account_id: Optional[int]
-        The account id used for the client.
-    client: AsyncClient
-        The httpx async client instance.
-    region: Region
-        The region used for the client.
-    lang: str
-        The language used for the client.
+        cookies (CookieTypes): The cookies used for the client.
+        headers (HeaderTypes): The headers used for the client.
+        account_id (Optional[int]): The account id used for the client.
+        player_id (Optional[int]): The player id used for the client.
+        region (Region): The region used for the client.
+        lang (str): The language used for the client.
     """
 
     _device_id = str(uuid.uuid3(uuid.NAMESPACE_URL, "SIMNet"))
@@ -68,13 +54,22 @@ class BaseClient(AsyncContextManager["BaseClient"]):
         player_id: Optional[int] = None,
         region: Region = Region.OVERSEAS,
         lang: str = "en-us",
+        timeout: Optional[TimeoutTypes] = None,
     ) -> None:
         """Initialize the client with the given parameters."""
+        if timeout is None:
+            timeout = Timeout(
+                connect=5.0,
+                read=5.0,
+                write=5.0,
+                pool=1.0,
+            )
+
         self.cookies = Cookies(cookies)
         self.headers = Headers(headers)
         self.player_id = player_id
         self.account_id = account_id
-        self.client = AsyncClient(cookies=self.cookies)
+        self.client = AsyncClient(cookies=self.cookies, timeout=timeout)
         self.region = region
         self.lang = lang
 
diff --git a/simnet/client/starrail.pyi b/simnet/client/starrail.pyi
index 1f15b80..b14e94d 100644
--- a/simnet/client/starrail.pyi
+++ b/simnet/client/starrail.pyi
@@ -3,7 +3,7 @@ from typing import Optional
 from simnet.client.chronicle.starrail import StarRailBattleChronicleClient
 from simnet.client.wish.starrail import WishClient
 from simnet.utils.enum_ import Region
-from simnet.utils.types import CookieTypes, HeaderTypes
+from simnet.utils.types import CookieTypes, HeaderTypes, TimeoutTypes
 
 
 class StarRailClient(StarRailBattleChronicleClient, WishClient):
@@ -15,4 +15,5 @@ class StarRailClient(StarRailBattleChronicleClient, WishClient):
         player_id: Optional[int] = None,
         region: Region = Region.OVERSEAS,
         lang: str = "en-us",
+        timeout: Optional[TimeoutTypes] = None,
     ): ...
diff --git a/simnet/utils/types.py b/simnet/utils/types.py
index b142069..82048f3 100644
--- a/simnet/utils/types.py
+++ b/simnet/utils/types.py
@@ -21,3 +21,8 @@ HeaderTypes = Union[
     Sequence[Tuple[str, str]],
     Sequence[Tuple[bytes, bytes]],
 ]
+TimeoutTypes = Union[
+    Optional[float],
+    Tuple[Optional[float], Optional[float], Optional[float], Optional[float]],
+    "Timeout",
+]