allow protoError and replyError to be any type of callable

chayim · Apr 15, 2014 · fcb22ad · fcb22ad
1 parent b278d64
commit fcb22ad
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 16 deletions.
diff --git a/src/reader.c b/src/reader.c
@@ -97,16 +97,23 @@ static PyObject *createDecodedString(hiredis_ReaderObject *self, const char *str
     return obj;
 }
 
+static void *createError(PyObject *errorCallable, char *errstr, size_t len) {
+    PyObject *obj;
+
+    PyObject *args = Py_BuildValue("(s#)", errstr, len);
+    assert(args != NULL); /* TODO: properly handle OOM etc */
+    obj = PyObject_CallObject(errorCallable, args);
+    assert(obj != NULL);
+    Py_DECREF(args);
+    return obj;
+}
+
 static void *createStringObject(const redisReadTask *task, char *str, size_t len) {
     hiredis_ReaderObject *self = (hiredis_ReaderObject*)task->privdata;
     PyObject *obj;
 
     if (task->type == REDIS_REPLY_ERROR) {
-        PyObject *args = Py_BuildValue("(s#)", str, len);
-        assert(args != NULL); /* TODO: properly handle OOM etc */
-        obj = PyObject_CallObject(self->replyErrorClass, args);
-        assert(obj != NULL);
-        Py_DECREF(args);
+        obj = createError(self->replyErrorClass, str, len);
     } else {
         obj = createDecodedString(self, str, len);
     }
@@ -153,15 +160,11 @@ static void Reader_dealloc(hiredis_ReaderObject *self) {
 }
 
 static int _Reader_set_exception(PyObject **target, PyObject *value) {
-    int subclass;
-    subclass = PyObject_IsSubclass(value, PyExc_Exception);
-
-    if (subclass == -1) {
-        return 0;
-    }
+    int callable;
+    callable = PyCallable_Check(value);
 
-    if (subclass == 0) {
-        PyErr_SetString(PyExc_TypeError, "Expected subclass of Exception");
+    if (callable == 0) {
+        PyErr_SetString(PyExc_TypeError, "Expected a callable");
         return 0;
     }
 
@@ -247,11 +250,17 @@ static PyObject *Reader_feed(hiredis_ReaderObject *self, PyObject *args) {
 
 static PyObject *Reader_gets(hiredis_ReaderObject *self) {
     PyObject *obj;
-    char *err;
+    PyObject *err;
+    char *errstr;
 
     if (redisReplyReaderGetReply(self->reader, (void**)&obj) == REDIS_ERR) {
-        err = redisReplyReaderGetError(self->reader);
-        PyErr_SetString(self->protocolErrorClass, err);
+        errstr = redisReplyReaderGetError(self->reader);
+        /* protocolErrorClass might be a callable. call it, then use it's type */
+        err = createError(self->protocolErrorClass, errstr, strlen(errstr));
+        obj = PyObject_Type(err);
+        PyErr_SetString(obj, errstr);
+        Py_DECREF(obj);
+        Py_DECREF(err);
         return NULL;
     }
 

diff --git a/test/reader.py b/test/reader.py
@@ -24,6 +24,11 @@ def test_protocol_error_with_custom_class(self):
     self.reader.feed(b"x")
     self.assertRaises(RuntimeError, self.reply)
 
+  def test_protocol_error_with_custom_callable(self):
+    self.reader = hiredis.Reader(protocolError=lambda err: RuntimeError(err))
+    self.reader.feed(b"x")
+    self.assertRaises(RuntimeError, self.reply)
+
   def test_fail_with_wrong_protocol_error_class(self):
     self.assertRaises(TypeError, hiredis.Reader, protocolError="wrong")
 
@@ -42,6 +47,14 @@ def test_error_string_with_custom_class(self):
     self.assertEquals(RuntimeError, type(error))
     self.assertEquals(("error",), error.args)
 
+  def test_error_string_with_custom_callable(self):
+    self.reader = hiredis.Reader(replyError=lambda err: RuntimeError(err))
+    self.reader.feed(b"-error\r\n")
+    error = self.reply()
+
+    self.assertEquals(RuntimeError, type(error))
+    self.assertEquals(("error",), error.args)
+
   def test_fail_with_wrong_reply_error_class(self):
     self.assertRaises(TypeError, hiredis.Reader, replyError="wrong")