Implemented pack_bytes and pack_command.

redis · chayim · Feb 2, 2023 · Jan 12, 2023 · Jan 14, 2023 · Jan 16, 2023
commit 8e122b895fc03944a02fa73338c801117034bc18
diff --git a/hiredis/__init__.py b/hiredis/__init__.py
@@ -1,6 +1,11 @@
-from .hiredis import Reader, HiredisError, ProtocolError, ReplyError
+from .hiredis import Reader, HiredisError, pack_bytes, pack_command, ProtocolError, ReplyError
 from .version import __version__
 
 __all__ = [
-  "Reader", "HiredisError", "ProtocolError", "ReplyError",
+  "Reader",
+  "HiredisError",
+  "pack_bytes",
+  "pack_command",
+  "ProtocolError",
+  "ReplyError",
   "__version__"]
diff --git a/hiredis/hiredis.pyi b/hiredis/hiredis.pyi
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional, Union, Tuple
 
 class HiredisError(Exception): ...
 class ProtocolError(HiredisError): ...
@@ -24,3 +24,9 @@ class Reader:
     def set_encoding(
         self, encoding: Optional[str] = ..., errors: Optional[str] = ...
     ) -> None: ...
+
+
+def pack_bytes(bytes: bytes): ...
+
+
+def pack_command(cmd: Tuple[str|int|float|bytes]): ...
diff --git a/src/hiredis.c b/src/hiredis.c
@@ -1,5 +1,6 @@
 #include "hiredis.h"
 #include "reader.h"
+#include "pack.h"
 
 static int hiredis_ModuleTraverse(PyObject *m, visitproc visit, void *arg) {
     Py_VISIT(GET_STATE(m)->HiErr_Base);
@@ -15,12 +16,47 @@ static int hiredis_ModuleClear(PyObject *m) {
     return 0;
 }
 
+static PyObject*
+py_pack_command(PyObject* self, PyObject* cmd)
+{
+    return pack_command(cmd);
+}
+
+static PyObject*
+py_pack_bytes(PyObject* self, PyObject* cmd)
+{
+    return pack_bytes(cmd);
+}
+
+PyDoc_STRVAR(pack_command_doc, "Pack ...... ");
+PyDoc_STRVAR(pack_bytes_doc, "Pack ...... ");
+
+PyMethodDef pack_command_method = {
+    "pack_command",                 /* The name as a C string. */
+    (PyCFunction) py_pack_command,  /* The C function to invoke. */
+    METH_O,                         /* Flags telling Python how to invoke */
+    pack_command_doc,               /* The docstring as a C string. */
+};
+
+PyMethodDef pack_bytes_method = {
+    "pack_bytes",                 /* The name as a C string. */
+    (PyCFunction) py_pack_bytes,  /* The C function to invoke. */
+    METH_O,                         /* Flags telling Python how to invoke */
+    pack_bytes_doc,               /* The docstring as a C string. */
+};
+
+PyMethodDef methods[] = {
+    {"pack_command", (PyCFunction) py_pack_command, METH_O, pack_command_doc},
+    {"pack_bytes", (PyCFunction) py_pack_bytes, METH_O, pack_bytes_doc},
+    {NULL},
+};
+
 static struct PyModuleDef hiredis_ModuleDef = {
     PyModuleDef_HEAD_INIT,
     MOD_HIREDIS,
     NULL,
     sizeof(struct hiredis_ModuleState), /* m_size */
-    NULL, /* m_methods */
+    methods, /* m_methods */
     NULL, /* m_reload */
     hiredis_ModuleTraverse, /* m_traverse */
     hiredis_ModuleClear, /* m_clear */

diff --git a/src/pack.c b/src/pack.c
@@ -0,0 +1,111 @@
+#include "pack.h"
+#include <hiredis/hiredis.h>
+
+static 
+sds add_to_buffer(sds curr_buff, char *cmd, uint len)
+{
+    assert(curr_buff);
+    assert(sds);
+    sds newbuf = sdscatlen(curr_buff, cmd, len);
+    if (newbuf ==  NULL) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+
+    return newbuf;
+}
+
+PyObject*
+pack_command(PyObject* cmd)
+{
+    char *buf = sdsempty();
+    assert(cmd);
+    char *obj_buf =  NULL;
+    Py_ssize_t obj_len = 0;
+
+    if (cmd == NULL || !PyTuple_Check(cmd)) {
+        PyErr_SetString(PyExc_TypeError, 
+                        "The argument must be a tuple of str, int, float or bytes.");
+        return NULL;
+    }
+
+    if (buf == NULL) {
+        return PyErr_NoMemory();
+    }
+
+    for (Py_ssize_t i = 0; i < PyTuple_Size(cmd); i++) {
+        PyObject *item = PyTuple_GetItem(cmd, i);
+        sds temp_ptr =  NULL; 
+        if (PyBytes_Check(item)) {
+            // check result?
+            PyBytes_AsStringAndSize(item, &obj_buf, &obj_len);
+            temp_ptr = add_to_buffer(buf, obj_buf, obj_len);
+            if (temp_ptr == NULL) {
+                return NULL;
+            }
+
+            buf = temp_ptr;
+
+        } else if (PyUnicode_Check(item)) {
+
+            obj_buf = PyUnicode_AsUTF8AndSize(item, &obj_len);
+            if (obj_buf == NULL) {
+                sds_free(buf);
+                return NULL;
+            }
+
+            temp_ptr = add_to_buffer(buf, obj_buf, obj_len);
+            if (temp_ptr == NULL) {
+                return NULL;
+            }
+
+            buf = temp_ptr;
+        } else if PyLong_Check(item) {
+        } else {
+            printf("not yet\n");
+        }
+
+        temp_ptr = add_to_buffer(buf, " ", 1);
+        if (temp_ptr == NULL) {
+            return NULL;
+        }
+        buf = temp_ptr;
+    }
+
+    obj_len = redisFormatCommand(&obj_buf, buf);
+
+    PyObject *result = PyBytes_FromStringAndSize(obj_buf, obj_len);
+
+    hi_free(obj_buf);
+    //sds_free(buf);
+
+    return result;
+}
+
+
+PyObject*
+pack_bytes(PyObject* cmd)
+{
+    assert(cmd);
+    if (cmd == NULL || !PyBytes_Check(cmd)) {
+        PyErr_SetString(PyExc_TypeError, 
+                        "The argument must be a tuple of str, int, float or bytes.");
+        return NULL;
+    }
+
+    char *obj_buf =  NULL;
+    Py_ssize_t obj_len = 0;
+
+    if(PyBytes_AsStringAndSize(cmd, &obj_buf, &obj_len) == -1) {
+        return NULL;
+    }
+
+    char * str_result = NULL;
+    obj_len = redisFormatCommand(&str_result, obj_buf);
+
+    PyObject *result = PyBytes_FromStringAndSize(str_result, obj_len);
+
+    hi_free(str_result);
+
+    return result;
+}
diff --git a/src/pack.h b/src/pack.h
@@ -0,0 +1,9 @@
+#ifndef __PACK_H
+#define __PACK_H
+
+#include <Python.h>
+
+extern PyObject* pack_command(PyObject* cmd);
+extern PyObject* pack_bytes(PyObject* bytes);
+
+#endif