support serializing and deserializing pysymbol

Author: isuruf

symengine/lib/pywrapper.cpp

@@ -1,4 +1,5 @@
 #include "pywrapper.h"
+#include <symengine/serialize-cereal.h>
 
 #if PY_MAJOR_VERSION >= 3
 #define PyInt_FromLong PyLong_FromLong
@@ -269,4 +270,80 @@ int PyFunction::compare(const Basic &o) const {
     return unified_compare(get_vec(), s.get_vec());
 }
 
+inline PyObject* get_pickle_module() {
+    static PyObject *module = NULL;
+    if (module == NULL) {
+        module = PyImport_ImportModule("pickle");
+    }
+    return module;
+}
+
+RCP<const Basic> load_basic(cereal::PortableBinaryInputArchive &ar, RCP<const Symbol> &)
+{
+    bool is_pysymbol;
+    std::string name;
+    ar(is_pysymbol);
+    ar(name);
+    if (is_pysymbol) {
+        std::string pickle_str;
+        ar(pickle_str);
+        PyObject *module = get_pickle_module();
+        PyObject *pickle_bytes = PyBytes_FromStringAndSize(pickle_str.data(), pickle_str.size());
+        PyObject *obj = PyObject_CallMethod(module, "loads", "O", pickle_bytes);
+        RCP<const Basic> result = make_rcp<PySymbol>(name, obj);
+        Py_XDECREF(pickle_bytes);
+        return result;
+    } else {
+        return symbol(name);
+    }
+}
+
+void save_basic(cereal::PortableBinaryOutputArchive &ar, const Symbol &b)
+{
+    bool is_pysymbol = is_a_sub<PySymbol>(b);
+    ar(is_pysymbol);
+    ar(b.__str__());
+    if (is_pysymbol) {
+        RCP<const PySymbol> p = rcp_static_cast<const PySymbol>(b.rcp_from_this());
+        PyObject *module = get_pickle_module();
+        PyObject *pickle_bytes = PyObject_CallMethod(module, "dumps", "O", p->get_py_object());
+        Py_ssize_t size;
+        char* buffer;
+        PyBytes_AsStringAndSize(pickle_bytes, &buffer, &size);
+        std::string pickle_str(buffer, size);
+        ar(pickle_str);
+        Py_XDECREF(pickle_bytes);
+    }
+}
+
+std::string wrapper_dumps(const Basic &x)
+{
+    std::cout << "qwe" << std::endl;
+    std::ostringstream oss;
+    unsigned short major = SYMENGINE_MAJOR_VERSION;
+    unsigned short minor = SYMENGINE_MINOR_VERSION;
+    cereal::PortableBinaryOutputArchive{oss}(major, minor,
+                                             x.rcp_from_this());
+    return oss.str();
+}
+
+RCP<const Basic> wrapper_loads(const std::string &serialized)
+{
+    unsigned short major, minor;
+    RCP<const Basic> obj;
+    std::istringstream iss(serialized);
+    cereal::PortableBinaryInputArchive iarchive{iss};
+    iarchive(major, minor);
+    if (major != SYMENGINE_MAJOR_VERSION or minor != SYMENGINE_MINOR_VERSION) {
+        throw SerializationError(StreamFmt()
+                                 << "SymEngine-" << SYMENGINE_MAJOR_VERSION
+                                 << "." << SYMENGINE_MINOR_VERSION
+                                 << " was asked to deserialize an object "
+                                 << "created using SymEngine-" << major << "."
+                                 << minor << ".");
+    }
+    iarchive(obj);
+    return obj;
+}
+
 } // SymEngine

symengine/lib/pywrapper.h

@@ -195,6 +195,9 @@ class PyFunction : public FunctionWrapper {
     virtual hash_t __hash__() const;
 };
 
+std::string wrapper_dumps(const Basic &x);
+RCP<const Basic> wrapper_loads(const std::string &s);
+
 }
 
 #endif //SYMENGINE_PYWRAPPER_H

symengine/lib/symengine.pxd

@@ -183,7 +183,6 @@ cdef extern from "<symengine/basic.h>" namespace "SymEngine":
         unsigned int hash() nogil except +
         vec_basic get_args() nogil
         int __cmp__(const Basic &o) nogil
-        string dumps() nogil except +
 
     ctypedef RCP[const Number] rcp_const_number "SymEngine::RCP<const SymEngine::Number>"
     ctypedef unordered_map[int, rcp_const_basic] umap_int_basic "SymEngine::umap_int_basic"
@@ -195,8 +194,6 @@ cdef extern from "<symengine/basic.h>" namespace "SymEngine":
     bool eq(const Basic &a, const Basic &b) nogil except +
     bool neq(const Basic &a, const Basic &b) nogil except +
 
-    RCP[const Basic] loads "SymEngine::Basic::loads"(const string &) nogil except +
-
     RCP[const Symbol] rcp_static_cast_Symbol "SymEngine::rcp_static_cast<const SymEngine::Symbol>"(rcp_const_basic &b) nogil
     RCP[const PySymbol] rcp_static_cast_PySymbol "SymEngine::rcp_static_cast<const SymEngine::PySymbol>"(rcp_const_basic &b) nogil
     RCP[const Integer] rcp_static_cast_Integer "SymEngine::rcp_static_cast<const SymEngine::Integer>"(rcp_const_basic &b) nogil
@@ -373,6 +370,9 @@ cdef extern from "pywrapper.h" namespace "SymEngine":
         PySymbol(string name, PyObject* pyobj)
         PyObject* get_py_object()
 
+    string wrapper_dumps(const Basic &x) nogil except +
+    rcp_const_basic wrapper_loads(const string &s) nogil except +
+
 cdef extern from "<symengine/integer.h>" namespace "SymEngine":
     cdef cppclass Integer(Number):
         Integer(int i) nogil

symengine/lib/symengine_wrapper.pyx

@@ -827,7 +827,7 @@ cdef list vec_pair_to_list(symengine.vec_pair& vec):
 
 
 def load_basic(bytes s):
-    return c2py(symengine.loads(s))
+    return c2py(symengine.wrapper_loads(s))
 
 
 repr_latex=[False]
@@ -841,7 +841,7 @@ cdef class Basic(object):
         return self.__str__()
 
     def __reduce__(self):
-        cdef bytes s = deref(self.thisptr).dumps()
+        cdef bytes s = symengine.wrapper_dumps(deref(self.thisptr))
         return (load_basic, (s,))
 
     def _repr_latex_(self):
@@ -1231,6 +1231,12 @@ cdef class Symbol(Expr):
         import sympy
         return sympy.Symbol(str(self))
 
+    def __reduce__(self):
+        if type(self) == Symbol:
+            return Basic.__reduce__(self)
+        else:
+            raise NotImplementedError("pickling for Symbol subclass not implemented")
+
     def _sage_(self):
         import sage.all as sage
         return sage.SR.symbol(str(self))

symengine/tests/test_pickling.py

@@ -1,4 +1,4 @@
-from symengine import symbols, sin, sinh, have_numpy, have_llvm, cos
+from symengine import symbols, sin, sinh, have_numpy, have_llvm, cos, Symbol
 import pickle
 import unittest
 
@@ -11,6 +11,23 @@ def test_basic():
     assert expr == expr2
 
 
+class MySymbol(Symbol):
+    def __init__(self, name, attr):
+        super().__init__(name=name)
+        self.attr = attr
+
+    def __reduce__(self):
+        return (self.__class__, (self.name, self.attr))
+
+
+def test_pysymbol():
+    a = MySymbol("hello", attr=1)
+    b = pickle.loads(pickle.dumps(a))
+    assert b.attr == 1
+    a._unsafe_reset()
+    b._unsafe_reset()
+
+
 @unittest.skipUnless(have_llvm, "No LLVM support")
 @unittest.skipUnless(have_numpy, "Numpy not installed")
 def test_llvm_double():
@@ -23,4 +40,3 @@ def test_llvm_double():
     ll = pickle.loads(ss)
     inp = [1, 2, 3]
     assert np.allclose(l(inp), ll(inp))
-