Introduce new type for RESP3 PUSH notifications (#208)

* Introduce new type for RESP3 PUSH notifications

Allow clients to distinguish between RESP3 arrays
and PUSH types by introducing PushNotification
type which subclasses list.

Fix #128

* Use simpler solution to preallocate PushNotification list

* Another attempt to make PushNotificationType compatible with PyPy

Author: uglide

hiredis/__init__.py

@@ -1,10 +1,11 @@
-from hiredis.hiredis import Reader, HiredisError, pack_command, ProtocolError, ReplyError
+from hiredis.hiredis import Reader, HiredisError, pack_command, ProtocolError, ReplyError, PushNotification
 from hiredis.version import __version__
 
 __all__ = [
   "Reader",
   "HiredisError",
   "pack_command",
   "ProtocolError",
+  "PushNotification",
   "ReplyError",
   "__version__"]

hiredis/hiredis.pyi

@@ -13,6 +13,10 @@ class ReplyError(HiredisError):
     ...
 
 
+class PushNotification(list):
+    ...
+
+
 class Reader:
     def __init__(
         self,

src/hiredis.c

@@ -59,6 +59,11 @@ PyMODINIT_FUNC PyInit_hiredis(void)
         return NULL;
     }
 
+    PushNotificationType.tp_base = &PyList_Type;
+    if (PyType_Ready(&PushNotificationType) < 0) {
+        return NULL;
+    }
+
     mod_hiredis = PyModule_Create(&hiredis_ModuleDef);
 
     /* Setup custom exceptions */
@@ -79,5 +84,8 @@ PyMODINIT_FUNC PyInit_hiredis(void)
     Py_INCREF(&hiredis_ReaderType);
     PyModule_AddObject(mod_hiredis, "Reader", (PyObject *)&hiredis_ReaderType);
 
+    Py_INCREF(&PushNotificationType);
+    PyModule_AddObject(mod_hiredis, "PushNotification", (PyObject *)&PushNotificationType);
+
     return mod_hiredis;
 }

src/reader.c

@@ -1,6 +1,7 @@
 #include "reader.h"
 
 #include <assert.h>
+#include <Python.h>
 
 static void Reader_dealloc(hiredis_ReaderObject *self);
 static int Reader_traverse(hiredis_ReaderObject *self, visitproc visit, void *arg);
@@ -14,6 +15,10 @@ static PyObject *Reader_len(hiredis_ReaderObject *self);
 static PyObject *Reader_has_data(hiredis_ReaderObject *self);
 static PyObject *Reader_set_encoding(hiredis_ReaderObject *self, PyObject *args, PyObject *kwds);
 
+static int PushNotificationType_init(PushNotificationObject *self, PyObject *args, PyObject *kwds);
+/* Create a new instance of PushNotificationType with preallocated number of elements */
+static PyObject* PushNotificationType_New(Py_ssize_t size);
+
 static PyMethodDef hiredis_ReaderMethods[] = {
     {"feed", (PyCFunction)Reader_feed, METH_VARARGS, NULL },
     {"gets", (PyCFunction)Reader_gets, METH_VARARGS, NULL },
@@ -66,6 +71,16 @@ PyTypeObject hiredis_ReaderType = {
     Reader_new,                   /*tp_new */
 };
 
+PyTypeObject PushNotificationType = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    .tp_name = MOD_HIREDIS ".PushNotification",
+    .tp_basicsize = sizeof(PushNotificationObject),
+    .tp_itemsize = 0,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+    .tp_doc = "Redis PUSH notification type",
+    .tp_init = (initproc) PushNotificationType_init,
+};
+
 static void *tryParentize(const redisReadTask *task, PyObject *obj) {
     PyObject *parent;
     if (task && task->parent) {
@@ -165,6 +180,9 @@ static void *createArrayObject(const redisReadTask *task, size_t elements) {
         case REDIS_REPLY_MAP:
             obj = PyDict_New();
             break;
+        case REDIS_REPLY_PUSH:
+            obj = PushNotificationType_New(elements);
+            break;
         default:
             obj = PyList_New(elements);
     }
@@ -199,6 +217,41 @@ static void freeObject(void *obj) {
     Py_XDECREF(obj);
 }
 
+static int PushNotificationType_init(PushNotificationObject *self, PyObject *args, PyObject *kwds) {
+    return PyList_Type.tp_init((PyObject *)self, args, kwds);
+}
+
+static PyObject* PushNotificationType_New(Py_ssize_t size) {
+    /* Check for negative size */
+    if (size < 0) {
+        PyErr_SetString(PyExc_SystemError, "negative list size");
+        return NULL;
+    }
+
+    /* Check for potential overflow */
+    if ((size_t)size > PY_SSIZE_T_MAX / sizeof(PyObject*)) {
+        return PyErr_NoMemory();
+    }
+
+#ifdef PYPY_VERSION
+    PyObject* obj = PyObject_CallObject((PyObject *) &PushNotificationType, NULL);    
+#else
+    PyObject* obj = PyType_GenericNew(&PushNotificationType, NULL, NULL);
+#endif
+    if (obj == NULL) {
+        return NULL;
+    }
+
+   int res = PyList_SetSlice(obj, PY_SSIZE_T_MAX, PY_SSIZE_T_MAX, PyList_New(size));
+
+   if (res == -1) {
+       Py_DECREF(obj);
+       return NULL;
+   }
+
+    return obj;
+}
+
 redisReplyObjectFunctions hiredis_ObjectFunctions = {
     createStringObject,  // void *(*createString)(const redisReadTask*, char*, size_t);
     createArrayObject,   // void *(*createArray)(const redisReadTask*, size_t);

src/reader.h

@@ -23,7 +23,12 @@ typedef struct {
     } error;
 } hiredis_ReaderObject;
 
+typedef struct {
+    PyListObject list;
+} PushNotificationObject;
+
 extern PyTypeObject hiredis_ReaderType;
+extern PyTypeObject PushNotificationType;
 extern redisReplyObjectFunctions hiredis_ObjectFunctions;
 
 #endif

tests/test_reader.py

@@ -185,9 +185,11 @@ def test_dict_with_unhashable_key(reader):
     with pytest.raises(TypeError):
       reader.gets()
 
-def test_vector(reader):
+def test_vector(reader):  
   reader.feed(b">4\r\n+pubsub\r\n+message\r\n+channel\r\n+message\r\n")
-  assert [b"pubsub", b"message", b"channel", b"message"] == reader.gets()
+  result = reader.gets()
+  assert isinstance(result, hiredis.PushNotification)
+  assert [b"pubsub", b"message", b"channel", b"message"] == result
 
 def test_verbatim_string(reader):
   value = b"text"