Add curl_infos parameter for extra curl_info extraction after perform

lexiforest · Nov 18, 2023 · 8b0c96a · 8b0c96a
1 parent 6116e73
commit 8b0c96a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/curl_cffi/requests/models.py b/curl_cffi/requests/models.py
@@ -60,6 +60,7 @@ def __init__(self, curl: Optional[Curl] = None, request: Optional[Request] = Non
         self.redirect_url = ""
         self.http_version = 0
         self.history = []
+        self.infos = {}
         self.queue: Optional[queue.Queue] = None
         self.stream_task = None
 

diff --git a/curl_cffi/requests/session.py b/curl_cffi/requests/session.py
@@ -161,6 +161,7 @@ def __init__(
         impersonate: Optional[Union[str, BrowserType]] = None,
         default_headers: bool = True,
         curl_options: Optional[dict] = None,
+        curl_infos: Optional[list] = None,
         http_version: Optional[CurlHttpVersion] = None,
         debug: bool = False,
         interface: Optional[str] = None,
@@ -177,6 +178,7 @@ def __init__(
         self.impersonate = impersonate
         self.default_headers = default_headers
         self.curl_options = curl_options or {}
+        self.curl_infos = curl_infos or []
         self.http_version = http_version
         self.debug = debug
         self.interface = interface
@@ -321,7 +323,6 @@ def _set_curl_options(
             else:
                 c.setopt(CurlOpt.CONNECTTIMEOUT_MS, int(timeout * 1000))  # type: ignore
 
-
         # allow_redirects
         c.setopt(CurlOpt.FOLLOWLOCATION, int(allow_redirects))
 
@@ -467,6 +468,9 @@ def _parse_response(self, curl, buffer, header_buffer):
         rsp.redirect_count = cast(int, c.getinfo(CurlInfo.REDIRECT_COUNT))
         rsp.redirect_url = cast(bytes, c.getinfo(CurlInfo.REDIRECT_URL)).decode()
 
+        for info in self.curl_infos:
+            rsp.infos[info] = c.getinfo(info)
+
         return rsp
 
 

diff --git a/tests/unittest/test_requests.py b/tests/unittest/test_requests.py
@@ -5,7 +5,7 @@
 import pytest
 
 from curl_cffi import requests, CurlOpt
-from curl_cffi.const import CurlECode
+from curl_cffi.const import CurlECode, CurlInfo
 
 
 def test_head(server):
@@ -567,3 +567,11 @@ def test_stream_options_persist(server):
         buffer.append(line)
     data = json.loads(b"".join(buffer))
     assert data["User-agent"][0] == "foo/1.0"
+
+
+def test_curl_infos(server):
+    s = requests.Session(curl_infos=[CurlInfo.PRIMARY_IP])
+
+    r = s.get(str(server.url))
+
+    assert r.infos[CurlInfo.PRIMARY_IP] == b"127.0.0.1"