Added GET argument to SET command (#1412)

Author: Jiekun

redis/client.py

@@ -555,6 +555,20 @@ def parse_module_result(response):
     return True
 
 
+def parse_set_result(response, **options):
+    """
+    Handle SET result since GET argument is available since Redis 6.2.
+    Parsing SET result into:
+    - BOOL
+    - String when GET argument is used
+    """
+    if options.get('get'):
+        # Redis will return a getCommand result.
+        # See `setGenericCommand` in t_string.c
+        return response
+    return response and str_if_bytes(response) == 'OK'
+
+
 class Redis:
     """
     Implementation of the Redis protocol.
@@ -683,7 +697,7 @@ class Redis:
         'SENTINEL SENTINELS': parse_sentinel_slaves_and_sentinels,
         'SENTINEL SET': bool_ok,
         'SENTINEL SLAVES': parse_sentinel_slaves_and_sentinels,
-        'SET': lambda r: r and str_if_bytes(r) == 'OK',
+        'SET': parse_set_result,
         'SLOWLOG GET': parse_slowlog_get,
         'SLOWLOG LEN': int,
         'SLOWLOG RESET': bool_ok,
@@ -1804,6 +1818,9 @@ def getset(self, name, value):
         """
         Sets the value at key ``name`` to ``value``
         and returns the old value at key ``name`` atomically.
+
+        As per Redis 6.2, GETSET is considered deprecated.
+        Please use SET with GET parameter in new code.
         """
         return self.execute_command('GETSET', name, value)
 
@@ -1964,7 +1981,7 @@ def restore(self, name, ttl, value, replace=False, absttl=False):
         return self.execute_command('RESTORE', *params)
 
     def set(self, name, value,
-            ex=None, px=None, nx=False, xx=False, keepttl=False):
+            ex=None, px=None, nx=False, xx=False, keepttl=False, get=False):
         """
         Set the value at key ``name`` to ``value``
 
@@ -1980,8 +1997,13 @@ def set(self, name, value,
 
         ``keepttl`` if True, retain the time to live associated with the key.
             (Available since Redis 6.0)
+
+        ``get`` if True, set the value at key ``name`` to ``value`` and return
+            the old value stored at key, or None when key did not exist.
+            (Available since Redis 6.2)
         """
         pieces = [name, value]
+        options = {}
         if ex is not None:
             pieces.append('EX')
             if isinstance(ex, datetime.timedelta):
@@ -2001,7 +2023,11 @@ def set(self, name, value,
         if keepttl:
             pieces.append('KEEPTTL')
 
-        return self.execute_command('SET', *pieces)
+        if get:
+            pieces.append('GET')
+            options["get"] = True
+
+        return self.execute_command('SET', *pieces, **options)
 
     def __setitem__(self, name, value):
         self.set(name, value)

tests/test_commands.py

@@ -1006,6 +1006,14 @@ def test_set_keepttl(self, r):
         assert r.get('a') == b'2'
         assert 0 < r.ttl('a') <= 10
 
+    @skip_if_server_version_lt('6.2.0')
+    def test_set_get(self, r):
+        assert r.set('a', 'True', get=True) is None
+        assert r.set('a', 'True', get=True) == b'True'
+        assert r.set('a', 'foo') is True
+        assert r.set('a', 'bar', get=True) == b'foo'
+        assert r.get('a') == b'bar'
+
     def test_setex(self, r):
         assert r.setex('a', 60, '1')
         assert r['a'] == b'1'