Support for incr decr

memcachepool/cache.py

@@ -16,23 +16,14 @@
 DEFAULT_ITEM_SIZE = 1000 * 1000
 
 
-# XXX using python-memcached style pickling
-# but maybe we could use something else like
-# json
-#
-# at least this makes it compatible with
-# existing data
-def serialize(data):
-    return pickle.dumps(data)
-
-
-def unserialize(data):
-    return pickle.loads(data)
-
-
 # XXX not sure if keeping the base BaseMemcachedCache class has anymore value
 class UMemcacheCache(MemcachedCache):
     "An implementation of a cache binding using python-memcached"
+
+    _FLAG_SERIALIZED = 1
+    _FLAG_INT = 1<<1
+    _FLAG_LONG = 1<<2
+
     def __init__(self, server, params):
         import umemcache
         super(MemcachedCache, self).__init__(server, params,
@@ -47,6 +38,18 @@ def __init__(self, server, params):
         self._pool = ClientPool(self._get_client, maxsize=self.maxsize)
         self._blacklist = {}
 
+    # XXX using python-memcached style pickling
+    # but maybe we could use something else like
+    # json
+    #
+    # at least this makes it compatible with
+    # existing data
+    def serialize(self, data):
+        return pickle.dumps(data)
+
+    def unserialize(self, data):
+        return pickle.loads(data)
+
     def _get_memcache_timeout(self, timeout):
         if timeout == 0:
             return timeout
@@ -100,12 +103,30 @@ def create_client(server):
         else:
             raise socket.timeout('No server left in the pool')
 
+    def _flag_for_value(self, value):
+        if isinstance(value, int):
+            return self._FLAG_INT
+        elif isinstance(value, long):
+            return self._FLAG_LONG
+        return self._FLAG_SERIALIZED
+
+    def _value_for_flag(self, value, flag):
+        if flag == self._FLAG_INT:
+            return int(value)
+        elif flag == self._FLAG_LONG:
+            return long(value)
+        return self.unserialize(value)
+
     def add(self, key, value, timeout=0, version=None):
-        value = serialize(value)
+        flag = _flag_for_value(value)
+        if flag == self._FLAG_SERIALIED:
+            value = self.serialize(value)
+        else:
+            value = '%d' % value
         key = self.make_key(key, version=version)
 
         with self._pool.reserve() as conn:
-            return conn.add(key, value, self._get_memcache_timeout(timeout))
+            return conn.add(key, value, self._get_memcache_timeout(timeout), flag)
 
     def get(self, key, default=None, version=None):
         key = self.make_key(key, version=version)
@@ -115,13 +136,17 @@ def get(self, key, default=None, version=None):
         if val is None:
             return default
 
-        return unserialize(val[0])
+        return self._value_for_flag(value=val[0], flag=val[1])
 
     def set(self, key, value, timeout=0, version=None):
-        value = serialize(value)
+        flag = self._flag_for_value(value)
+        if flag == self._FLAG_SERIALIZED:
+            value = self.serialize(value)
+        else:
+            value = '%d' % value
         key = self.make_key(key, version=version)
         with self._pool.reserve() as conn:
-            conn.set(key, value, self._get_memcache_timeout(timeout))
+            conn.set(key, value, self._get_memcache_timeout(timeout), flag)
 
     def delete(self, key, version=None):
         key = self.make_key(key, version=version)
@@ -140,14 +165,14 @@ def get_many(self, keys, version=None):
                 res = conn.get(key)
                 if res is None:
                     continue
-                ret[key] = res[0]
+                ret[key] = res
 
         if ret:
             res = {}
             m = dict(zip(new_keys, keys))
 
             for k, v in ret.items():
-                res[m[k]] = unserialize(v)
+                res[m[k]] = self._value_for_flag(value=v[0], flag=v[1])
 
             return res
 
@@ -193,11 +218,16 @@ def set_many(self, data, timeout=0, version=None):
         safe_data = {}
         for key, value in data.items():
             key = self.make_key(key, version=version)
-            safe_data[key] = serialize(value)
+            flag = self._flag_for_value(value)
+            if flag == self._FLAG_SERIALIZED:
+                value = self.serialize(value)
+            else:
+                value = '%d' % value
+            safe_data[key] = value
 
         with self._pool.reserve() as conn:
             for key, value in safe_data.items():
-                conn.set(key, value, self._get_memcache_timeout(timeout))
+                conn.set(key, value, self._get_memcache_timeout(timeout), flag)
 
     def delete_many(self, keys, version=None):
         with self._pool.reserve() as conn:

memcachepool/tests/test_cache.py

@@ -38,6 +38,32 @@ def test_many(self):
         cache.delete_many(['a', 'b'])
         self.assertEqual(cache.get_many(['a', 'b']), {})
 
+    def test_incr_decr(self):
+        # make sure all the 'many' APIs work
+        from memcachepool.cache import UMemcacheCache
+
+        # creating the cache class
+        cache = UMemcacheCache('127.0.0.1:11211', {})
+        cache.set('a', 1)
+        cache.incr('a', 1)
+        self.assertEquals(cache.get('a'), 2)
+        cache.decr('a', 1)
+        self.assertEquals(cache.get('a'), 1)
+
+    def test_types(self):
+        # make sure all the 'many' APIs work
+        from memcachepool.cache import UMemcacheCache
+
+        # creating the cache class
+        cache = UMemcacheCache('127.0.0.1:11211', {})
+        cache.set('a', int(1))
+        self.assertEquals(cache.get('a'), 1)
+        self.assertTrue(isinstance(cache.get('a'), int))
+
+        cache.set('a', long(1))
+        self.assertEquals(cache.get('a'), 1)
+        self.assertTrue(isinstance(cache.get('a'), long))
+
     def test_loadbalance(self):
         from memcachepool.cache import UMemcacheCache