From a46d0335537b614a4fb3b1f09bb6310d94644f98 Mon Sep 17 00:00:00 2001
From: Camillo Lugaresi <camillo@openai.com>
Date: Wed, 22 May 2024 16:05:30 +0200
Subject: [PATCH] Support pickling HTTPStatusError

---
 httpx/_exceptions.py     | 17 +++++++++++++++++
 tests/test_exceptions.py | 17 +++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/httpx/_exceptions.py b/httpx/_exceptions.py
index 77f45a6d39..7d7f9c7df7 100644
--- a/httpx/_exceptions.py
+++ b/httpx/_exceptions.py
@@ -34,6 +34,7 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import typing
 
 if typing.TYPE_CHECKING:
@@ -267,6 +268,22 @@ def __init__(self, message: str, *, request: Request, response: Response) -> Non
         self.request = request
         self.response = response
 
+    def __reduce__(self) -> tuple[typing.Any, ...]:
+        # BaseException has a custom __reduce__ that can't handle subclasses with
+        # required keyword-only arguments.
+        return (
+            functools.partial(
+                self.__class__, request=self.request, response=self.response
+            ),
+            self.args,
+            # In case user code adds attributes to the exception instance.
+            {
+                k: v
+                for k, v in self.__dict__.items()
+                if k not in ("_request", "response")
+            },
+        )
+
 
 class InvalidURL(Exception):
     """
diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
index 60c8721c02..42c728d9f5 100644
--- a/tests/test_exceptions.py
+++ b/tests/test_exceptions.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import pickle
 import typing
 
 import httpcore
@@ -61,3 +62,19 @@ def test_request_attribute() -> None:
     request = httpx.Request("GET", "https://www.example.com")
     exc = httpx.ReadTimeout("Read operation timed out", request=request)
     assert exc.request == request
+
+
+def test_pickle_error(server):
+    with httpx.Client() as client:
+        response = client.request("GET", server.url.copy_with(path="/status/404"))
+        with pytest.raises(httpx.HTTPStatusError) as exc_info:
+            response.raise_for_status()
+        error = exc_info.value
+        assert isinstance(error, httpx.HTTPStatusError)
+        pickled_error = pickle.dumps(error)
+        unpickled_error = pickle.loads(pickled_error)
+        # Note that the unpickled error will not be equal to the original error
+        # because requests and responses are compared by identity.
+        assert str(unpickled_error) == str(error)
+        assert unpickled_error.request is not None
+        assert unpickled_error.response is not None