From 213b58b9f28cdaa9f05e8de46c973debb9d7e232 Mon Sep 17 00:00:00 2001
From: Guillaume Pujol <guill.p.linux@gmail.com>
Date: Sat, 5 Aug 2023 21:14:03 +0200
Subject: [PATCH] add support for `pickle` (#5)

---
 README.md                  |  2 +-
 binapy/parsing/__init__.py |  1 +
 binapy/parsing/pickle.py   | 33 +++++++++++++++++++++++++++++++++
 tests/test_pickle.py       | 20 ++++++++++++++++++++
 4 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 binapy/parsing/pickle.py
 create mode 100644 tests/test_pickle.py

diff --git a/README.md b/README.md
index 153a134..2616925 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ isinstance(bp, bytes)
 ## Features
 
 - Fluent interface, based on a `bytes` subclass
-- Provides a convenient interface over `hashlib`, `base64`, `zlib`, `urllib.parse`, `json` and more
+- Provides a convenient interface over `hashlib`, `base64`, `zlib`, `urllib.parse`, `json`, `pickle` and more
 - Easy to extend with new formats
 
 ## TODO
diff --git a/binapy/parsing/__init__.py b/binapy/parsing/__init__.py
index 6979ecc..6554ac1 100644
--- a/binapy/parsing/__init__.py
+++ b/binapy/parsing/__init__.py
@@ -1,2 +1,3 @@
 """This module contains helpers for parsing or serializing data in various formats."""
 from .json import *
+from .pickle import *
diff --git a/binapy/parsing/pickle.py b/binapy/parsing/pickle.py
new file mode 100644
index 0000000..19eb9c5
--- /dev/null
+++ b/binapy/parsing/pickle.py
@@ -0,0 +1,33 @@
+"""This module implements helpers for (de)serializing to/from Python `pickle` objects."""
+
+import pickle
+from typing import Any, Optional
+
+from binapy import binapy_parser, binapy_serializer
+
+
+@binapy_serializer("pickle")
+def to_pickle(
+    o: object,
+    protocol: Optional[int] = None,
+    fix_imports: bool = True,
+    buffer_callback: Any = None,
+) -> bytes:
+    """Serialize an object using `pickle.dumps()`."""
+    return pickle.dumps(
+        o, protocol=protocol, fix_imports=fix_imports, buffer_callback=buffer_callback
+    )
+
+
+@binapy_parser("pickle")
+def from_pickle(
+    bp: bytes,
+    fix_imports: bool = True,
+    encoding: str = "ASCII",
+    errors: str = "strict",
+    buffers: Any = None,
+) -> object:
+    """Deserialize an object using `pickle.loads()`."""
+    return pickle.loads(
+        bp, fix_imports=fix_imports, encoding=encoding, errors=errors, buffers=buffers
+    )
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
new file mode 100644
index 0000000..6e7c843
--- /dev/null
+++ b/tests/test_pickle.py
@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import secrets
+
+from binapy import BinaPy
+
+
+class PickleTester:
+    def __init__(self) -> None:
+        self.random = secrets.token_hex()
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, PickleTester) and other.random == self.random
+
+
+o = PickleTester()
+
+
+def test_pickle() -> None:
+    assert BinaPy.serialize_to("pickle", o).parse_from("pickle") == o