Adding new hash commands with expiration options - HGETDEL, HGETEX, HSETEX (#3570)

Author: petyaslavova

Diff for: redis/commands/core.py

@@ -3,6 +3,7 @@
 import datetime
 import hashlib
 import warnings
+from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -44,6 +45,10 @@
     TimeoutSecT,
     ZScoreBoundT,
 )
+from redis.utils import (
+    deprecated_function,
+    extract_expire_flags,
+)
 
 from .helpers import list_or_args
 
@@ -1837,10 +1842,10 @@ def getdel(self, name: KeyT) -> ResponseT:
     def getex(
         self,
         name: KeyT,
-        ex: Union[ExpiryT, None] = None,
-        px: Union[ExpiryT, None] = None,
-        exat: Union[AbsExpiryT, None] = None,
-        pxat: Union[AbsExpiryT, None] = None,
+        ex: Optional[ExpiryT] = None,
+        px: Optional[ExpiryT] = None,
+        exat: Optional[AbsExpiryT] = None,
+        pxat: Optional[AbsExpiryT] = None,
         persist: bool = False,
     ) -> ResponseT:
         """
@@ -1863,41 +1868,19 @@ def getex(
 
         For more information see https://redis.io/commands/getex
         """
-
         opset = {ex, px, exat, pxat}
         if len(opset) > 2 or len(opset) > 1 and persist:
             raise DataError(
                 "``ex``, ``px``, ``exat``, ``pxat``, "
                 "and ``persist`` are mutually exclusive."
             )
 
-        pieces: list[EncodableT] = []
-        # similar to set command
-        if ex is not None:
-            pieces.append("EX")
-            if isinstance(ex, datetime.timedelta):
-                ex = int(ex.total_seconds())
-            pieces.append(ex)
-        if px is not None:
-            pieces.append("PX")
-            if isinstance(px, datetime.timedelta):
-                px = int(px.total_seconds() * 1000)
-            pieces.append(px)
-        # similar to pexpireat command
-        if exat is not None:
-            pieces.append("EXAT")
-            if isinstance(exat, datetime.datetime):
-                exat = int(exat.timestamp())
-            pieces.append(exat)
-        if pxat is not None:
-            pieces.append("PXAT")
-            if isinstance(pxat, datetime.datetime):
-                pxat = int(pxat.timestamp() * 1000)
-            pieces.append(pxat)
+        exp_options: list[EncodableT] = extract_expire_flags(ex, px, exat, pxat)
+
         if persist:
-            pieces.append("PERSIST")
+            exp_options.append("PERSIST")
 
-        return self.execute_command("GETEX", name, *pieces)
+        return self.execute_command("GETEX", name, *exp_options)
 
     def __getitem__(self, name: KeyT):
         """
@@ -2255,14 +2238,14 @@ def set(
         self,
         name: KeyT,
         value: EncodableT,
-        ex: Union[ExpiryT, None] = None,
-        px: Union[ExpiryT, None] = None,
+        ex: Optional[ExpiryT] = None,
+        px: Optional[ExpiryT] = None,
         nx: bool = False,
         xx: bool = False,
         keepttl: bool = False,
         get: bool = False,
-        exat: Union[AbsExpiryT, None] = None,
-        pxat: Union[AbsExpiryT, None] = None,
+        exat: Optional[AbsExpiryT] = None,
+        pxat: Optional[AbsExpiryT] = None,
     ) -> ResponseT:
         """
         Set the value at key ``name`` to ``value``
@@ -2292,36 +2275,21 @@ def set(
 
         For more information see https://redis.io/commands/set
         """
+        opset = {ex, px, exat, pxat}
+        if len(opset) > 2 or len(opset) > 1 and keepttl:
+            raise DataError(
+                "``ex``, ``px``, ``exat``, ``pxat``, "
+                "and ``keepttl`` are mutually exclusive."
+            )
+
+        if nx and xx:
+            raise DataError("``nx`` and ``xx`` are mutually exclusive.")
+
         pieces: list[EncodableT] = [name, value]
         options = {}
-        if ex is not None:
-            pieces.append("EX")
-            if isinstance(ex, datetime.timedelta):
-                pieces.append(int(ex.total_seconds()))
-            elif isinstance(ex, int):
-                pieces.append(ex)
-            elif isinstance(ex, str) and ex.isdigit():
-                pieces.append(int(ex))
-            else:
-                raise DataError("ex must be datetime.timedelta or int")
-        if px is not None:
-            pieces.append("PX")
-            if isinstance(px, datetime.timedelta):
-                pieces.append(int(px.total_seconds() * 1000))
-            elif isinstance(px, int):
-                pieces.append(px)
-            else:
-                raise DataError("px must be datetime.timedelta or int")
-        if exat is not None:
-            pieces.append("EXAT")
-            if isinstance(exat, datetime.datetime):
-                exat = int(exat.timestamp())
-            pieces.append(exat)
-        if pxat is not None:
-            pieces.append("PXAT")
-            if isinstance(pxat, datetime.datetime):
-                pxat = int(pxat.timestamp() * 1000)
-            pieces.append(pxat)
+
+        pieces.extend(extract_expire_flags(ex, px, exat, pxat))
+
         if keepttl:
             pieces.append("KEEPTTL")
 
@@ -4940,6 +4908,16 @@ def pfmerge(self, dest: KeyT, *sources: KeyT) -> ResponseT:
 AsyncHyperlogCommands = HyperlogCommands
 
 
+class HashDataPersistOptions(Enum):
+    # set the value for each provided key to each
+    # provided value only if all do not already exist.
+    FNX = "FNX"
+
+    # set the value for each provided key to each
+    # provided value only if all already exist.
+    FXX = "FXX"
+
+
 class HashCommands(CommandsProtocol):
     """
     Redis commands for Hash data type.
@@ -4980,6 +4958,80 @@ def hgetall(self, name: str) -> Union[Awaitable[dict], dict]:
         """
         return self.execute_command("HGETALL", name, keys=[name])
 
+    def hgetdel(
+        self, name: str, *keys: str
+    ) -> Union[
+        Awaitable[Optional[List[Union[str, bytes]]]], Optional[List[Union[str, bytes]]]
+    ]:
+        """
+        Return the value of ``key`` within the hash ``name`` and
+        delete the field in the hash.
+        This command is similar to HGET, except for the fact that it also deletes
+        the key on success from the hash with the provided ```name```.
+
+        Available since Redis 8.0
+        For more information see https://redis.io/commands/hgetdel
+        """
+        if len(keys) == 0:
+            raise DataError("'hgetdel' should have at least one key provided")
+
+        return self.execute_command("HGETDEL", name, "FIELDS", len(keys), *keys)
+
+    def hgetex(
+        self,
+        name: KeyT,
+        *keys: str,
+        ex: Optional[ExpiryT] = None,
+        px: Optional[ExpiryT] = None,
+        exat: Optional[AbsExpiryT] = None,
+        pxat: Optional[AbsExpiryT] = None,
+        persist: bool = False,
+    ) -> Union[
+        Awaitable[Optional[List[Union[str, bytes]]]], Optional[List[Union[str, bytes]]]
+    ]:
+        """
+        Return the values of ``key`` and ``keys`` within the hash ``name``
+        and optionally set their expiration.
+
+        ``ex`` sets an expire flag on ``kyes`` for ``ex`` seconds.
+
+        ``px`` sets an expire flag on ``keys`` for ``px`` milliseconds.
+
+        ``exat`` sets an expire flag on ``keys`` for ``ex`` seconds,
+        specified in unix time.
+
+        ``pxat`` sets an expire flag on ``keys`` for ``ex`` milliseconds,
+        specified in unix time.
+
+        ``persist`` remove the time to live associated with the ``keys``.
+
+        Available since Redis 8.0
+        For more information see https://redis.io/commands/hgetex
+        """
+        if not keys:
+            raise DataError("'hgetex' should have at least one key provided")
+
+        opset = {ex, px, exat, pxat}
+        if len(opset) > 2 or len(opset) > 1 and persist:
+            raise DataError(
+                "``ex``, ``px``, ``exat``, ``pxat``, "
+                "and ``persist`` are mutually exclusive."
+            )
+
+        exp_options: list[EncodableT] = extract_expire_flags(ex, px, exat, pxat)
+
+        if persist:
+            exp_options.append("PERSIST")
+
+        return self.execute_command(
+            "HGETEX",
+            name,
+            *exp_options,
+            "FIELDS",
+            len(keys),
+            *keys,
+        )
+
     def hincrby(
         self, name: str, key: str, amount: int = 1
     ) -> Union[Awaitable[int], int]:
@@ -5034,8 +5086,10 @@ def hset(
 
         For more information see https://redis.io/commands/hset
         """
+
         if key is None and not mapping and not items:
             raise DataError("'hset' with no key value pairs")
+
         pieces = []
         if items:
             pieces.extend(items)
@@ -5047,6 +5101,89 @@ def hset(
 
         return self.execute_command("HSET", name, *pieces)
 
+    def hsetex(
+        self,
+        name: str,
+        key: Optional[str] = None,
+        value: Optional[str] = None,
+        mapping: Optional[dict] = None,
+        items: Optional[list] = None,
+        ex: Optional[ExpiryT] = None,
+        px: Optional[ExpiryT] = None,
+        exat: Optional[AbsExpiryT] = None,
+        pxat: Optional[AbsExpiryT] = None,
+        data_persist_option: Optional[HashDataPersistOptions] = None,
+        keepttl: bool = False,
+    ) -> Union[Awaitable[int], int]:
+        """
+        Set ``key`` to ``value`` within hash ``name``
+
+        ``mapping`` accepts a dict of key/value pairs that will be
+        added to hash ``name``.
+
+        ``items`` accepts a list of key/value pairs that will be
+        added to hash ``name``.
+
+        ``ex`` sets an expire flag on ``keys`` for ``ex`` seconds.
+
+        ``px`` sets an expire flag on ``keys`` for ``px`` milliseconds.
+
+        ``exat`` sets an expire flag on ``keys`` for ``ex`` seconds,
+            specified in unix time.
+
+        ``pxat`` sets an expire flag on ``keys`` for ``ex`` milliseconds,
+            specified in unix time.
+
+        ``data_persist_option`` can be set to ``FNX`` or ``FXX`` to control the
+            behavior of the command.
+            ``FNX`` will set the value for each provided key to each
+                provided value only if all do not already exist.
+            ``FXX`` will set the value for each provided key to each
+                provided value only if all already exist.
+
+        ``keepttl`` if True, retain the time to live associated with the keys.
+
+        Returns the number of fields that were added.
+
+        Available since Redis 8.0
+        For more information see https://redis.io/commands/hsetex
+        """
+        if key is None and not mapping and not items:
+            raise DataError("'hsetex' with no key value pairs")
+
+        if items and len(items) % 2 != 0:
+            raise DataError(
+                "'hsetex' with odd number of items. "
+                "'items' must contain a list of key/value pairs."
+            )
+
+        opset = {ex, px, exat, pxat}
+        if len(opset) > 2 or len(opset) > 1 and keepttl:
+            raise DataError(
+                "``ex``, ``px``, ``exat``, ``pxat``, "
+                "and ``keepttl`` are mutually exclusive."
+            )
+
+        exp_options: list[EncodableT] = extract_expire_flags(ex, px, exat, pxat)
+        if data_persist_option:
+            exp_options.append(data_persist_option.value)
+
+        if keepttl:
+            exp_options.append("KEEPTTL")
+
+        pieces = []
+        if items:
+            pieces.extend(items)
+        if key is not None:
+            pieces.extend((key, value))
+        if mapping:
+            for pair in mapping.items():
+                pieces.extend(pair)
+
+        return self.execute_command(
+            "HSETEX", name, *exp_options, "FIELDS", int(len(pieces) / 2), *pieces
+        )
+
     def hsetnx(self, name: str, key: str, value: str) -> Union[Awaitable[bool], bool]:
         """
         Set ``key`` to ``value`` within hash ``name`` if ``key`` does not
@@ -5056,19 +5193,18 @@ def hsetnx(self, name: str, key: str, value: str) -> Union[Awaitable[bool], bool
         """
         return self.execute_command("HSETNX", name, key, value)
 
+    @deprecated_function(
+        version="4.0.0",
+        reason="Use 'hset' instead.",
+        name="hmset",
+    )
     def hmset(self, name: str, mapping: dict) -> Union[Awaitable[str], str]:
         """
         Set key to value within hash ``name`` for each corresponding
         key and value from the ``mapping`` dict.
 
         For more information see https://redis.io/commands/hmset
         """
-        warnings.warn(
-            f"{self.__class__.__name__}.hmset() is deprecated. "
-            f"Use {self.__class__.__name__}.hset() instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
         if not mapping:
             raise DataError("'hmset' with 'mapping' of length 0")
         items = []