Add support for PEXPIREAT's options

Author: dogukanteber

redis/commands/core.py

@@ -1767,18 +1767,32 @@ def pexpire(self, name: KeyT, time: ExpiryT) -> ResponseT:
             time = int(time.total_seconds() * 1000)
         return self.execute_command("PEXPIRE", name, time)
 
-    def pexpireat(self, name: KeyT, when: AbsExpiryT) -> ResponseT:
-        """
-        Set an expire flag on key ``name``. ``when`` can be represented
-        as an integer representing unix time in milliseconds (unix time * 1000)
-        or a Python datetime object.
+    def pexpireat(
+            self,
+            name: KeyT,
+            when: AbsExpiryT,
+            option: str = None
+        ) -> ResponseT:
+        """
+        Set an expire flag on key ``name`` with given ``option``. ``when``
+        can be represented as an integer representing unix time in
+        milliseconds (unix time * 1000) or a Python datetime object.
+        
+        Valid options are:
+            NX -> Set expiry only when the key has no expiry
+            XX -> Set expiry only when the key has an existing expiry
+            GT -> Set expiry only when the new expiry is greater than current one
+            LT -> Set expiry only when the new expiry is less than current one
 
         For more information check https://redis.io/commands/pexpireat
         """
         if isinstance(when, datetime.datetime):
             ms = int(when.microsecond / 1000)
             when = int(time.mktime(when.timetuple())) * 1000 + ms
-        return self.execute_command("PEXPIREAT", name, when)
+        exp_option = list()
+        if option is not None:
+            exp_option.append(option)
+        return self.execute_command("PEXPIREAT", name, when, *exp_option)
 
     def psetex(
         self,

tests/test_commands.py

@@ -1232,6 +1232,41 @@ def test_pexpireat_unixtime(self, r):
         assert r.pexpireat("a", expire_at_seconds) is True
         assert 0 < r.pttl("a") <= 61000
 
+    @skip_if_server_version_lt("7.0.0")
+    def test_pexpireat_option_nx(self, r):
+        assert r.set("key", "val") is True
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=1)
+        assert r.pexpireat("key", expire_at, "NX") is True
+        assert r.pexpireat("key", expire_at, "NX") is False
+
+    @skip_if_server_version_lt("7.0.0")
+    def test_pexpireat_option_xx(self, r):
+        assert r.set("key", "val") is True
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=1)
+        assert r.pexpireat("key", expire_at, "XX") is False
+        assert r.pexpireat("key", expire_at) is True
+        assert r.pexpireat("key", expire_at, "XX") is True
+
+    @skip_if_server_version_lt("7.0.0")
+    def test_pexpireat_option_gt(self, r):
+        assert r.set("key", "val") is True
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=2)
+        assert r.pexpireat("key", expire_at) is True
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=1)
+        assert r.pexpireat("key", expire_at, "GT") is False
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=3)
+        assert r.pexpireat("key", expire_at, "GT") is True
+
+    @skip_if_server_version_lt("7.0.0")
+    def test_pexpireat_option_lt(self, r):
+        assert r.set("key", "val") is True
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=2)
+        assert r.pexpireat("key", expire_at) is True
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=3)
+        assert r.pexpireat("key", expire_at, "LT") is False
+        expire_at = redis_server_time(r) + datetime.timedelta(minutes=1)
+        assert r.pexpireat("key", expire_at, "LT") is True
+
     @skip_if_server_version_lt("2.6.0")
     def test_psetex(self, r):
         assert r.psetex("a", 1000, "value")