From 4b913ef9df287756e497cbbe61df03660af150b6 Mon Sep 17 00:00:00 2001 From: Sergey Prokazov Date: Mon, 30 Jan 2023 02:57:56 -0600 Subject: [PATCH] Add pack_command to support writing via hiredis-py (#147) Co-authored-by: Sergey Prokazov Co-authored-by: zalmane Co-authored-by: Chayim --- .gitignore | 1 + CHANGELOG.md | 2 + hiredis/__init__.py | 8 ++- hiredis/hiredis.pyi | 22 ++++++-- setup.py | 119 ++++++++++++++++++++++++++++---------------- src/hiredis.c | 24 ++++++++- src/pack.c | 106 +++++++++++++++++++++++++++++++++++++++ src/pack.h | 8 +++ tests/test_pack.py | 47 +++++++++++++++++ 9 files changed, 287 insertions(+), 50 deletions(-) mode change 100755 => 100644 setup.py create mode 100644 src/pack.c create mode 100644 src/pack.h create mode 100644 tests/test_pack.py diff --git a/.gitignore b/.gitignore index dbe349b..86cb2c8 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ MANIFEST .venv **/*.so +hiredis.egg-info diff --git a/CHANGELOG.md b/CHANGELOG.md index 96ed1e7..2e577c0 100644 --- 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 index f01b169..10afc02 100644 --- 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 index 61c0346..d75e88d 100644 --- 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 old mode 100755 new mode 100644 index 01677ce..c506c17 --- 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 index 9a0c88b..c96097d 100644 --- 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 new file mode 100644 index 0000000..0272a26 --- /dev/null +++ b/src/pack.c @@ -0,0 +1,106 @@ +#include "pack.h" +#include +#include + +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; +} \ No newline at end of file diff --git a/src/pack.h b/src/pack.h new file mode 100644 index 0000000..950193b --- /dev/null +++ b/src/pack.h @@ -0,0 +1,8 @@ +#ifndef __PACK_H +#define __PACK_H + +#include + +extern PyObject* pack_command(PyObject* cmd); + +#endif \ No newline at end of file diff --git a/tests/test_pack.py b/tests/test_pack.py new file mode 100644 index 0000000..2c3dbfd --- /dev/null +++ b/tests/test_pack.py @@ -0,0 +1,47 @@ +import hiredis +import pytest + +testdata = [ + # all_types + (('HSET', 'foo', 'key', 'value1', b'key_b', b'bytes str', 'key_mv', memoryview(b'bytes str'), b'key_i', 67, 'key_f', 3.14159265359), + b'*12\r\n$4\r\nHSET\r\n$3\r\nfoo\r\n$3\r\nkey\r\n$6\r\nvalue1\r\n$5\r\nkey_b\r\n$9\r\nbytes str\r\n$6\r\nkey_mv\r\n$9\r\nbytes str\r\n$5\r\nkey_i\r\n$2\r\n67\r\n$5\r\nkey_f\r\n$13\r\n3.14159265359\r\n'), + # spaces_in_cmd + (('COMMAND', 'GETKEYS', 'EVAL', 'return {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}', 2, 'key1', 'key2', 'first', 'second'), + b'*9\r\n$7\r\nCOMMAND\r\n$7\r\nGETKEYS\r\n$4\r\nEVAL\r\n$40\r\nreturn {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}\r\n$1\r\n2\r\n$4\r\nkey1\r\n$4\r\nkey2\r\n$5\r\nfirst\r\n$6\r\nsecond\r\n'), + # null_chars + (('SET', 'a', bytes(b'\xaa\x00\xffU')), + b'*3\r\n$3\r\nSET\r\n$1\r\na\r\n$4\r\n\xaa\x00\xffU\r\n'), + # encoding + (('SET', 'a', "йדב문諺"), + b'*3\r\n$3\r\nSET\r\n$1\r\na\r\n$12\r\n\xD0\xB9\xD7\x93\xD7\x91\xEB\xAC\xB8\xE8\xAB\xBA\r\n'), + # big_int + (('SET', 'a', 2**128), + b'*3\r\n$3\r\nSET\r\n$1\r\na\r\n$39\r\n340282366920938463463374607431768211456\r\n'), +] + +testdata_ids = [ + "all_types", + "spaces_in_cmd", + "null_chars", + "encoding", + "big_int", +] + + +@pytest.mark.parametrize("cmd,expected_packed_cmd", testdata, ids=testdata_ids) +def test_basic(cmd, expected_packed_cmd): + packed_cmd = hiredis.pack_command(cmd) + assert packed_cmd == expected_packed_cmd + + +def test_wrong_type(): + with pytest.raises(TypeError): + hiredis.pack_command(("HSET", "foo", True)) + + class Bar: + def __init__(self) -> None: + self.key = "key" + self.value = 36 + + with pytest.raises(TypeError): + hiredis.pack_command(("HSET", "foo", Bar()))