diff --git a/hiredis/__init__.py b/hiredis/__init__.py index 623ee6b..a0eff07 100644 --- a/hiredis/__init__.py +++ b/hiredis/__init__.py @@ -1,4 +1,4 @@ -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__ = [ @@ -6,5 +6,6 @@ "HiredisError", "pack_command", "ProtocolError", + "PushNotification", "ReplyError", "__version__"] diff --git a/hiredis/hiredis.pyi b/hiredis/hiredis.pyi index d75e88d..e7d852e 100644 --- a/hiredis/hiredis.pyi +++ b/hiredis/hiredis.pyi @@ -13,6 +13,10 @@ class ReplyError(HiredisError): ... +class PushNotification(list): + ... + + class Reader: def __init__( self, diff --git a/src/hiredis.c b/src/hiredis.c index c96097d..f858072 100644 --- a/src/hiredis.c +++ b/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; } diff --git a/src/reader.c b/src/reader.c index e0052f7..2e32617 100644 --- a/src/reader.c +++ b/src/reader.c @@ -1,6 +1,7 @@ #include "reader.h" #include +#include 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); diff --git a/src/reader.h b/src/reader.h index b09e7fa..c9359c1 100644 --- a/src/reader.h +++ b/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 diff --git a/tests/test_reader.py b/tests/test_reader.py index 3694084..5d6652e 100644 --- a/tests/test_reader.py +++ b/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"