Update enum implemantation.

2 files changed, 87 insertions, 16 deletions

diff --git a/libshiboken/sbkenum.cpp b/libshiboken/sbkenum.cpp
index 059c41749..aa5f80bf7 100644
--- a/libshiboken/sbkenum.cpp
+++ b/libshiboken/sbkenum.cpp
@@ -100,11 +100,6 @@ static PyObject* enum_int(PyObject* v)
 #endif
 }
 
-static int enum_bool(PyObject* v)
-{
-    return (SBK_ENUM(v)->ob_value > 0);
-}
-
 static long getNumberValue(PyObject* v)
 {
     PyObject* number = PyNumber_Long(v);
@@ -113,19 +108,89 @@ static long getNumberValue(PyObject* v)
     return result;
 }
 
+
+static PyObject*
+enum_and(PyObject *self, PyObject *b)
+{
+    if (!PyNumber_Check(b)) {
+        Py_INCREF(Py_NotImplemented);
+        return Py_NotImplemented;
+    }
+
+    long valA = SBK_ENUM(self)->ob_value;
+    long valB = getNumberValue(b);
+    return PyLong_FromLong(valA & valB);
+}
+
+static PyObject*
+enum_or(PyObject *self, PyObject *b)
+{
+    if (!PyNumber_Check(b)) {
+        Py_INCREF(Py_NotImplemented);
+        return Py_NotImplemented;
+    }
+
+    long valA = SBK_ENUM(self)->ob_value;
+    long valB = getNumberValue(b);
+    return PyLong_FromLong(valA | valB);
+}
+
+static PyObject*
+enum_xor(PyObject *self, PyObject *b)
+{
+    if (!PyNumber_Check(b)) {
+        Py_INCREF(Py_NotImplemented);
+        return Py_NotImplemented;
+    }
+
+    long valA = SBK_ENUM(self)->ob_value;
+    long valB = getNumberValue(b);
+    return PyLong_FromLong(valA ^ valB);
+}
+
+
+static int enum_bool(PyObject* v)
+{
+    return (SBK_ENUM(v)->ob_value > 0);
+}
+
+static PyObject* enum_add(PyObject* self, PyObject* v)
+{
+    long valA = SBK_ENUM(self)->ob_value;
+    long valB = getNumberValue(v);
+    return PyLong_FromLong(valA + valB);
+}
+
+static PyObject* enum_subtract(PyObject* self, PyObject* v)
+{
+    long valA = SBK_ENUM(self)->ob_value;
+    long valB = getNumberValue(v);
+    return PyLong_FromLong(valA - valB);
+}
+
+static PyObject* enum_multiply(PyObject* self, PyObject* v)
+{
+    long valA = SBK_ENUM(self)->ob_value;
+    long valB = getNumberValue(v);
+    return PyLong_FromLong(valA * valB);
+}
+
+
 static PyObject *
 enum_richcompare(PyObject *self, PyObject *other, int op)
 {
-    printf("COMPARE\n");
+    printf("enum_richcompare\n");
     int result = 0;
     if (!PyNumber_Check(other)) {
-        PyErr_BadArgument();
-        return NULL;
+        Py_INCREF(Py_NotImplemented);
+        return Py_NotImplemented;
     }
 
     long valA = SBK_ENUM(self)->ob_value;
     long valB = getNumberValue(other);
 
+    printf("enum_richcompare2: %ld/%ld\n", valA, valB);
+
     if (self == other) {
         result = 1;
     } else  {
@@ -167,9 +232,9 @@ static PyGetSetDef SbkEnumGetSetList[] = {
 };
 
 static PyNumberMethods enum_as_number = {
-     /* nb_add */                   0,
-     /* nb_subtract */              0,
-     /* nb_multiply */              0,
+     /* nb_add */                   enum_add,
+     /* nb_subtract */              enum_subtract,
+     /* nb_multiply */              enum_multiply,
 #ifndef IS_PY3K
      /* nb_divide */                0,
 #endif
@@ -177,15 +242,15 @@ static PyNumberMethods enum_as_number = {
      /* nb_divmod */                0,
      /* nb_power */                 0,
      /* nb_negative */              0,
-     /* nb_positive */              0,
+     /* nb_positive */              enum_int,
      /* nb_absolute */              0,
      /* nb_bool */                  enum_bool,
      /* nb_invert */                0,
      /* nb_lshift */                0,
      /* nb_rshift */                0,
-     /* nb_and */                   0,
-     /* nb_xor */                   0,
-     /* nb_or */                    0,
+     /* nb_and */                   enum_and,
+     /* nb_xor */                   enum_xor,
+     /* nb_or */                    enum_or,
 #ifndef IS_PY3K
      /* nb_coerce */                0,
 #endif
@@ -216,7 +281,7 @@ static PyNumberMethods enum_as_number = {
      /* nb_inplace_floor_divide */  0,
      /* nb_inplace_true_divide */   0,
 
-     /* nb_index */                 0
+     /* nb_index */                 enum_int
 };
 
 PyTypeObject SbkEnumType_Type = {
@@ -288,6 +353,11 @@ private:
 
 namespace Enum {
 
+bool check(PyObject* pyObj)
+{
+    return Py_TYPE(pyObj->ob_type) == &SbkEnumType_Type;
+}
+
 PyObject* getEnumItemFromValue(PyTypeObject* enumType, long itemValue)
 {
     PyObject *key, *value;
diff --git a/libshiboken/sbkenum.h b/libshiboken/sbkenum.h
index 510ab6b87..d6456a918 100644
--- a/libshiboken/sbkenum.h
+++ b/libshiboken/sbkenum.h
@@ -44,6 +44,7 @@ inline bool isShibokenEnum(PyObject* pyObj)
 
 namespace Enum
 {
+    LIBSHIBOKEN_API bool check(PyObject* obj);
     /**
      *  Creates a new enum type (and its flags type, if any is given)
      *  and registers it to Python and adds it to \p module.