handle errors gracefully

Author: isuruf

symengine/lib/pywrapper.cpp

@@ -275,6 +275,9 @@ inline PyObject* get_pickle_module() {
     if (module == NULL) {
         module = PyImport_ImportModule("pickle");
     }
+    if (module == NULL) {
+        throw SymEngineException("error importing pickle module.")
+    }
     return module;
 }
 
@@ -290,6 +293,9 @@ RCP<const Basic> load_basic(cereal::PortableBinaryInputArchive &ar, RCP<const Sy
         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);
+        if (obj == NULL) {
+            throw SymEngineException("error when loading pickled symbol subclass object");
+        }
         RCP<const Basic> result = make_rcp<PySymbol>(name, obj);
         Py_XDECREF(pickle_bytes);
         return result;
@@ -307,6 +313,9 @@ void save_basic(cereal::PortableBinaryOutputArchive &ar, const Symbol &b)
         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());
+        if (pickle_bytes == NULL) {
+            throw SymEngineException("error when pickling symbol subclass object");
+        }
         Py_ssize_t size;
         char* buffer;
         PyBytes_AsStringAndSize(pickle_bytes, &buffer, &size);

symengine/tests/test_pickling.py

@@ -1,4 +1,5 @@
 from symengine import symbols, sin, sinh, have_numpy, have_llvm, cos, Symbol
+from symengine.utilities import raises
 import pickle
 import unittest
 
@@ -11,22 +12,28 @@ def test_basic():
     assert expr == expr2
 
 
-class MySymbol(Symbol):
+class MySymbolBase(Symbol):
     def __init__(self, name, attr):
         super().__init__(name=name)
         self.attr = attr
 
+
+class MySymbol(MySymbolBase):
     def __reduce__(self):
         return (self.__class__, (self.name, self.attr))
 
 
 def test_pysymbol():
     a = MySymbol("hello", attr=1)
-    b = pickle.loads(pickle.dumps(a))
+    b = pickle.loads(pickle.dumps(a + 2)) - 2
     assert b.attr == 1
     a._unsafe_reset()
     b._unsafe_reset()
 
+    a = MySymbolBase("hello", attr=1)
+    raises(NotImplementedError, lambda: pickle.dumps(a + 2))
+    a._unsafe_reset()
+
 
 @unittest.skipUnless(have_llvm, "No LLVM support")
 @unittest.skipUnless(have_numpy, "Numpy not installed")