Python: add ZRANDMEMBER command (#276)

CHANGELOG.md

@@ -24,6 +24,7 @@
 * Python: Added BLPOP and BRPOP commands ([#1369](https://github.com/aws/glide-for-redis/pull/1369))
 * Python: Added ZRANGESTORE command ([#1377](https://github.com/aws/glide-for-redis/pull/1377))
 * Python: Added ZDIFFSTORE command ([#1378](https://github.com/aws/glide-for-redis/pull/1378))
+* Python: Added ZRANDMEMBER command (TODO: add PR link)
 
 
 #### Fixes

glide-core/src/client/value_conversion.rs

@@ -20,7 +20,8 @@ pub(crate) enum ExpectedReturnType {
     ArrayOfDoubleOrNull,
     Lolwut,
     ArrayOfArraysOfDoubleOrNull,
-    ArrayOfKeyValuePairs,
+    ArrayOfPairs,
+    ArrayOfMemberScorePairs,
 }
 
 pub(crate) fn convert_to_expected_type(
@@ -270,23 +271,25 @@ pub(crate) fn convert_to_expected_type(
                     .into()),
             }
         }
-        ExpectedReturnType::ArrayOfKeyValuePairs => match value {
-            Value::Nil => Ok(value),
-            Value::Array(ref array) if array.is_empty() || matches!(array[0], Value::Array(_)) => {
-                Ok(value)
-            }
-            Value::Array(array)
-                if matches!(array[0], Value::BulkString(_) | Value::SimpleString(_)) =>
-            {
-                convert_flat_array_to_key_value_pairs(array)
-            }
-            _ => Err((
-                ErrorKind::TypeError,
-                "Response couldn't be converted to an array of key-value pairs",
-                format!("(response was {:?})", value),
-            )
-                .into()),
-        },
+        // Used by HRANDFIELD when the WITHVALUES arg is passed.
+        // The server response can be a nil value, an empty array, a flat array of key-value pairs, or a two-dimensional array of key-value pairs.
+        // The conversions we do here are as follows:
+        //
+        // - if the server returned nil, return nil
+        // - if the server returned an empty array, return an empty array
+        // - otherwise, return a two-dimensional array of key-value pairs
+        ExpectedReturnType::ArrayOfPairs => convert_to_array_of_pairs(value, None),
+        // Used by ZRANDMEMBER when the WITHSCORES arg is passed.
+        // The server response can be a nil value, an empty array, a flat array of member-score pairs, or a two-dimensional array of member-score pairs.
+        // The server response scores can be strings or doubles. The conversions we do here are as follows:
+        //
+        // - if the server returned nil, return nil
+        // - if the server returned an empty array, return an empty array
+        // - otherwise, return a two-dimensional array of member-score pairs, where scores are of type double
+        ExpectedReturnType::ArrayOfMemberScorePairs => {
+            // RESP2 returns scores as strings, but we want scores as type double.
+            convert_to_array_of_pairs(value, Some(ExpectedReturnType::Double))
+        }
     }
 }
 
@@ -367,10 +370,53 @@ fn convert_array_to_map(
     Ok(Value::Map(map))
 }
 
-/// Converts a flat array of values to a two-dimensional array, where the inner arrays are two-length arrays representing key-value pairs. Normally a map would be more suitable for these responses, but some commands (eg HRANDFIELD) may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map.
+/// Used by commands like ZRANDMEMBER and HRANDFIELD. Normally a map would be more suitable for these key-value responses, but these commands may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map.
+///
+/// Converts a server response as follows:
+/// - if the server returned nil, return nil.
+/// - if the server returned an empty array, return an empty array.
+/// - if the server returned a flat array (RESP2), convert it to a two-dimensional array, where the inner arrays are length=2 arrays representing key-value pairs.
+/// - if the server returned a two-dimensional array (RESP3), return the response as is, since it is already in the correct format.
+/// - otherwise, return an error.
+///
+/// `response` is a server response that we should attempt to convert as described above.
+/// `value_expected_return_type` indicates the desired return type of the values in the key-value pairs. The values will only be converted if the response was a flat array, since RESP3 already returns an array of pairs with values already of the correct type.
+fn convert_to_array_of_pairs(
+    response: Value,
+    value_expected_return_type: Option<ExpectedReturnType>,
+) -> RedisResult<Value> {
+    match response {
+        Value::Nil => Ok(response),
+        Value::Array(ref array) if array.is_empty() || matches!(array[0], Value::Array(_)) => {
+            // The server response is an empty array or a RESP3 array of pairs. In RESP3, the values in the pairs are
+            // already of the correct type, so we do not need to convert them and `response` is in the correct format.
+            Ok(response)
+        }
+        Value::Array(array)
+            if array.len() % 2 == 0
+                && matches!(array[0], Value::BulkString(_) | Value::SimpleString(_)) =>
+        {
+            // The server response is a RESP2 flat array with keys at even indices and their associated values at
+            // odd indices.
+            convert_flat_array_to_array_of_pairs(array, value_expected_return_type)
+        }
+        _ => Err((
+            ErrorKind::TypeError,
+            "Response couldn't be converted to an array of member-score pairs",
+            format!("(response was {:?})", response),
+        )
+            .into()),
+    }
+}
+
+/// Converts a flat array of values to a two-dimensional array, where the inner arrays are length=2 arrays representing key-value pairs. Normally a map would be more suitable for these responses, but some commands (eg HRANDFIELD) may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map.
 ///
 /// `array` is a flat array containing keys at even-positioned elements and their associated values at odd-positioned elements.
-fn convert_flat_array_to_key_value_pairs(array: Vec<Value>) -> RedisResult<Value> {
+/// `value_expected_return_type` indicates the desired return type of the values in the key-value pairs.
+fn convert_flat_array_to_array_of_pairs(
+    array: Vec<Value>,
+    value_expected_return_type: Option<ExpectedReturnType>,
+) -> RedisResult<Value> {
     if array.len() % 2 != 0 {
         return Err((
             ErrorKind::TypeError,
@@ -381,7 +427,9 @@ fn convert_flat_array_to_key_value_pairs(array: Vec<Value>) -> RedisResult<Value
 
     let mut result = Vec::with_capacity(array.len() / 2);
     for i in (0..array.len()).step_by(2) {
-        let pair = vec![array[i].clone(), array[i + 1].clone()];
+        let key = array[i].clone();
+        let value = convert_to_expected_type(array[i + 1].clone(), value_expected_return_type)?;
+        let pair = vec![key, value];
         result.push(Value::Array(pair));
     }
     Ok(Value::Array(result))
@@ -407,10 +455,10 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
         b"GEOPOS" => Some(ExpectedReturnType::ArrayOfArraysOfDoubleOrNull),
         b"HRANDFIELD" => cmd
             .position(b"WITHVALUES")
-            .map(|_| ExpectedReturnType::ArrayOfKeyValuePairs),
+            .map(|_| ExpectedReturnType::ArrayOfPairs),
         b"ZRANDMEMBER" => cmd
             .position(b"WITHSCORES")
-            .map(|_| ExpectedReturnType::ArrayOfKeyValuePairs),
+            .map(|_| ExpectedReturnType::ArrayOfMemberScorePairs),
         b"ZADD" => cmd
             .position(b"INCR")
             .map(|_| ExpectedReturnType::DoubleOrNull),
@@ -517,33 +565,20 @@ mod tests {
     }
 
     #[test]
-    fn convert_array_of_kv_pairs() {
+    fn convert_to_array_of_pairs_return_type() {
         assert!(matches!(
             expected_type_for_cmd(
                 redis::cmd("HRANDFIELD")
                     .arg("key")
                     .arg("1")
                     .arg("withvalues")
             ),
-            Some(ExpectedReturnType::ArrayOfKeyValuePairs)
+            Some(ExpectedReturnType::ArrayOfPairs)
         ));
 
         assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key").arg("1")).is_none());
         assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key")).is_none());
 
-        assert!(matches!(
-            expected_type_for_cmd(
-                redis::cmd("ZRANDMEMBER")
-                    .arg("key")
-                    .arg("1")
-                    .arg("withscores")
-            ),
-            Some(ExpectedReturnType::ArrayOfKeyValuePairs)
-        ));
-
-        assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none());
-        assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none());
-
         let flat_array = Value::Array(vec![
             Value::BulkString(b"key1".to_vec()),
             Value::BulkString(b"value1".to_vec()),
@@ -561,38 +596,71 @@ mod tests {
             ]),
         ]);
         let converted_flat_array =
-            convert_to_expected_type(flat_array, Some(ExpectedReturnType::ArrayOfKeyValuePairs))
-                .unwrap();
+            convert_to_expected_type(flat_array, Some(ExpectedReturnType::ArrayOfPairs)).unwrap();
         assert_eq!(two_dimensional_array, converted_flat_array);
 
         let converted_two_dimensional_array = convert_to_expected_type(
             two_dimensional_array.clone(),
-            Some(ExpectedReturnType::ArrayOfKeyValuePairs),
+            Some(ExpectedReturnType::ArrayOfPairs),
         )
         .unwrap();
         assert_eq!(two_dimensional_array, converted_two_dimensional_array);
 
         let empty_array = Value::Array(vec![]);
-        let converted_empty_array = convert_to_expected_type(
-            empty_array.clone(),
-            Some(ExpectedReturnType::ArrayOfKeyValuePairs),
-        )
-        .unwrap();
+        let converted_empty_array =
+            convert_to_expected_type(empty_array.clone(), Some(ExpectedReturnType::ArrayOfPairs))
+                .unwrap();
         assert_eq!(empty_array, converted_empty_array);
 
         let converted_nil_value =
-            convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfKeyValuePairs))
-                .unwrap();
+            convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfPairs)).unwrap();
         assert_eq!(Value::Nil, converted_nil_value);
 
-        let array_of_doubles = Value::Array(vec![Value::Double(5.5)]);
+        let flat_array_unexpected_length =
+            Value::Array(vec![Value::BulkString(b"somekey".to_vec())]);
         assert!(convert_to_expected_type(
-            array_of_doubles,
-            Some(ExpectedReturnType::ArrayOfKeyValuePairs)
+            flat_array_unexpected_length,
+            Some(ExpectedReturnType::ArrayOfPairs)
         )
         .is_err());
     }
 
+    #[test]
+    fn convert_to_member_score_pairs_return_type() {
+        assert!(matches!(
+            expected_type_for_cmd(
+                redis::cmd("ZRANDMEMBER")
+                    .arg("key")
+                    .arg("1")
+                    .arg("withscores")
+            ),
+            Some(ExpectedReturnType::ArrayOfMemberScorePairs)
+        ));
+
+        assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none());
+        assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none());
+
+        // convert_to_array_of_pairs_return_type already tests most functionality since the conversion for ArrayOfPairs
+        // and ArrayOfMemberScorePairs is mostly the same. Here we also test that the scores are converted to double
+        // when the server response was a RESP2 flat array.
+        let flat_array = Value::Array(vec![
+            Value::BulkString(b"one".to_vec()),
+            Value::BulkString(b"1.0".to_vec()),
+            Value::BulkString(b"two".to_vec()),
+            Value::BulkString(b"2.0".to_vec()),
+        ]);
+        let expected_response = Value::Array(vec![
+            Value::Array(vec![Value::BulkString(b"one".to_vec()), Value::Double(1.0)]),
+            Value::Array(vec![Value::BulkString(b"two".to_vec()), Value::Double(2.0)]),
+        ]);
+        let converted_flat_array = convert_to_expected_type(
+            flat_array,
+            Some(ExpectedReturnType::ArrayOfMemberScorePairs),
+        )
+        .unwrap();
+        assert_eq!(expected_response, converted_flat_array);
+    }
+
     #[test]
     fn convert_zadd_only_if_incr_is_included() {
         assert!(matches!(

python/python/glide/async_commands/core.py

@@ -983,7 +983,7 @@ async def hrandfield_count(self, key: str, count: int) -> List[str]:
             key (str): The key of the hash.
             count (int): The number of field names to return.
                 If `count` is positive, returns unique elements.
-                If negative, allows for duplicates.
+                If `count` is negative, allows for duplicates elements.
 
         Returns:
             List[str]: A list of random field names from the hash.
@@ -1010,7 +1010,7 @@ async def hrandfield_withvalues(self, key: str, count: int) -> List[List[str]]:
             key (str): The key of the hash.
             count (int): The number of field names to return.
                 If `count` is positive, returns unique elements.
-                If negative, allows for duplicates.
+                If `count` is negative, allows for duplicates elements.
 
         Returns:
             List[List[str]]: A list of `[field_name, value]` lists, where `field_name` is a random field name from the
@@ -2746,6 +2746,89 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int:
             ),
         )
 
+    async def zrandmember(self, key: str) -> Optional[str]:
+        """
+        Returns a random element from the sorted set stored at 'key'.
+
+        See https://valkey.io/commands/zrandmember for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+
+        Returns:
+            Optional[str]: A string representing a random element from the sorted set.
+                If the sorted set does not exist or is empty, the response will be None.
+
+        Examples:
+            >>> await client.zrandmember("my_sorted_set")
+                "GLIDE"  # "GLIDE" is a random member of "my_sorted_set".
+            >>> await client.zrandmember("non_existing_sorted_set")
+                None  # "non_existing_sorted_set" is not an existing key, so None was returned.
+        """
+        return cast(
+            Optional[str],
+            await self._execute_command(RequestType.ZRandMember, [key]),
+        )
+
+    async def zrandmember_count(self, key: str, count: int) -> List[str]:
+        """
+        Retrieves random elements from the sorted set stored at 'key'.
+
+        See https://valkey.io/commands/zrandmember for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+            count (int): The number of elements to return.
+                If `count` is positive, returns unique elements.
+                If `count` is negative, allows for duplicates elements.
+
+        Returns:
+            List[str]: A list of elements from the sorted set.
+                If the sorted set does not exist or is empty, the response will be an empty list.
+
+        Examples:
+            >>> await client.zrandmember("my_sorted_set", -3)
+                ["GLIDE", "GLIDE", "PYTHON"]  # "GLIDE" and "PYTHON" are random members of "my_sorted_set".
+            >>> await client.zrandmember("non_existing_sorted_set", 3)
+                []  # "non_existing_sorted_set" is not an existing key, so an empty list was returned.
+        """
+        return cast(
+            List[str],
+            await self._execute_command(RequestType.ZRandMember, [key, str(count)]),
+        )
+
+    async def zrandmember_withscores(
+        self, key: str, count: int
+    ) -> List[List[Union[str, float]]]:
+        """
+        Retrieves random elements along with their scores from the sorted set stored at 'key'.
+
+        See https://valkey.io/commands/zrandmember for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+            count (int): The number of elements to return.
+                If `count` is positive, returns unique elements.
+                If `count` is negative, allows for duplicates elements.
+
+        Returns:
+            List[List[Union[str, float]]]: A list of `[member, score]` lists, where `member` is a random member from
+                the sorted set and `score` is the associated score.
+                If the sorted set does not exist or is empty, the response will be an empty list.
+
+        Examples:
+            >>> await client.zrandmember_withscores("my_sorted_set", -3)
+                [["GLIDE", 1.0], ["GLIDE", 1.0], ["PYTHON", 2.0]]  # "GLIDE" and "PYTHON" are random members of "my_sorted_set", and have scores of 1.0 and 2.0, respectively.
+            >>> await client.zrandmember_withscores("non_existing_sorted_set", 3)
+                []  # "non_existing_sorted_set" is not an existing key, so an empty list was returned.
+        """
+        return cast(
+            List[List[Union[str, float]]],
+            await self._execute_command(
+                RequestType.ZRandMember, [key, str(count), "WITHSCORES"]
+            ),
+        )
+
     async def invoke_script(
         self,
         script: Script,