Add PySymbol, a Symbol subclass that can be subclassed in Python

Author: isuruf

symengine/lib/symengine.pxd

@@ -130,6 +130,7 @@ cdef extern from "<symengine/symengine_rcp.h>" namespace "SymEngine":
         T& operator*() 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
     RCP[const Rational] rcp_static_cast_Rational "SymEngine::rcp_static_cast<const SymEngine::Rational>"(RCP[const Basic] &b) nogil
     RCP[const Complex] rcp_static_cast_Complex "SymEngine::rcp_static_cast<const SymEngine::Complex>"(RCP[const Basic] &b) nogil
@@ -257,6 +258,7 @@ cdef extern from "<symengine/basic.h>" namespace "SymEngine":
     bool is_a_Log "SymEngine::is_a<SymEngine::Log>"(const Basic &b) nogil
     bool is_a_PyNumber "SymEngine::is_a<SymEngine::PyNumber>"(const Basic &b) nogil
     bool is_a_ATan2 "SymEngine::is_a<SymEngine::ATan2>"(const Basic &b) nogil
+    bool is_a_PySymbol "SymEngine::is_a_sub<SymEngine::PySymbol>"(const Basic &b) nogil
 
     RCP[const Basic] expand(RCP[const Basic] &o) nogil except +
 
@@ -292,6 +294,11 @@ cdef extern from "<symengine/pywrapper.h>" namespace "SymEngine":
     cdef cppclass PyFunction:
         PyObject* get_py_object()
 
+cdef extern from "<symengine/pywrapper.h>" namespace "SymEngine":
+    cdef cppclass PySymbol(Symbol):
+        PySymbol(string name, PyObject* pyobj)
+        PyObject* get_py_object()
+
 cdef extern from "<symengine/integer.h>" namespace "SymEngine":
     cdef cppclass Integer(Number):
         Integer(int i) nogil
@@ -376,6 +383,7 @@ cdef extern from "<symengine/pow.h>" namespace "SymEngine":
 cdef extern from "<symengine/basic.h>" namespace "SymEngine":
     # We need to specialize these for our classes:
     RCP[const Basic] make_rcp_Symbol "SymEngine::make_rcp<const SymEngine::Symbol>"(string name) nogil
+    RCP[const Basic] make_rcp_PySymbol "SymEngine::make_rcp<const SymEngine::PySymbol>"(string name, PyObject * pyobj) nogil
     RCP[const Basic] make_rcp_Constant "SymEngine::make_rcp<const SymEngine::Constant>"(string name) nogil
     RCP[const Basic] make_rcp_Integer "SymEngine::make_rcp<const SymEngine::Integer>"(int i) nogil
     RCP[const Basic] make_rcp_Integer "SymEngine::make_rcp<const SymEngine::Integer>"(integer_class i) nogil

symengine/lib/symengine/pywrapper.h

@@ -8,6 +8,38 @@
 
 namespace SymEngine {
 
+/*
+ * PySymbol is a subclass of Symbol that keeps a reference to a Python object.
+ * When subclassing a Symbol from Python, the information stored in subclassed
+ * object is lost because all the arithmetic and function evaluations happen on
+ * the C++ side. The object returned by `(x + 1) - 1` is wrapped in the Python
+ * class Symbol and therefore the fact that `x` is a subclass of Symbol is lost.
+ *
+ * By subclassing in the C++ side and keeping a python object reference, the
+ * subclassed python object can be returned instead of wrapping in a Python
+ * class Symbol.
+ *
+ * TODO: Python object and C++ object both keep a reference to each other as one
+ * must be alive when the other is alive. This creates a cyclic reference and
+ * should be fixed.
+*/
+
+class PySymbol : public Symbol {
+private:
+    PyObject* obj;
+public:
+    PySymbol(const std::string& name, PyObject* obj) : Symbol(name), obj(obj) {
+        Py_INCREF(obj);
+    }
+    PyObject* get_py_object() const {
+        return obj;
+    }
+    virtual ~PySymbol() {
+        // TODO: This is never called because of the cyclic reference.
+        Py_DECREF(obj);
+    }
+};
+
 /*
  * This module provides classes to wrap Python objects defined in SymPy
  * or Sage into SymEngine.

symengine/lib/symengine_wrapper.pyx

@@ -34,6 +34,8 @@ cdef c2py(RCP[const symengine.Basic] o):
     elif (symengine.is_a_Complex(deref(o))):
         r = Complex.__new__(Complex)
     elif (symengine.is_a_Symbol(deref(o))):
+        if (symengine.is_a_PySymbol(deref(o))):
+            return <object>(deref(symengine.rcp_static_cast_PySymbol(o)).get_py_object())
         r = Symbol.__new__(Symbol)
     elif (symengine.is_a_Constant(deref(o))):
         r = Constant.__new__(Constant)
@@ -650,11 +652,20 @@ def series(ex, x=None, x0=0, n=6, method='sympy', removeO=False):
 
 cdef class Symbol(Basic):
 
+    """
+    Symbol is a class to store a symbolic variable with a given name.
+
+    Note: Subclassing `Symbol` will not work properly. Use `PySymbol`
+          which is a subclass of `Symbol` for subclassing.
+    """
     def __cinit__(self, name = None):
         if name is None:
             return
         self.thisptr = symengine.make_rcp_Symbol(name.encode("utf-8"))
 
+    def __init__(self, name = None):
+        return
+
     def _sympy_(self):
         cdef RCP[const symengine.Symbol] X = symengine.rcp_static_cast_Symbol(self.thisptr)
         import sympy
@@ -677,6 +688,15 @@ cdef class Symbol(Basic):
     def is_Symbol(self):
         return True
 
+
+cdef class PySymbol(Symbol):
+    def __init__(self, name, *args, **kwargs):
+        super(PySymbol, self).__init__(name)
+        if name is None:
+            return
+        self.thisptr = symengine.make_rcp_PySymbol(name.encode("utf-8"), <PyObject*>self)
+
+
 def symarray(prefix, shape, **kwargs):
     """ Creates an nd-array of symbols