From 1c83ad0de262b2a55c1b60c838caaf887a372a83 Mon Sep 17 00:00:00 2001
From: lericson <ludvig@lericson.se>
Date: Thu, 5 Nov 2009 19:09:09 +0100
Subject: [PATCH] Add specialized exception classes.

Useful when you need to catch certain errors like when a key isn't found
during an increment/decrement operation.
---
 _pylibmcmodule.c | 25 ++++++++++++++++++++++---
 _pylibmcmodule.h | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 tests.py         | 10 ++++++++++
 3 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/_pylibmcmodule.c b/_pylibmcmodule.c
index 674d97f..40093c7 100644
--- a/_pylibmcmodule.c
+++ b/_pylibmcmodule.c
@@ -890,8 +890,18 @@ static PyObject *PylibMC_ErrFromMemcached(PylibMC_Client *self, const char *what
         PyErr_Format(PyExc_RuntimeError, "error == 0? %s:%d",
                      __FILE__, __LINE__);
     } else { 
-        PyErr_Format(PylibMCExc_MemcachedError, "error %d from %s: %s",
-                error, what, memcached_strerror(self->mc, error));
+        PylibMC_McErr *err;
+        PyObject *exc = (PyObject *)PylibMCExc_MemcachedError;
+
+        for (err = PylibMCExc_mc_errs; err->name != NULL; err++) {
+            if (err->rc == error) {
+                exc = err->exc;
+                break;
+            }
+        }
+
+        PyErr_Format(exc, "error %d from %s: %s", error, what,
+                     memcached_strerror(self->mc, error));
     }
     return NULL;
 }
@@ -976,6 +986,7 @@ static PyMethodDef PylibMC_functions[] = {
 PyMODINIT_FUNC init_pylibmc(void) {
     PyObject *module;
     PylibMC_Behavior *b;
+    PylibMC_McErr *err;
     char name[128];
 
     if (PyType_Ready(&PylibMC_ClientType) < 0) {
@@ -1006,12 +1017,20 @@ Oh, and: plankton.\n");
         return;
     }
 
+    PyModule_AddStringConstant(module, "__version__", PYLIBMC_VERSION);
+
     PylibMCExc_MemcachedError = PyErr_NewException(
             "_pylibmc.MemcachedError", NULL, NULL);
     PyModule_AddObject(module, "MemcachedError",
                        (PyObject *)PylibMCExc_MemcachedError);
 
-    PyModule_AddStringConstant(module, "__version__", PYLIBMC_VERSION);
+    for (err = PylibMCExc_mc_errs; err->name != NULL; err++) {
+        char excnam[64];
+        strncpy(excnam, "_pylibmc.", 64);
+        strncat(excnam, err->name, 64);
+        err->exc = PyErr_NewException(excnam, PylibMCExc_MemcachedError, NULL);
+        PyModule_AddObject(module, err->name, (PyObject *)err->exc);
+    }
 
     Py_INCREF(&PylibMC_ClientType);
     PyModule_AddObject(module, "client", (PyObject *)&PylibMC_ClientType);
diff --git a/_pylibmcmodule.h b/_pylibmcmodule.h
index a0a39a3..ddf161d 100644
--- a/_pylibmcmodule.h
+++ b/_pylibmcmodule.h
@@ -68,6 +68,53 @@ typedef memcached_return (*_PylibMC_SetCommand)(memcached_st *, const char *,
 
 /* {{{ Exceptions */
 static PyObject *PylibMCExc_MemcachedError;
+
+/* Mapping of memcached_return value -> Python exception object. */
+typedef struct {
+    memcached_return rc;
+    char *name;
+    PyObject *exc;
+} PylibMC_McErr;
+
+static PylibMC_McErr PylibMCExc_mc_errs[] = {
+    { MEMCACHED_FAILURE, "Failure", NULL },
+    { MEMCACHED_HOST_LOOKUP_FAILURE, "HostLookupError", NULL },
+    { MEMCACHED_CONNECTION_FAILURE, "ConnectionError", NULL },
+    { MEMCACHED_CONNECTION_BIND_FAILURE, "ConnectionBindError", NULL },
+    { MEMCACHED_WRITE_FAILURE, "WriteError", NULL },
+    { MEMCACHED_READ_FAILURE, "ReadError", NULL },
+    { MEMCACHED_UNKNOWN_READ_FAILURE, "UnknownReadFailure", NULL },
+    { MEMCACHED_PROTOCOL_ERROR, "ProtocolError", NULL },
+    { MEMCACHED_CLIENT_ERROR, "ClientError", NULL },
+    { MEMCACHED_SERVER_ERROR, "ServerError", NULL },
+    { MEMCACHED_CONNECTION_SOCKET_CREATE_FAILURE, "SocketCreateError", NULL },
+    { MEMCACHED_DATA_EXISTS, "DataExists", NULL },
+    { MEMCACHED_DATA_DOES_NOT_EXIST, "DataDoesNotExist", NULL },
+    //{ MEMCACHED_NOTSTORED, "NotStored", NULL },
+    //{ MEMCACHED_STORED, "Stored", NULL },
+    { MEMCACHED_NOTFOUND, "NotFound", NULL },
+    { MEMCACHED_MEMORY_ALLOCATION_FAILURE, "AllocationError", NULL },
+    //{ MEMCACHED_PARTIAL_READ, "PartialRead", NULL },
+    { MEMCACHED_SOME_ERRORS, "SomeErrors", NULL },
+    { MEMCACHED_NO_SERVERS, "NoServers", NULL },
+    //{ MEMCACHED_END, "", NULL },
+    //{ MEMCACHED_DELETED, "", NULL },
+    //{ MEMCACHED_VALUE, "", NULL },
+    //{ MEMCACHED_STAT, "", NULL },
+    //{ MEMCACHED_ITEM, "", NULL },
+    //{ MEMCACHED_ERRNO, "", NULL },
+    { MEMCACHED_FAIL_UNIX_SOCKET, "UnixSocketError", NULL },
+    { MEMCACHED_NOT_SUPPORTED, "NotSupportedError", NULL },
+    { MEMCACHED_FETCH_NOTFINISHED, "FetchNotFinished", NULL },
+    //{ MEMCACHED_TIMEOUT, "TimeoutError", NULL },
+    //{ MEMCACHED_BUFFERED, "Buffer, NULL },
+    { MEMCACHED_BAD_KEY_PROVIDED, "BadKeyProvided", NULL },
+    { MEMCACHED_INVALID_HOST_PROTOCOL, "InvalidHostProtocolError", NULL },
+    //{ MEMCACHED_SERVER_MARKED_DEAD,
+    { MEMCACHED_UNKNOWN_STAT_KEY, "UnknownStatKey", NULL },
+    { MEMCACHED_E2BIG, "TooBigError", NULL },
+    { 0, NULL, NULL }
+};
 /* }}} */
 
 /* {{{ Behavior statics */
diff --git a/tests.py b/tests.py
index a545900..15d9124 100644
--- a/tests.py
+++ b/tests.py
@@ -147,6 +147,16 @@ True
 True
 >>> del c2
 
+Per-error exceptions
+>>> c.incr("test")
+Traceback (most recent call last):
+  ...
+NotFound: error 16 from memcached_increment: NOT FOUND
+>>> c.incr(chr(0))
+Traceback (most recent call last):
+  ...
+ProtocolError: error 8 from memcached_increment: PROTOCOL ERROR
+
 Behaviors.
 >>> c.set_behaviors({"tcp_nodelay": True, "hash": 6})
 >>> list(sorted((k, v) for (k, v) in c.get_behaviors().items()