Make Request and Response picklable

Author: hannseman

httpx/_models.py

@@ -786,6 +786,13 @@ def __repr__(self) -> str:
 
 
 class Request:
+    __attrs__ = [
+        "_content",
+        "method",
+        "url",
+        "headers",
+    ]
+
     def __init__(
         self,
         method: typing.Union[str, bytes],
@@ -893,8 +900,27 @@ def __repr__(self) -> str:
         url = str(self.url)
         return f"<{class_name}({self.method!r}, {url!r})>"
 
+    def __getstate__(self) -> typing.Dict[str, typing.Any]:
+        return {attr: getattr(self, attr, None) for attr in self.__attrs__}
+
+    def __setstate__(self, state: typing.Dict[str, typing.Any]) -> None:
+        for name, value in state.items():
+            setattr(self, name, value)
+        self.stream = ByteStream(b"")
+
 
 class Response:
+    __attrs__ = [
+        "_content",
+        "_request",
+        "status_code",
+        "headers",
+        "next_request",
+        "extensions",
+        "history",
+        "_num_bytes_downloaded",
+    ]
+
     def __init__(
         self,
         status_code: int,
@@ -1151,6 +1177,16 @@ def num_bytes_downloaded(self) -> int:
     def __repr__(self) -> str:
         return f"<Response [{self.status_code} {self.reason_phrase}]>"
 
+    def __getstate__(self) -> typing.Dict[str, typing.Any]:
+        return {attr: getattr(self, attr, None) for attr in self.__attrs__}
+
+    def __setstate__(self, state: typing.Dict[str, typing.Any]) -> None:
+        for name, value in state.items():
+            setattr(self, name, value)
+        self.is_closed = True
+        self.is_stream_consumed = True
+        self.stream = ByteStream(b"")
+
     def read(self) -> bytes:
         """
         Read and return the response content.

tests/models/test_requests.py

@@ -1,3 +1,4 @@
+import pickle
 import typing
 
 import pytest
@@ -173,3 +174,34 @@ def test_url():
     assert request.url.port is None
     assert request.url.path == "/abc"
     assert request.url.raw_path == b"/abc?foo=bar"
+
+
+def test_request_picklable():
+    request = httpx.Request("POST", "http://example.org", json={"test": 123})
+    request.read()
+    pickle_request = pickle.loads(pickle.dumps(request))
+    assert pickle_request.method == "POST"
+    assert pickle_request.url.path == "/"
+    assert pickle_request.headers["Content-Type"] == "application/json"
+    assert pickle_request.content == b'{"test": 123}'
+    assert pickle_request.stream is not None
+    assert request.headers == {
+        "Host": "example.org",
+        "Content-Type": "application/json",
+        "content-length": "13",
+    }
+
+
+def test_request_generator_content_picklable():
+    def content():
+        yield b"test 123"  # pragma: nocover
+
+    request = httpx.Request("POST", "http://example.org", content=content())
+    pickle_request = pickle.loads(pickle.dumps(request))
+    assert pickle_request.stream is not None
+    assert pickle_request.content is None
+
+    request = httpx.Request("POST", "http://example.org", content=content())
+    request.read()
+    pickle_request = pickle.loads(pickle.dumps(request))
+    assert pickle_request.content == b"test 123"

tests/models/test_responses.py

@@ -1,4 +1,5 @@
 import json
+import pickle
 from unittest import mock
 
 import brotli
@@ -853,3 +854,36 @@ def content():
     headers = {"Content-Length": "8"}
     response = httpx.Response(200, content=content(), headers=headers)
     assert response.headers == {"Content-Length": "8"}
+
+
+def test_response_picklable():
+    response = httpx.Response(
+        200,
+        content=b"Hello, world!",
+        request=httpx.Request("GET", "https://example.org"),
+    )
+    pickle_response = pickle.loads(pickle.dumps(response))
+    assert pickle_response.is_closed is True
+    assert pickle_response.is_stream_consumed is True
+    assert pickle_response.next_request is None
+    assert pickle_response.stream is not None
+    assert pickle_response.content == b"Hello, world!"
+    assert pickle_response.status_code == 200
+    assert pickle_response.request.url == response.request.url
+    assert pickle_response.extensions == {}
+    assert pickle_response.history == []
+
+
+@pytest.mark.asyncio
+async def test_response_async_streaming_picklable():
+    response = httpx.Response(200, content=async_streaming_body())
+    pickle_response = pickle.loads(pickle.dumps(response))
+    assert pickle_response.content is None
+    assert pickle_response._num_bytes_downloaded == 0
+    assert pickle_response.headers == {"Transfer-Encoding": "chunked"}
+
+    response = httpx.Response(200, content=async_streaming_body())
+    await response.aread()
+    pickle_response = pickle.loads(pickle.dumps(response))
+    assert pickle_response.content == b"Hello, world!"
+    assert pickle_response._num_bytes_downloaded == 13