Implement zlib compression.

_pylibmcmodule.c

@@ -32,6 +32,12 @@
  */
 
 #include "_pylibmcmodule.h"
+#ifdef USE_ZLIB
+#  include <zlib.h>
+#  define ZLIB_BUFSZ (1 << 14)
+#  define _ZLIB_ERR(s, rc) \
+  PyErr_Format(PylibMCExc_MemcachedError, "zlib error %d in " s, rc);
+#endif
 
 
 /* {{{ _pylibmc.client implementation */
@@ -155,12 +161,160 @@ error:
     return -1;
 }
 
+/* {{{ Compression helpers */
+#ifdef USE_ZLIB
+static PyObject *_PylibMC_Deflate(PyObject *value) {
+    int rc;
+    char *out;
+    z_stream strm;
+    Py_ssize_t length, out_sz;
+
+    if (!PyString_Check(value)) {
+        return NULL;
+    }
+
+    /* Don't ask me about this one. Got it from zlibmodule.c in Python 2.6. */
+    length = PyString_GET_SIZE(value);
+    out_sz = length + length / 1000 + 12 + 1;
+
+    if ((out = PyMem_New(char, out_sz)) == NULL) {
+        return NULL;
+    }
+
+    strm.avail_in = length;
+    strm.avail_out = out_sz;
+    strm.next_in = (Byte *)PyString_AS_STRING(value);
+    strm.next_out = (Byte *)out;
+
+    strm.zalloc = (alloc_func)NULL;
+    strm.zfree = (free_func)Z_NULL;
+
+    /* TODO Expose compression level somehow. */
+    if ((rc = deflateInit((z_streamp)&strm, Z_BEST_SPEED)) != Z_OK) {
+        _ZLIB_ERR("deflateInit", rc);
+        goto error;
+    }
+
+    if ((rc = deflate((z_streamp)&strm, Z_FINISH)) != Z_STREAM_END) {
+        _ZLIB_ERR("deflate", rc);
+        goto error;
+    }
+
+    if ((rc = deflateEnd((z_streamp)&strm)) != Z_OK) {
+        _ZLIB_ERR("deflateEnd", rc);
+        goto error;
+    }
+
+    return PyString_FromStringAndSize(out, strm.total_out);
+error:
+    PyMem_Del(out);
+    return NULL;
+}
+
+static PyObject *_PylibMC_Inflate(char *value, size_t size) {
+    int rc;
+    char *out;
+    PyObject *out_obj;
+    Py_ssize_t rvalsz;
+    z_stream strm;
+
+    /* Output buffer */
+    rvalsz = ZLIB_BUFSZ;
+    out_obj = PyString_FromStringAndSize(NULL, rvalsz);
+    if (out_obj == NULL) {
+        return NULL;
+    }
+    out = PyString_AS_STRING(out_obj);
+
+    /* Set up zlib stream. */
+    strm.avail_in = size;
+    strm.avail_out = (uInt)rvalsz;
+    strm.next_in = (Byte *)value;
+    strm.next_out = (Byte *)out;
+
+    strm.zalloc = (alloc_func)NULL;
+    strm.zfree = (free_func)Z_NULL;
+
+    /* TODO Add controlling of windowBits with inflateInit2? */
+    if ((rc = inflateInit((z_streamp)&strm)) != Z_OK) {
+        _ZLIB_ERR("inflateInit", rc);
+        goto error;
+    }
+
+    do {
+        rc = inflate((z_streamp)&strm, Z_FINISH);
+        printf("(loop) rc = %d\n", rc);
+
+        switch (rc) {
+        case Z_STREAM_END:
+            printf("Z_STREAM_END\n");
+            break;
+        /* When a Z_BUF_ERROR occurs, we should be out of memory.
+         * This is also true for Z_OK, hence the fall-through. */
+        case Z_BUF_ERROR:
+            if (strm.avail_out) {
+                _ZLIB_ERR("inflate", rc);
+                inflateEnd(&strm);
+                goto error;
+            }
+        /* Fall-through */
+        case Z_OK:
+            printf("Z_OK\n");
+            if (_PyString_Resize(&out_obj, (Py_ssize_t)(rvalsz << 1)) < 0) {
+                inflateEnd(&strm);
+                goto error;
+            }
+            /* Wind forward */
+            out = PyString_AS_STRING(out_obj);
+            strm.next_out = (Byte *)(out + rvalsz);
+            strm.avail_out = rvalsz;
+            rvalsz = rvalsz << 1;
+            break;
+        default:
+            inflateEnd(&strm);
+            goto error;
+        }
+    } while (rc != Z_STREAM_END);
+
+    if ((rc = inflateEnd(&strm)) != Z_OK) {
+        _ZLIB_ERR("inflateEnd", rc);
+        goto error;
+    }
+
+    printf("size %d -> size %d\n", (int)size, (int)strm.total_out);
+
+    _PyString_Resize(&out_obj, strm.total_out);
+    return out_obj;
+error:
+    Py_DECREF(out_obj);
+    return NULL;
+}
+#endif
+
 static PyObject *_PylibMC_parse_memcached_value(char *value, size_t size,
         uint32_t flags) {
     PyObject *retval, *tmp;
 
+#if USE_ZLIB
+    PyObject *inflated = NULL;
+
+    /* Decompress value if necessary. */
+    if (flags & PYLIBMC_FLAG_ZLIB) {
+        inflated = _PylibMC_Inflate(value, size);
+        value = PyString_AS_STRING(inflated);
+        size = PyString_GET_SIZE(inflated);
+    }
+#else
+    if (flags & PYLIBMC_FLAG_ZLIB) {
+        PyErr_SetString(PylibMCExc_MemcachedError,
+            "value for key compressed, unable to inflate");
+        return NULL;
+    }
+#endif
+
     retval = NULL;
-    switch (flags) {
+
+    switch (flags & PYLIBMC_FLAG_TYPES) {
         case PYLIBMC_FLAG_PICKLE:
             retval = _PylibMC_Unpickle(value, size);
             break;
@@ -183,6 +337,12 @@ static PyObject *_PylibMC_parse_memcached_value(char *value, size_t size,
                     "unknown memcached key flags %u", flags);
     }
 
+#if USE_ZLIB
+    if (inflated != NULL) {
+        Py_DECREF(inflated);
+    }
+#endif
+
     return retval;
 }
 
@@ -229,18 +389,26 @@ static PyObject *_PylibMC_RunSetCommand(PylibMC_Client *self,
     PyObject *retval = NULL;
     PyObject *store_val = NULL;
     unsigned int time = 0;
+    unsigned int min_compress = 0;
     uint32_t store_flags = 0;
 
-    static char *kws[] = { "key", "val", "time", NULL };
+    static char *kws[] = { "key", "val", "time", "min_compress_len", NULL };
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#O|I", kws,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#O|II", kws,
-                &key, &key_sz, &val, &time)) {
+                &key, &key_sz, &val, &time, &min_compress)) {
         return NULL;
     }
     if (!_PylibMC_CheckKeyStringAndSize(key, key_sz)) {
         return NULL;
     }
 
+#ifndef USE_ZLIB
+    if (min_compress) {
+        PyErr_SetString(PyExc_TypeError, "min_compress_len without zlib");
+        return NULL;
+    }
+#endif
+
     /* Adapt val to a str. */
     if (PyString_Check(val)) {
         store_val = val;
@@ -270,6 +438,22 @@ static PyObject *_PylibMC_RunSetCommand(PylibMC_Client *self,
         return NULL;
     }
 
+#ifdef USE_ZLIB
+    if (min_compress && PyString_GET_SIZE(store_val) > min_compress) {
+        PyObject *deflated = _PylibMC_Deflate(val);
+
+        if (deflated != NULL) {
+            Py_DECREF(store_val);
+            store_val = deflated;
+            store_flags |= PYLIBMC_FLAG_ZLIB;
+        } else {
+            /* XXX What to do here, really? */
+            PyErr_Print();
+            PyErr_Clear();
+        }
+    }
+#endif
+
     rc = f(self->mc, key, key_sz,
            PyString_AS_STRING(store_val), PyString_GET_SIZE(store_val),
            time, store_flags);
@@ -1029,6 +1213,14 @@ Oh, and: plankton.\n");
 
     PyModule_AddStringConstant(module, "__version__", PYLIBMC_VERSION);
 
+#ifdef USE_ZLIB
+    Py_INCREF(Py_True);
+    PyModule_AddObject(module, "support_compression", Py_True);
+#else
+    Py_INCREF(Py_False);
+    PyModule_AddObject(module, "support_compression", Py_False);
+#endif
+
     PylibMCExc_MemcachedError = PyErr_NewException(
             "_pylibmc.MemcachedError", NULL, NULL);
     PyModule_AddObject(module, "MemcachedError",

_pylibmcmodule.h

@@ -53,12 +53,20 @@ typedef ssize_t Py_ssize_t;
 #define PYLIBMC_SERVER_UDP   (1 << 1)
 #define PYLIBMC_SERVER_UNIX  (1 << 2)
 
-/* Key flags from python-memcache. */
+/* {{{ Key flags from python-memcached
+ * Some flags (like the compression one, ZLIB) are combined with others.
+ */
 #define PYLIBMC_FLAG_NONE    0
 #define PYLIBMC_FLAG_PICKLE  (1 << 0)
 #define PYLIBMC_FLAG_INTEGER (1 << 1)
 #define PYLIBMC_FLAG_LONG    (1 << 2)
-#define PYLIBMC_FLAG_BOOL    (1 << 3)
+/* Note: this is an addition! python-memcached doesn't handle bools. */
+#define PYLIBMC_FLAG_BOOL    (1 << 4)
+#define PYLIBMC_FLAG_TYPES   (PYLIBMC_FLAG_PICKLE | PYLIBMC_FLAG_INTEGER | \
+                              PYLIBMC_FLAG_LONG | PYLIBMC_FLAG_BOOL)
+/* Modifier flags */
+#define PYLIBMC_FLAG_ZLIB    (1 << 3)
+/* }}} */
 
 #define PYLIBMC_INC  (1 << 0)
 #define PYLIBMC_DEC  (1 << 1)

pylibmc.py

@@ -26,6 +26,7 @@ import _pylibmc
 
 __all__ = ["hashers", "distributions", "Client"]
 __version__ = _pylibmc.__version__
+support_compression = _pylibmc.support_compression
 
 errors = tuple(e for (n, e) in _pylibmc.exceptions)
 # *Cough* Uhm, not the prettiest of things but this unpacks all exception

setup.py

@@ -1,6 +1,10 @@
 import os
 from distutils.core import setup, Extension
 
+use_zlib = True  # TODO Configurable somehow?
+
+libs = ["memcached"]
+defs = []
 incdirs = []
 libdirs = []
 
@@ -9,12 +13,16 @@ if "LIBMEMCACHED_DIR" in os.environ:
     incdirs.append(os.path.join(libdir, "include"))
     libdirs.append(os.path.join(libdir, "lib"))
 
-readme_text = open("README.rst", "U").read()
+if use_zlib:
-version = open("pylibmc-version.h", "U").read().strip().split("\"")[1]
+    libs.append("z")
+    defs.append(("USE_ZLIB", None))
  
 pylibmc_ext = Extension("_pylibmc", ["_pylibmcmodule.c"],
-                        libraries=["memcached"],
+                        libraries=libs, include_dirs=incdirs,
-                        include_dirs=incdirs, library_dirs=libdirs)
+                        library_dirs=libdirs, define_macros=defs)
+
+readme_text = open("README.rst", "U").read()
+version = open("pylibmc-version.h", "U").read().strip().split("\"")[1]
 
 setup(name="pylibmc", version=version,
       url="http://lericson.blogg.se/code/category/pylibmc.html",

tests.py

@@ -342,6 +342,7 @@ if __name__ == "__main__":
     print "Starting tests with _pylibmc at", _pylibmc.__file__
     print "Reported libmemcached version:", _pylibmc.libmemcached_version
     print "Reported pylibmc version:", _pylibmc.__version__
+    print "Support compression:", _pylibmc.support_compression
 
     if not is_alive(test_server):
         raise SystemExit("Test server (%r) not alive." % (test_server,))