Implement rich compare

Author: atpalmer

src/cstring.c

@@ -26,14 +26,6 @@ static void cstring_dealloc(PyObject *self) {
     Py_TYPE(self)->tp_free(self);
 }
 
-static PyObject *cstring_str(PyObject *self) {
-    return PyUnicode_FromString(CSTRING_VALUE(self));
-}
-
-static Py_ssize_t cstring_len(PyObject *self) {
-    return Py_SIZE(self) - 1;
-}
-
 static PyTypeObject cstring_type;
 
 static int _ensure_cstring(PyObject *self) {
@@ -46,6 +38,42 @@ static int _ensure_cstring(PyObject *self) {
     return 0;
 }
 
+static PyObject *cstring_str(PyObject *self) {
+    return PyUnicode_FromString(CSTRING_VALUE(self));
+}
+
+static PyObject *cstring_richcompare(PyObject *self, PyObject *other, int op) {
+    if(!_ensure_cstring(other))
+        return NULL;
+
+    const char *left = CSTRING_VALUE(self);
+    const char *right = CSTRING_VALUE(other);
+
+    for(;*left && *right && *left == *right; ++left, ++right)
+        ;
+
+    switch (op) {
+    case Py_EQ:
+        return PyBool_FromLong(*left == *right);
+    case Py_NE:
+        return PyBool_FromLong(*left != *right);
+    case Py_LT:
+        return PyBool_FromLong(*left < *right);
+    case Py_GT:
+        return PyBool_FromLong(*left > *right);
+    case Py_LE:
+        return PyBool_FromLong(*left <= *right);
+    case Py_GE:
+        return PyBool_FromLong(*left >= *right);
+    default:
+        assert(0);
+    }
+}
+
+static Py_ssize_t cstring_len(PyObject *self) {
+    return Py_SIZE(self) - 1;
+}
+
 static PyObject *cstring_concat(PyObject *left, PyObject *right) {
     if(!_ensure_cstring(left))
         return NULL;
@@ -113,6 +141,7 @@ static PyTypeObject cstring_type = {
     .tp_flags = Py_TPFLAGS_DEFAULT,
     .tp_new = cstring_new,
     .tp_dealloc = cstring_dealloc,
+    .tp_richcompare = cstring_richcompare,
     .tp_str = cstring_str,
     .tp_as_sequence = &cstring_as_sequence,
 };

test/test_cstring.py

@@ -7,6 +7,44 @@ def test_str():
     assert str(result) == 'hello, world'
 
 
+# Rich Compare
+
+
+def test_eq():
+    a = cstring('hello')
+    b = cstring('hello')
+    assert a is not b
+    assert a == b
+
+
+def test_neq():
+    assert cstring('hello') != cstring('world')
+
+
+def test_ne():
+    assert cstring('hello') != cstring('world')
+
+
+def test_lt():
+    assert cstring('a') < cstring('b')
+    assert cstring('a') < cstring('aa')
+
+
+def test_le():
+    assert cstring('a') <= cstring('a')
+    assert cstring('a') <= cstring('b')
+
+
+def test_gt():
+    assert cstring('b') > cstring('a')
+    assert cstring('aa') > cstring('a')
+
+
+def test_ge():
+    assert cstring('a') >= cstring('a')
+    assert cstring('b') >= cstring('a')
+
+
 # Sequence API
 
 
@@ -64,3 +102,5 @@ def test_contains_True():
 def test_contains_False():
     assert cstring('hello') not in cstring('world')
 
+
+