Skip to content

Commit

Permalink
Make Request and Response picklable
Browse files Browse the repository at this point in the history
  • Loading branch information
hannseman committed Apr 16, 2021
1 parent ed19995 commit e69822a
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 0 deletions.
36 changes: 36 additions & 0 deletions httpx/_models.py
Expand Up @@ -786,6 +786,13 @@ def __repr__(self) -> str:


class Request:
__attrs__ = [
"_content",
"method",
"url",
"headers",
]

def __init__(
self,
method: typing.Union[str, bytes],
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down
32 changes: 32 additions & 0 deletions tests/models/test_requests.py
@@ -1,3 +1,4 @@
import pickle
import typing

import pytest
Expand Down Expand Up @@ -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"
34 changes: 34 additions & 0 deletions tests/models/test_responses.py
@@ -1,4 +1,5 @@
import json
import pickle
from unittest import mock

import brotli
Expand Down Expand Up @@ -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

0 comments on commit e69822a

Please sign in to comment.