Implement extended set options added in Redis 2.6.12.

SPECS :
- http://redis.io/commands/set
- https://github.com/antirez/redis/issues/931

2 files changed, 71 insertions, 4 deletions

diff --git a/redis/client.py b/redis/client.py
index 941fb1d..0644214 100644
--- a/redis/client.py
+++ b/redis/client.py
@@ -213,7 +213,8 @@ class StrictRedis(object):
         string_keys_to_dict('ZSCORE ZINCRBY', float_or_none),
         string_keys_to_dict(
             'FLUSHALL FLUSHDB LSET LTRIM MSET RENAME '
-            'SAVE SELECT SET SHUTDOWN SLAVEOF WATCH UNWATCH',
+#            'SAVE SELECT SET SHUTDOWN SLAVEOF WATCH UNWATCH',
+            'SAVE SELECT SHUTDOWN SLAVEOF WATCH UNWATCH',
             lambda r: nativestr(r) == 'OK'
         ),
         string_keys_to_dict('BLPOP BRPOP', lambda r: r and tuple(r) or None),
@@ -242,6 +243,7 @@ class StrictRedis(object):
             'PING': lambda r: nativestr(r) == 'PONG',
             'RANDOMKEY': lambda r: r and r or None,
             'SCRIPT': parse_script,
+            'SET': lambda r: r and r == 'OK'  or None,
             'TIME': lambda x: (int(x[0]), int(x[1]))
         }
     )
@@ -707,9 +709,38 @@ class StrictRedis(object):
         "Rename key ``src`` to ``dst`` if ``dst`` doesn't already exist"
         return self.execute_command('RENAMENX', src, dst)
 
-    def set(self, name, value):
-        "Set the value at key ``name`` to ``value``"
-        return self.execute_command('SET', name, value)
+    def set(self, name, value, ex=None, px=None, nx=False, xx=False):
+        """
+        Set the value at key ``name`` to ``value``
+
+        ``ex`` sets an expire flag on key ``name`` for ``ex`` seconds.
+
+        ``px`` sets an expire flag on key ``name`` for ``px`` milliseconds.
+
+        ``nx`` if set to True, set the value at key ``name`` to ``value`` if it
+            does not already exist.
+
+        ``xx`` if set to True, set the value at key ``name`` to ``value`` if it
+            already exists.
+        """
+        pieces = [name, value]
+        if ex:
+            pieces.append('EX')
+            if isinstance(ex, datetime.timedelta):
+                ex = ex.seconds + ex.days * 24 * 3600
+            pieces.append(ex)
+        if px:
+            pieces.append('PX')
+            if isinstance(px, datetime.timedelta):
+                ms = int(px.microseconds / 1000)
+                px = (px.seconds + px.days * 24 * 3600) * 1000 + ms
+            pieces.append(px)
+
+        if nx:
+            pieces.append('NX')
+        if xx:
+            pieces.append('XX')
+        return self.execute_command('SET', *pieces)
     __setitem__ = set
 
     def setbit(self, name, offset, value):
diff --git a/tests/server_commands.py b/tests/server_commands.py
index 0225242..4a48ffb 100644
--- a/tests/server_commands.py
+++ b/tests/server_commands.py
@@ -456,6 +456,42 @@ class ServerCommandsTestCase(unittest.TestCase):
         self.assertEquals(self.client['a'], b('1'))
         self.assertEquals(self.client['b'], b('2'))
 
+    def test_set_nx(self):
+        self.assertEquals(self.client.set('foo', '1', nx=True), True)
+        self.assertEquals(self.client.set('foo', '2', nx=True), None)
+        self.assertEquals(self.client.get('foo'), b('1'))
+
+    def test_set_xx(self):
+        self.assertEquals(self.client.set('foo', '1', xx=True), None)
+        self.assertEquals(self.client.get('foo'), None)
+        self.client.set('foo', 'bar')
+        self.assertEquals(self.client.set('foo', '2', xx=True), True)
+        self.assertEquals(self.client.get('foo'), b('2'))
+
+    def test_set_px(self):
+        self.assertEquals(self.client.set('foo', '1', px=10000), True)
+        self.assertEquals(self.client['foo'], b('1'))
+        self.assert_(self.client.pttl('foo'), 10000)
+        self.assert_(self.client.ttl('foo'), 10)
+        # expire given a timeelta
+        expire_at = datetime.timedelta(milliseconds=1000)
+        self.assertEquals(self.client.set('foo', '1', px=expire_at), True)
+        self.assert_(self.client.pttl('foo'), 1000)
+        self.assert_(self.client.ttl('foo'), 1)
+
+    def test_set_ex(self):
+        self.assertEquals(self.client.set('foo', '1', ex=10), True)
+        self.assertEquals(self.client.ttl('foo'), 10)
+        # expire given a timeelta
+        expire_at = datetime.timedelta(seconds=60)
+        self.assertEquals(self.client.set('foo', '1', ex=expire_at), True)
+        self.assertEquals(self.client.ttl('foo'), 60)
+
+    def test_set_multipleoptions(self):
+        self.client['foo'] = 'val'
+        self.assertEquals(self.client.set('foo', 'bar', xx=True, px=10000), True)
+        self.assertEquals(self.client.ttl('foo'), 10)
+
     def test_setex(self):
         self.assertEquals(self.client.setex('a', '1', 60), True)
         self.assertEquals(self.client['a'], b('1'))