Custom response objects

curl_cffi/requests/session.py

@@ -21,6 +21,7 @@
     TypedDict,
     Union,
     cast,
+    Type,
 )
 from urllib.parse import ParseResult, parse_qsl, quote, unquote, urlencode, urljoin, urlparse
 
@@ -86,6 +87,7 @@ class BaseSessionParams(TypedDict, total=False):
         debug: bool
         interface: Optional[str]
         cert: Optional[Union[str, Tuple[str, str]]]
+        response_class: Optional[Type[Response]]
 
 else:
     ProxySpec = Dict[str, str]
@@ -221,6 +223,7 @@ def __init__(
         debug: bool = False,
         interface: Optional[str] = None,
         cert: Optional[Union[str, Tuple[str, str]]] = None,
+        response_class: Optional[Type[Response]] = None,
     ):
         self.headers = Headers(headers)
         self.cookies = Cookies(cookies)
@@ -245,6 +248,13 @@ def __init__(
         self.interface = interface
         self.cert = cert
 
+        if response_class is None:
+            response_class = Response
+        elif issubclass(response_class, Response) is False:
+            raise TypeError( "`response_class` must be a subclass of `curl_cffi.requests.models.Response`"
+                            f" not of type `{response_class}`" )
+        self.response_class = response_class
+
         if proxy and proxies:
             raise TypeError("Cannot specify both 'proxy' and 'proxies'")
         if proxy:
@@ -699,7 +709,7 @@ def qput(chunk):
 
     def _parse_response(self, curl, buffer, header_buffer, default_encoding):
         c = curl
-        rsp = Response(c)
+        rsp = self.response_class(c)
         rsp.url = cast(bytes, c.getinfo(CurlInfo.EFFECTIVE_URL)).decode()
         if buffer:
             rsp.content = buffer.getvalue()

examples/custom_response_class.py

@@ -0,0 +1,14 @@
+from curl_cffi import requests
+
+class CustomResponse(requests.Response):
+    @property
+    def status(self):
+        return self.status_code
+
+    def custom_method():
+        return "this is a custom method"
+
+session = requests.Session(response_class=CustomResponse)
+response: CustomResponse = session.get("http://example.com")
+print(response.status)
+print(response.custom_method())

tests/integration/test_response_class.py

@@ -0,0 +1,27 @@
+from curl_cffi import requests
+
+def test_default_response():
+    response = requests.get("http://example.com")
+    assert type(response) == requests.Response
+    print(response.status_code)
+
+class CustomResponse(requests.Response):
+    @property
+    def status(self):
+        return self.status_code
+
+def test_custom_response():
+    session = requests.Session(response_class=CustomResponse)
+    response = session.get("http://example.com")
+    assert isinstance(response, CustomResponse)
+    assert hasattr(response, "status")
+    print(response.status)
+
+class WrongTypeResponse: pass
+
+def test_wrong_type_custom_response():
+    try:
+        requests.Session(response_class=WrongTypeResponse)
+        assert False, "session was created without raising issue for wrong response class type"
+    except TypeError:
+        print("Wrong response class type detected")