diff --git a/python-bindings/pythonbindings.cpp b/python-bindings/pythonbindings.cpp index 7fa977bed..cdc50cd54 100644 --- a/python-bindings/pythonbindings.cpp +++ b/python-bindings/pythonbindings.cpp @@ -373,6 +373,16 @@ PYBIND11_MODULE(blspy, m) // PythonGIL release_lock; return G1Element::FromBytes(Bytes(data_ptr, G1Element::SIZE)); }) + .def(py::pickle( + [](const G1Element &dp) { // __getstate__ + return py::make_tuple(dp.Serialize()); + }, + [](py::tuple t) { // __setstate__ + if (t.size() != 1) + throw std::runtime_error("Invalid state!"); + auto vecBytes = t[0].cast>(); + return G1Element::FromByteVector(vecBytes); + })) .def("generator", &G1Element::Generator) .def("from_message", py::overload_cast&, const uint8_t*, int>(&G1Element::FromMessage)) .def("pair", &G1Element::Pair) diff --git a/python-bindings/test.py b/python-bindings/test.py index 8fe58f0d1..2d93043c2 100644 --- a/python-bindings/test.py +++ b/python-bindings/test.py @@ -1,5 +1,6 @@ # flake8: noqa: E501 import binascii +import pickle from copy import deepcopy from blspy import ( @@ -335,11 +336,24 @@ def test_aggregate_verify_zero_items(): assert AugSchemeMPL.aggregate_verify([], [], G2Element()) +def test_pickle_support(): + seed = bytes([ + 0, 50, 6, 244, 24, 199, 1, 25, 52, 88, 192, 19, 18, 12, 89, 6, + 220, 18, 102, 58, 209, 82, 12, 62, 89, 110, 182, 9, 44, 20, 254, 22 + ]) + key: PrivateKey = AugSchemeMPL.key_gen(seed) + g1: G1Element = key.get_g1() + g1_data = pickle.dumps(g1) + g1_restored = pickle.loads(g1_data) + assert g1 == g1_restored + + test_schemes() test_vectors_invalid() test_vectors_valid() test_readme() test_aggregate_verify_zero_items() +test_pickle_support() print("\nAll tests passed.")