From 6116e735cc2e49c76252c8bfd7723f3133bbdf2f Mon Sep 17 00:00:00 2001
From: Yifei Kong <kong@yifei.me>
Date: Sat, 18 Nov 2023 15:23:37 +0800
Subject: [PATCH] Set connect_time only in stream mode, and timeout=None will
 wait indefinitely

---
 curl_cffi/requests/session.py | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/curl_cffi/requests/session.py b/curl_cffi/requests/session.py
index 0d933a82..4d27c093 100644
--- a/curl_cffi/requests/session.py
+++ b/curl_cffi/requests/session.py
@@ -126,6 +126,9 @@ def _peek_aio_queue(q: asyncio.Queue, default=None):
         return default
 
 
+not_set = object()
+
+
 class BaseSession:
     """Provide common methods for setting curl options and reading info in sessions."""
 
@@ -136,11 +139,11 @@ class BaseSession:
         "proxies",
         "params",
         "verify",
+        "timeout",
         "cert",
         "trust_env",  # TODO
         "max_redirects",
         "impersonate",
-        "timeout",
     ]
 
     def __init__(
@@ -190,7 +193,7 @@ def _set_curl_options(
         cookies: Optional[CookieTypes] = None,
         files: Optional[Dict] = None,
         auth: Optional[Tuple[str, str]] = None,
-        timeout: Optional[Union[float, Tuple[float, float]]] = None,
+        timeout: Optional[Union[float, Tuple[float, float], object]] = not_set,
         allow_redirects: bool = True,
         max_redirects: Optional[int] = None,
         proxies: Optional[dict] = None,
@@ -301,14 +304,23 @@ def _set_curl_options(
             c.setopt(CurlOpt.PASSWORD, password.encode())  # type: ignore
 
         # timeout
-        timeout = timeout or self.timeout
+        if timeout is not_set:
+            timeout = self.timeout
+        if timeout is None:
+            timeout = 0  # indefinitely
+
         if isinstance(timeout, tuple):
             connect_timeout, read_timeout = timeout
             all_timeout = connect_timeout + read_timeout
             c.setopt(CurlOpt.CONNECTTIMEOUT_MS, int(connect_timeout * 1000))
-            c.setopt(CurlOpt.TIMEOUT_MS, int(all_timeout * 1000))
+            if not stream:
+                c.setopt(CurlOpt.TIMEOUT_MS, int(all_timeout * 1000))
         else:
-            c.setopt(CurlOpt.TIMEOUT_MS, int(timeout * 1000))
+            if not stream:
+                c.setopt(CurlOpt.TIMEOUT_MS, int(timeout * 1000))  # type: ignore
+            else:
+                c.setopt(CurlOpt.CONNECTTIMEOUT_MS, int(timeout * 1000))  # type: ignore
+
 
         # allow_redirects
         c.setopt(CurlOpt.FOLLOWLOCATION, int(allow_redirects))
@@ -487,7 +499,7 @@ def __init__(
             proxies: dict of proxies to use, format: {"http": proxy_url, "https": proxy_url}.
             params: query string for the session.
             verify: whether to verify https certs.
-            timeout: how many seconds to wait before giving up.
+            timeout: how many seconds to wait before giving up. In stream mode, only connect_timeout will be set.
             trust_env: use http_proxy/https_proxy and other environments, default True.
             max_redirects: max redirect counts, default unlimited(-1).
             impersonate: which browser version to impersonate in the session.