Add pack_command to support writing via hiredis-py (#147)

Co-authored-by: Sergey Prokazov <sergey.prokazov@redis.com> Co-authored-by: zalmane <oren.elias@gmail.com> Co-authored-by: Chayim <chayim@users.noreply.github.com>
redis · Jan 30, 2023 · 4b913ef · 4b913ef
1 parent 1b2e6fc
commit 4b913ef
Show file tree

Hide file tree

Showing 9 changed files with 287 additions and 50 deletions.
diff --git a/.gitignore b/.gitignore
@@ -4,3 +4,4 @@
 MANIFEST
 .venv
 **/*.so
+hiredis.egg-info
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,5 @@
+* Implement pack_command that serializes redis-py command to the RESP bytes object.
+
 ### 2.1.1 (2023-10-01)
 
 * Restores publishing of source distribution (#139)

diff --git a/hiredis/__init__.py b/hiredis/__init__.py
@@ -1,6 +1,10 @@
-from .hiredis import Reader, HiredisError, ProtocolError, ReplyError
+from .hiredis import Reader, HiredisError, pack_command, ProtocolError, ReplyError
 from .version import __version__
 
 __all__ = [
-  "Reader", "HiredisError", "ProtocolError", "ReplyError",
+  "Reader",
+  "HiredisError",
+  "pack_command",
+  "ProtocolError",
+  "ReplyError",
   "__version__"]
diff --git a/hiredis/hiredis.pyi b/hiredis/hiredis.pyi
@@ -1,8 +1,17 @@
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional, Union, Tuple
+
+
+class HiredisError(Exception):
+    ...
+
+
+class ProtocolError(HiredisError):
+    ...
+
+
+class ReplyError(HiredisError):
+    ...
 
-class HiredisError(Exception): ...
-class ProtocolError(HiredisError): ...
-class ReplyError(HiredisError): ...
 
 class Reader:
     def __init__(
@@ -13,6 +22,7 @@ class Reader:
         errors: Optional[str] = ...,
         notEnoughData: Any = ...,
     ) -> None: ...
+
     def feed(
         self, __buf: Union[str, bytes], __off: int = ..., __len: int = ...
     ) -> None: ...
@@ -21,6 +31,10 @@ class Reader:
     def getmaxbuf(self) -> int: ...
     def len(self) -> int: ...
     def has_data(self) -> bool: ...
+
     def set_encoding(
         self, encoding: Optional[str] = ..., errors: Optional[str] = ...
     ) -> None: ...
+
+
+def pack_command(cmd: Tuple[str | int | float | bytes | memoryview]): ...
diff --git a/setup.py b/setup.py
@@ -1,56 +1,89 @@
 #!/usr/bin/env python
 
 try:
-  from setuptools import setup, Extension
+    from setuptools import setup, Extension
 except ImportError:
-  from distutils.core import setup, Extension
-import sys, importlib, os, glob, io
+    from distutils.core import setup, Extension
+import importlib
+import glob
+import io
+import sys
+
 
 def version():
-  loader = importlib.machinery.SourceFileLoader("hiredis.version", "hiredis/version.py")
-  module = loader.load_module()
-  return module.__version__
+    loader = importlib.machinery.SourceFileLoader(
+        "hiredis.version", "hiredis/version.py")
+    module = loader.load_module()
+    return module.__version__
+
+
+def get_sources():
+    hiredis_sources = ("alloc", "async", "hiredis", "net", "read", "sds", "sockcompat")
+    return sorted(glob.glob("src/*.c") + ["vendor/hiredis/%s.c" % src for src in hiredis_sources])
+
+
+def get_linker_args():
+    if 'win32' in sys.platform or 'darwin' in sys.platform:
+        return []
+    else:
+        return ["-Wl,-Bsymbolic",]
+
+
+def get_compiler_args():
+    if 'win32' in sys.platform:
+        return []
+    else:
+        return ["-std=c99",]
+
+
+def get_libraries():
+    if 'win32' in sys.platform:
+        return ["ws2_32",]
+    else:
+        return []
+
 
 ext = Extension("hiredis.hiredis",
-  sources=sorted(glob.glob("src/*.c") +
-                 ["vendor/hiredis/%s.c" % src for src in ("alloc", "read", "sds")]),
-  extra_compile_args=["-std=c99"],
-  include_dirs=["vendor"])
+                sources=get_sources(),
+                extra_compile_args=get_compiler_args(),
+                extra_link_args=get_linker_args(),
+                libraries=get_libraries(),
+                include_dirs=["vendor"])
 
 setup(
-  name="hiredis",
-  version=version(),
-  description="Python wrapper for hiredis",
-  long_description=io.open('README.md', 'rt', encoding='utf-8').read(),
-  long_description_content_type='text/markdown',
-  url="https://github.com/redis/hiredis-py",
-  author="Jan-Erik Rediger, Pieter Noordhuis",
-  author_email="janerik@fnordig.de, pcnoordhuis@gmail.com",
-  keywords=["Redis"],
-  license="BSD",
-  packages=["hiredis"],
-  package_data={"hiredis": ["hiredis.pyi", "py.typed"]},
-  ext_modules=[ext],
-  python_requires=">=3.7",
-  project_urls={
+    name="hiredis",
+    version=version(),
+    description="Python wrapper for hiredis",
+    long_description=io.open('README.md', 'rt', encoding='utf-8').read(),
+    long_description_content_type='text/markdown',
+    url="https://github.com/redis/hiredis-py",
+    author="Jan-Erik Rediger, Pieter Noordhuis",
+    author_email="janerik@fnordig.de, pcnoordhuis@gmail.com",
+    keywords=["Redis"],
+    license="BSD",
+    packages=["hiredis"],
+    package_data={"hiredis": ["hiredis.pyi", "py.typed"]},
+    ext_modules=[ext],
+    python_requires=">=3.7",
+    project_urls={
         "Changes": "https://github.com/redis/hiredis-py/releases",
         "Issue tracker": "https://github.com/redis/hiredis-py/issues",
-  },
-  classifiers=[
-    'Development Status :: 5 - Production/Stable',
-    'Intended Audience :: Developers',
-    'License :: OSI Approved :: BSD License',
-    'Operating System :: MacOS',
-    'Operating System :: POSIX',
-    'Programming Language :: C',
-    'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3 :: Only',
-    'Programming Language :: Python :: 3.7',
-    'Programming Language :: Python :: 3.8',
-    'Programming Language :: Python :: 3.9',
-    'Programming Language :: Python :: 3.10',
-    'Programming Language :: Python :: 3.11',
-    'Programming Language :: Python :: Implementation :: CPython',
-    'Topic :: Software Development',
-  ],
+    },
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Developers',
+        'License :: OSI Approved :: BSD License',
+        'Operating System :: MacOS',
+        'Operating System :: POSIX',
+        'Programming Language :: C',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3 :: Only',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
+        'Programming Language :: Python :: Implementation :: CPython',
+        'Topic :: Software Development',
+    ],
 )
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,33 @@ static int hiredis_ModuleClear(PyObject *m) {
     return 0;
 }
 
+static PyObject*
+py_pack_command(PyObject* self, PyObject* cmd)
+{
+    return pack_command(cmd);
+}
+
+PyDoc_STRVAR(pack_command_doc, "Pack a series of arguments into the Redis protocol");
+
+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 methods[] = {
+    {"pack_command", (PyCFunction) py_pack_command, METH_O, pack_command_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,106 @@
+#include "pack.h"
+#include <hiredis/hiredis.h>
+#include <hiredis/sdsalloc.h>
+
+PyObject *
+pack_command(PyObject *cmd)
+{
+    assert(cmd);
+    PyObject *result = NULL;
+
+    if (cmd == NULL || !PyTuple_Check(cmd))
+    {
+        PyErr_SetString(PyExc_TypeError,
+                        "The argument must be a tuple of str, int, float or bytes.");
+        return NULL;
+    }
+
+    int tokens_number = PyTuple_Size(cmd);
+    sds *tokens = s_malloc(sizeof(sds) * tokens_number);
+    if (tokens == NULL)
+    {
+        return PyErr_NoMemory();
+    }
+
+    memset(tokens, 0, sizeof(sds) * tokens_number);
+
+    size_t *lengths = hi_malloc(sizeof(size_t) * tokens_number);
+    if (lengths == NULL)
+    {
+        sds_free(tokens);
+        return PyErr_NoMemory();
+    }
+
+    Py_ssize_t len = 0;
+
+    for (Py_ssize_t i = 0; i < PyTuple_Size(cmd); i++)
+    {
+        PyObject *item = PyTuple_GetItem(cmd, i);
+
+        if (PyBytes_Check(item))
+        {
+            char *bytes = NULL;
+            Py_buffer buffer;
+            PyObject_GetBuffer(item, &buffer, PyBUF_SIMPLE);
+            PyBytes_AsStringAndSize(item, &bytes, &len);
+            tokens[i] = sdsempty();
+            tokens[i] = sdscpylen(tokens[i], bytes, len);
+            lengths[i] = buffer.len;
+            PyBuffer_Release(&buffer);
+        }
+        else if (PyUnicode_Check(item))
+        {
+            const char *bytes = PyUnicode_AsUTF8AndSize(item, &len);
+            if (bytes == NULL)
+            {
+                // PyUnicode_AsUTF8AndSize sets an exception.
+                goto cleanup;
+            }
+
+            tokens[i] = sdsnewlen(bytes, len);
+            lengths[i] = len;
+        }
+        else if (PyMemoryView_Check(item))
+        {
+            Py_buffer *p_buf = PyMemoryView_GET_BUFFER(item);
+            tokens[i] = sdsnewlen(p_buf->buf, p_buf->len);
+            lengths[i] = p_buf->len;
+        }
+        else
+        {
+            if (PyLong_CheckExact(item) || PyFloat_Check(item))
+            {
+                PyObject *repr = PyObject_Repr(item);
+                const char *bytes = PyUnicode_AsUTF8AndSize(repr, &len);
+
+                tokens[i] = sdsnewlen(bytes, len);
+                lengths[i] = len;
+                Py_DECREF(repr);
+            }
+            else
+            {
+                PyErr_SetString(PyExc_TypeError,
+                                "A tuple item must be str, int, float or bytes.");
+                goto cleanup;
+            }
+        }
+    }
+
+    char *resp_bytes = NULL;
+
+    len = redisFormatCommandArgv(&resp_bytes, tokens_number, (const char **)tokens, lengths);
+
+    if (len == -1)
+    {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "Failed to serialize the command.");
+        goto cleanup;
+    }
+
+    result = PyBytes_FromStringAndSize(resp_bytes, len);
+    hi_free(resp_bytes);
+cleanup:
+    sdsfreesplitres(tokens, tokens_number);
+    hi_free(lengths);
+    return result;
+}
diff --git a/src/pack.h b/src/pack.h
@@ -0,0 +1,8 @@
+#ifndef __PACK_H
+#define __PACK_H
+
+#include <Python.h>
+
+extern PyObject* pack_command(PyObject* cmd);
+
+#endif
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,3 +4,4 @@ @@
     MANIFEST
     .venv
     **/*.so
+    hiredis.egg-info