Support pickling ApiResponse (#100)

Author: pquentin

elastic_transport/_response.py

@@ -23,6 +23,7 @@
     Iterator,
     List,
     NoReturn,
+    Tuple,
     TypeVar,
     Union,
     overload,
@@ -97,6 +98,12 @@ def __getitem__(self, item: Any) -> Any:
     def __getattr__(self, attr: str) -> Any:
         return getattr(self._body, attr)
 
+    def __getstate__(self) -> Tuple[_BodyType, ApiResponseMeta]:
+        return self._body, self._meta
+
+    def __setstate__(self, state: Tuple[_BodyType, ApiResponseMeta]) -> None:
+        self._body, self._meta = state
+
     def __len__(self) -> int:
         return len(self._body)
 

tests/test_response.py

@@ -15,6 +15,8 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import pickle
+
 import pytest
 
 from elastic_transport import (
@@ -139,3 +141,19 @@ def test_constructor_allowed():
 
     resp = ObjectApiResponse(meta=meta, body={}, body_cls=int)
     assert resp == {}
+
+
+@pytest.mark.parametrize(
+    "response_cls, body",
+    [
+        (TextApiResponse, "Hello World"),
+        (BinaryApiResponse, b"Hello World"),
+        (ObjectApiResponse, {"Hello": "World"}),
+        (ListApiResponse, ["Hello", "World"]),
+    ],
+)
+def test_pickle(response_cls, body):
+    resp = response_cls(meta=meta, body=body)
+    pickled_resp = pickle.loads(pickle.dumps(resp))
+    assert pickled_resp == resp
+    assert pickled_resp.meta == resp.meta