Allow to specify value returned by gets when not enough data is ava…

…ilable RESP3 has native boolean types, to be able to differentiate between insufficient data in buffer and actual False value add option to specify custom sentinel returned by gets when not enough data is available. CamelCase is used to conform to other arguments naming. Closes: #114
redis · Feb 15, 2022 · 0b1ac03 · 0b1ac03
1 parent dd15f01
commit 0b1ac03
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -47,9 +47,12 @@ Example:
 b'hello'
 ```
 
-When the buffer does not contain a full reply, `gets` returns `False`. This
-means extra data is needed and `feed` should be called again before calling
-`gets` again:
+When the buffer does not contain a full reply, `gets` returns `False`.
+This means extra data is needed and `feed` should be called again before calling
+`gets` again. Alternatively you could provide custom sentinel object via parameter,
+which is useful for RESP3 protocol where native boolean types are supported:
+
+Example:
 
 ```python
 >>> reader.feed("*2\r\n$5\r\nhello\r\n")
@@ -58,6 +61,9 @@ False
 >>> reader.feed("$5\r\nworld\r\n")
 >>> reader.gets()
 [b'hello', b'world']
+>>> reader = hiredis.Reader(notEnoughData=Ellipsis)
+>>> reader.gets()
+Ellipsis
 ```
 
 #### Unicode

diff --git a/hiredis/hiredis.pyi b/hiredis/hiredis.pyi
@@ -11,6 +11,7 @@ class Reader:
         replyError: Callable[[str], Exception] = ...,
         encoding: Optional[str] = ...,
         errors: Optional[str] = ...,
+        notEnoughData: Any = ...,
     ) -> None: ...
     def feed(
         self, __buf: Union[str, bytes], __off: int = ..., __len: int = ...

diff --git a/src/reader.c b/src/reader.c
@@ -214,6 +214,7 @@ static void Reader_dealloc(hiredis_ReaderObject *self) {
     redisReaderFree(self->reader);
     Py_XDECREF(self->protocolErrorClass);
     Py_XDECREF(self->replyErrorClass);
+    Py_XDECREF(self->notEnoughDataObject);
 
     ((PyObject *)self)->ob_type->tp_free((PyObject*)self);
 }
@@ -268,14 +269,15 @@ static int _Reader_set_encoding(hiredis_ReaderObject *self, char *encoding, char
 }
 
 static int Reader_init(hiredis_ReaderObject *self, PyObject *args, PyObject *kwds) {
-    static char *kwlist[] = { "protocolError", "replyError", "encoding", "errors", NULL };
+    static char *kwlist[] = { "protocolError", "replyError", "encoding", "errors", "notEnoughData", NULL };
     PyObject *protocolErrorClass = NULL;
     PyObject *replyErrorClass = NULL;
+    PyObject *notEnoughData = NULL;
     char *encoding = NULL;
     char *errors = NULL;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOzz", kwlist,
-        &protocolErrorClass, &replyErrorClass, &encoding, &errors))
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOzzO", kwlist,
+        &protocolErrorClass, &replyErrorClass, &encoding, &errors, &notEnoughData))
             return -1;
 
     if (protocolErrorClass)
@@ -286,6 +288,9 @@ static int Reader_init(hiredis_ReaderObject *self, PyObject *args, PyObject *kwd
         if (!_Reader_set_exception(&self->replyErrorClass, replyErrorClass))
             return -1;
 
+    if (notEnoughData)
+        self->notEnoughDataObject = notEnoughData;
+
     return _Reader_set_encoding(self, encoding, errors);
 }
 
@@ -299,11 +304,13 @@ static PyObject *Reader_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 
         self->encoding = NULL;
         self->errors = "strict";  // default to "strict" to mimic Python
+        self->notEnoughDataObject = Py_False;
         self->shouldDecode = 1;
         self->protocolErrorClass = HIREDIS_STATE->HiErr_ProtocolError;
         self->replyErrorClass = HIREDIS_STATE->HiErr_ReplyError;
         Py_INCREF(self->protocolErrorClass);
         Py_INCREF(self->replyErrorClass);
+        Py_INCREF(self->notEnoughDataObject);
 
         self->error.ptype = NULL;
         self->error.pvalue = NULL;
@@ -368,7 +375,7 @@ static PyObject *Reader_gets(hiredis_ReaderObject *self, PyObject *args) {
     }
 
     if (obj == NULL) {
-        Py_RETURN_FALSE;
+        return self->notEnoughDataObject;
     } else {
         /* Restore error when there is one. */
         if (self->error.ptype != NULL) {

diff --git a/src/reader.h b/src/reader.h
@@ -11,6 +11,7 @@ typedef struct {
     int shouldDecode;
     PyObject *protocolErrorClass;
     PyObject *replyErrorClass;
+    PyObject *notEnoughDataObject;
 
     /* 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

diff --git a/test/reader.py b/test/reader.py
@@ -314,3 +314,7 @@ def test_reader_has_data(self):
     self.assertEquals(True, self.reader.has_data())
     self.reply()
     self.assertEquals(False, self.reader.has_data())
+
+  def test_custom_not_enough_data(self):
+    self.reader = hiredis.Reader(notEnoughData=Ellipsis)
+    assert self.reader.gets() is Ellipsis