Merge pull request #3 from valkey-io/mkmkme/fix-resp3-attempt2

Introduce a new way of handling array objects

libvalkey/libvalkey.pyi

@@ -21,6 +21,7 @@ class Reader:
         encoding: Optional[str] = ...,
         errors: Optional[str] = ...,
         notEnoughData: Any = ...,
+        listOnly: bool = ...,
     ) -> None: ...
 
     def feed(

src/reader.c

@@ -14,6 +14,7 @@ static PyObject *Reader_getmaxbuf(libvalkey_ReaderObject *self);
 static PyObject *Reader_len(libvalkey_ReaderObject *self);
 static PyObject *Reader_has_data(libvalkey_ReaderObject *self);
 static PyObject *Reader_set_encoding(libvalkey_ReaderObject *self, PyObject *args, PyObject *kwds);
+static PyObject *Reader_listOnly(PyObject *self, void *closure);
 
 static PyMethodDef libvalkey_ReaderMethods[] = {
     {"feed", (PyCFunction)Reader_feed, METH_VARARGS, NULL },
@@ -26,6 +27,11 @@ static PyMethodDef libvalkey_ReaderMethods[] = {
     { NULL }  /* Sentinel */
 };
 
+static PyGetSetDef libvalkey_ReaderGetSet[] = {
+    {"listOnly", (getter)Reader_listOnly, NULL, NULL, NULL},
+    {NULL}  /* Sentinel */
+};
+
 PyTypeObject libvalkey_ReaderType = {
     PyVarObject_HEAD_INIT(NULL, 0)
     MOD_LIBVALKEY ".Reader",        /*tp_name*/
@@ -56,7 +62,7 @@ PyTypeObject libvalkey_ReaderType = {
     0,                            /*tp_iternext */
     libvalkey_ReaderMethods,        /*tp_methods */
     0,                            /*tp_members */
-    0,                            /*tp_getset */
+    libvalkey_ReaderGetSet,        /*tp_getset */
     0,                            /*tp_base */
     0,                            /*tp_dict */
     0,                            /*tp_descr_get */
@@ -67,30 +73,88 @@ PyTypeObject libvalkey_ReaderType = {
     Reader_new,                   /*tp_new */
 };
 
+static int tryParentize_impl(const valkeyReadTask *task, PyObject *obj,
+                             PyObject *parent, libvalkey_ReaderObject *self) {
+    switch (task->parent->type) {
+        case VALKEY_REPLY_MAP:
+            if (task->idx % 2 == 0) {
+                /* Save the object as a key. */
+                self->pendingObject = obj;
+            } else {
+                if (self->pendingObject == NULL) {
+                    Py_DECREF(obj);
+                    return -1;
+                }
+                if (PyDict_SetItem(parent, self->pendingObject, obj) < 0) {
+                    Py_DECREF(obj);
+                    Py_DECREF(self->pendingObject);
+                    self->pendingObject = NULL;
+                    return -1;
+                }
+                self->pendingObject = NULL;
+            }
+            break;
+        case VALKEY_REPLY_SET:
+            assert(PyAnySet_CheckExact(parent));
+            if (PySet_Add(parent, obj) < 0) {
+                Py_DECREF(obj);
+                return -1;
+            }
+            break;
+        default:
+            assert(PyList_CheckExact(parent));
+            if (PyList_SetItem(parent, task->idx, obj) < 0) {
+                Py_DECREF(obj);
+                return -1;
+            }
+    }
+    return 0;
+}
+
+static int tryParentize_ListOnly(const valkeyReadTask *task, PyObject *obj,
+                                 PyObject *parent, libvalkey_ReaderObject *self) {
+    switch (task->parent->type) {
+        case VALKEY_REPLY_MAP:
+            if (task->idx % 2 == 0) {
+                /* Set a temporary item to save the object as a key. */
+                self->pendingObject = PyTuple_New(2);
+                if (self->pendingObject == NULL) {
+                    Py_DECREF(obj);
+                    return -1;
+                }
+                PyTuple_SET_ITEM(self->pendingObject, 0, obj);
+            } else {
+                if (self->pendingObject == NULL) {
+                    Py_DECREF(obj);
+                    return -1;
+                }
+                PyTuple_SET_ITEM(self->pendingObject, 1, obj);
+                int res = PyList_Append(parent, self->pendingObject);
+                Py_DECREF(self->pendingObject);
+                self->pendingObject = NULL;
+                if (res < 0)
+                    return -1;
+            }
+            break;
+        default:
+            assert(PyList_CheckExact(parent));
+            if (PyList_SetItem(parent, task->idx, obj) < 0) {
+                Py_DECREF(obj);
+                return -1;
+            }
+    }
+    return 0;
+}
+
 static void *tryParentize(const valkeyReadTask *task, PyObject *obj) {
-    PyObject *parent;
+    libvalkey_ReaderObject *self = (libvalkey_ReaderObject*)task->privdata;
     if (task && task->parent) {
-        parent = (PyObject*)task->parent->obj;
-        switch (task->parent->type) {
-            case VALKEY_REPLY_MAP:
-                if (task->idx % 2 == 0) {
-                    /* Set a temporary item to save the object as a key. */
-                    PyDict_SetItem(parent, obj, Py_None);
-                } else {
-                    /* Pop the temporary item and set proper key and value. */
-                    PyObject *last_item = PyObject_CallMethod(parent, "popitem", NULL);
-                    PyObject *last_key = PyTuple_GetItem(last_item, 0);
-                    PyDict_SetItem(parent, last_key, obj);
-                }
-                break;
-            case VALKEY_REPLY_SET:
-                assert(PyAnySet_CheckExact(parent));
-                PySet_Add(parent, obj);
-                break;
-            default:
-                assert(PyList_CheckExact(parent));
-                PyList_SET_ITEM(parent, task->idx, obj);
-        }
+        PyObject *parent = (PyObject*)task->parent->obj;
+        int res = self->listOnly
+            ? tryParentize_ListOnly(task, obj, parent, self)
+            : tryParentize_impl(task, obj, parent, self);
+        if (res < 0)
+            return NULL;
     }
     return obj;
 }
@@ -158,16 +222,24 @@ static void *createStringObject(const valkeyReadTask *task, char *str, size_t le
 }
 
 static void *createArrayObject(const valkeyReadTask *task, size_t elements) {
+    libvalkey_ReaderObject *self = (libvalkey_ReaderObject*)task->privdata;
     PyObject *obj;
-    switch (task->type) {
-        case VALKEY_REPLY_MAP:
-            obj = PyDict_New();
-            break;
-        case VALKEY_REPLY_SET:
-            obj = PySet_New(NULL);
-            break;
-        default:
-            obj = PyList_New(elements);
+    if (self->listOnly) {
+        /* For map, don't preallocate listand use append later. */
+        if (task->type == VALKEY_REPLY_MAP)
+            elements = 0;
+        obj = PyList_New(elements);
+    } else {
+        switch (task->type) {
+            case VALKEY_REPLY_MAP:
+                obj = PyDict_New();
+                break;
+            case VALKEY_REPLY_SET:
+                obj = PySet_New(NULL);
+                break;
+            default:
+                obj = PyList_New(elements);
+        }
     }
     return tryParentize(task, obj);
 }
@@ -279,15 +351,24 @@ static int _Reader_set_encoding(libvalkey_ReaderObject *self, char *encoding, ch
 }
 
 static int Reader_init(libvalkey_ReaderObject *self, PyObject *args, PyObject *kwds) {
-    static char *kwlist[] = { "protocolError", "replyError", "encoding", "errors", "notEnoughData", NULL };
+    static char *kwlist[] = {
+        "protocolError",
+        "replyError",
+        "encoding",
+        "errors",
+        "notEnoughData",
+        "listOnly",
+        NULL,
+    };
     PyObject *protocolErrorClass = NULL;
     PyObject *replyErrorClass = NULL;
     PyObject *notEnoughData = NULL;
     char *encoding = NULL;
     char *errors = NULL;
+    int listOnly = 0;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOzzO", kwlist,
-        &protocolErrorClass, &replyErrorClass, &encoding, &errors, &notEnoughData))
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOzzOp", kwlist,
+        &protocolErrorClass, &replyErrorClass, &encoding, &errors, &notEnoughData, &listOnly))
             return -1;
 
     if (protocolErrorClass)
@@ -305,6 +386,8 @@ static int Reader_init(libvalkey_ReaderObject *self, PyObject *args, PyObject *k
         Py_INCREF(self->notEnoughDataObject);
     }
 
+    self->listOnly = listOnly;
+
     return _Reader_set_encoding(self, encoding, errors);
 }
 
@@ -322,6 +405,8 @@ static PyObject *Reader_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         self->shouldDecode = 1;
         self->protocolErrorClass = LIBVALKEY_STATE->VkErr_ProtocolError;
         self->replyErrorClass = LIBVALKEY_STATE->VkErr_ReplyError;
+        self->pendingObject = NULL;
+        self->listOnly = 0;
         Py_INCREF(self->protocolErrorClass);
         Py_INCREF(self->replyErrorClass);
         Py_INCREF(self->notEnoughDataObject);
@@ -367,18 +452,27 @@ static PyObject *Reader_feed(libvalkey_ReaderObject *self, PyObject *args) {
 
 static PyObject *Reader_gets(libvalkey_ReaderObject *self, PyObject *args) {
     PyObject *obj;
-    PyObject *err;
-    char *errstr;
 
     self->shouldDecode = 1;
     if (!PyArg_ParseTuple(args, "|i", &self->shouldDecode)) {
         return NULL;
     }
 
     if (valkeyReaderGetReply(self->reader, (void**)&obj) == VALKEY_ERR) {
-        errstr = valkeyReaderGetError(self->reader);
-        /* protocolErrorClass might be a callable. call it, then use it's type */
-        err = createError(self->protocolErrorClass, errstr, strlen(errstr));
+        PyObject *err = NULL;
+        char *errstr = valkeyReaderGetError(self->reader);
+        /* This is a hack to avoid
+         * "SystemError: class returned a result with an exception set".
+         * It is caused by the fact that one of callbacks already set an
+         * exception (i.e. failed PyDict_SetItem). Calling createError()
+         * after an exception is set will raise SystemError as per
+         * https://github.com/python/cpython/issues/67759.
+         * Hence, only call createError() if there is no exception set.
+         */
+        if (PyErr_Occurred() == NULL) {
+            /* protocolErrorClass might be a callable. call it, then use it's type */
+            err = createError(self->protocolErrorClass, errstr, strlen(errstr));
+        }
         if (err != NULL) {
             obj = PyObject_Type(err);
             PyErr_SetString(obj, errstr);
@@ -454,3 +548,10 @@ static PyObject *Reader_set_encoding(libvalkey_ReaderObject *self, PyObject *arg
     Py_RETURN_NONE;
 
 }
+
+static PyObject *Reader_listOnly(PyObject *obj, void *closure) {
+    libvalkey_ReaderObject *self = (libvalkey_ReaderObject*)obj;
+    PyObject *result = PyBool_FromLong(self->listOnly);
+    Py_INCREF(result);
+    return result;
+}

src/reader.h

@@ -13,6 +13,9 @@ typedef struct {
     PyObject *protocolErrorClass;
     PyObject *replyErrorClass;
     PyObject *notEnoughDataObject;
+    int listOnly;
+
+    PyObject *pendingObject;
 
     /* Stores error object in between incomplete calls to #gets, in order to
      * only set the error once a full reply has been read. Otherwise, the

tests/test_reader.py

@@ -1,9 +1,9 @@
 import libvalkey
 import pytest
 
-@pytest.fixture()
-def reader():
-  return libvalkey.Reader()
+@pytest.fixture(params=[True, False])
+def reader(request):
+  return libvalkey.Reader(listOnly=request.param)
 
 # def reply():
 #   return reader.gets()
@@ -135,12 +135,57 @@ def test_none(reader):
   assert reader.gets() is None
 
 def test_set(reader):
-  reader.feed(b"~3\r\n+tangerine\r\n_\r\n,10.5\r\n")
-  assert {b"tangerine", None, 10.5} == reader.gets()
+    reader.feed(b"~3\r\n+tangerine\r\n_\r\n,10.5\r\n")
+    expected = (
+      [b"tangerine", None, 10.5]
+      if reader.listOnly
+      else {b"tangerine", None, 10.5}
+    )
+    assert expected == reader.gets()
 
 def test_dict(reader):
   reader.feed(b"%2\r\n+radius\r\n,4.5\r\n+diameter\r\n:9\r\n")
-  assert {b"radius": 4.5, b"diameter": 9} == reader.gets()
+  expected = (
+      [(b"radius", 4.5), (b"diameter", 9)]
+      if reader.listOnly
+      else {b"radius": 4.5, b"diameter": 9}
+  )
+  assert expected == reader.gets()
+
+def test_set_with_nested_dict(reader):
+  reader.feed(b"~2\r\n+tangerine\r\n%1\r\n+a\r\n:1\r\n")
+  if reader.listOnly:
+    assert [b"tangerine", [(b"a", 1)]] == reader.gets()
+  else:
+    with pytest.raises(TypeError):
+      reader.gets()
+
+def test_dict_with_nested_set(reader):
+  reader.feed(b"%1\r\n+a\r\n~2\r\n:1\r\n:2\r\n")
+  expected = (
+      [(b"a", [1, 2])]
+      if reader.listOnly
+      else {b"a": {1, 2}}
+  )
+  assert expected == reader.gets()
+
+def test_map_inside_list(reader):
+  reader.feed(b"*1\r\n%1\r\n+a\r\n:1\r\n")
+  expected = (
+      [[(b"a", 1)]]
+      if reader.listOnly
+      else [{b"a": 1}]
+  )
+  assert expected == reader.gets()
+
+def test_map_inside_set(reader):
+  reader.feed(b"~1\r\n%1\r\n+a\r\n:1\r\n")
+  if reader.listOnly:
+    assert [[(b"a", 1)]] == reader.gets()
+  else:
+    # Map inside set is not allowed in Python
+    with pytest.raises(TypeError):
+      reader.gets()
 
 def test_vector(reader):
   reader.feed(b">4\r\n+pubsub\r\n+message\r\n+channel\r\n+message\r\n")