Skip to content

Commit

Permalink
Python: add BITCOUNT command
Browse files Browse the repository at this point in the history
  • Loading branch information
aaron-congo committed Jun 17, 2024
1 parent 0b30286 commit 4da8f52
Show file tree
Hide file tree
Showing 7 changed files with 191 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
* Node: Added OBJECT IDLETIME command ([#1567](https://github.com/aws/glide-for-redis/pull/1567))
* Node: Added OBJECT REFCOUNT command ([#1568](https://github.com/aws/glide-for-redis/pull/1568))
* Python: Added SETBIT command ([#1571](https://github.com/aws/glide-for-redis/pull/1571))
* Python: Added BITCOUNT command ([#1592](https://github.com/aws/glide-for-redis/pull/1592))

### Breaking Changes
* Node: Update XREAD to return a Map of Map ([#1494](https://github.com/aws/glide-for-redis/pull/1494))
Expand Down
1 change: 1 addition & 0 deletions python/python/glide/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0

from glide.async_commands.bitmap import BitmapIndexType, OffsetOptions
from glide.async_commands.command_args import Limit, ListDirection, OrderBy
from glide.async_commands.core import (
ConditionalChange,
Expand Down
50 changes: 50 additions & 0 deletions python/python/glide/async_commands/bitmap.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
# Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0
from enum import Enum
from typing import List, Optional


class BitmapIndexType(Enum):
"""
Enumeration specifying if index arguments are BYTE indexes or BIT indexes. Can be specified in `OffsetOptions`,
which is an optional argument to the `BITCOUNT` command.
Since: Redis version 7.0.0.
"""

BYTE = "BYTE"
"""
Specifies that indexes provided to `OffsetOptions` are byte indexes.
"""
BIT = "BIT"
"""
Specifies that indexes provided to `OffsetOptions` are bit indexes.
"""


class OffsetOptions:
def __init__(
self, start: int, end: int, index_type: Optional[BitmapIndexType] = None
):
"""
Represents offsets specifying a string interval to analyze in the `BITCOUNT` command. The offsets are
zero-based indexes, with `0` being the first index of the string, `1` being the next index and so on.
The offsets can also be negative numbers indicating offsets starting at the end of the string, with `-1` being
the last index of the string, `-2` being the penultimate, and so on.
Args:
start (int): The starting offset index.
end (int): The ending offset index.
index_type (Optional[BitmapIndexType]): The index offset type. This option can only be specified if you are
using Redis version 7.0.0 or above. Could be either `BitmapIndexType.BYTE` or `BitmapIndexType.BIT`.
If no index type is provided, the indexes will be assumed to be byte indexes.
"""
self.start = start
self.end = end
self.index_type = index_type

def to_args(self) -> List[str]:
args = [str(self.start), str(self.end)]
if self.index_type is not None:
args.append(self.index_type.value)

return args
36 changes: 36 additions & 0 deletions python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
get_args,
)

from glide.async_commands.bitmap import OffsetOptions
from glide.async_commands.command_args import Limit, ListDirection, OrderBy
from glide.async_commands.sorted_set import (
AggregationType,
Expand Down Expand Up @@ -4033,6 +4034,41 @@ async def pfmerge(self, destination: str, source_keys: List[str]) -> TOK:
),
)

async def bitcount(self, key: str, options: Optional[OffsetOptions] = None) -> int:
"""
Counts the number of set bits (population counting) in the string stored at `key`. The `options` argument can
optionally be provided to count the number of bits in a specific string interval.
See https://valkey.io/commands/bitcount for more details.
Args:
key (str): The key for the string to count the set bits of.
options (Optional[OffsetOptions]): The offset options.
Returns:
int: If `options` is provided, returns the number of set bits in the string interval specified by `options`.
If `options` is not provided, returns the number of set bits in the string stored at `key`.
Otherwise, if `key` is missing, returns `0` as it is treated as an empty string.
Examples:
>>> await client.bitcount("my_key1")
2 # The string stored at "my_key1" contains 2 set bits.
>>> await client.bitcount("my_key2", OffsetOptions(1, 3))
2 # The second to fourth bytes of the string stored at "my_key2" contain 2 set bits.
>>> await client.bitcount("my_key3", OffsetOptions(1, 1, BitmapIndexType.BIT))
1 # Indicates that the second bit of the string stored at "my_key3" is set.
>>> await client.bitcount("my_key3", OffsetOptions(-1, -1, BitmapIndexType.BIT))
1 # Indicates that the last bit of the string stored at "my_key3" is set.
"""
args = [key]
if options is not None:
args = args + options.to_args()

return cast(
int,
await self._execute_command(RequestType.BitCount, args),
)

async def setbit(self, key: str, offset: int, value: int) -> int:
"""
Sets or clears the bit at `offset` in the string value stored at `key`. The `offset` is a zero-based index,
Expand Down
25 changes: 25 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import threading
from typing import List, Mapping, Optional, Tuple, TypeVar, Union

from glide.async_commands.bitmap import OffsetOptions
from glide.async_commands.command_args import Limit, ListDirection, OrderBy
from glide.async_commands.core import (
ConditionalChange,
Expand Down Expand Up @@ -2775,6 +2776,30 @@ def pfmerge(
"""
return self.append_command(RequestType.PfMerge, [destination] + source_keys)

def bitcount(
self: TTransaction, key: str, options: Optional[OffsetOptions] = None
) -> TTransaction:
"""
Counts the number of set bits (population counting) in a string stored at `key`. The `options` argument can
optionally be provided to count the number of bits in a specific string interval.
See https://valkey.io/commands/bitcount for more details.
Args:
key (str): The key for the string to count the set bits of.
options (Optional[OffsetOptions]): The offset options.
Command response:
int: If `options` is provided, returns the number of set bits in the string interval specified by `options`.
If `options` is not provided, returns the number of set bits in the string stored at `key`.
Otherwise, if `key` is missing, returns `0` as it is treated as an empty string.
"""
args = [key]
if options is not None:
args = args + options.to_args()

return self.append_command(RequestType.BitCount, args)

def setbit(self: TTransaction, key: str, offset: int, value: int) -> TTransaction:
"""
Sets or clears the bit at `offset` in the string value stored at `key`. The `offset` is a zero-based index,
Expand Down
65 changes: 65 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

import pytest
from glide import ClosingError, RequestError, Script
from glide.async_commands.bitmap import BitmapIndexType, OffsetOptions
from glide.async_commands.command_args import Limit, ListDirection, OrderBy
from glide.async_commands.core import (
ConditionalChange,
Expand Down Expand Up @@ -4029,6 +4030,70 @@ async def test_pfmerge(self, redis_client: TRedisClient):
with pytest.raises(RequestError):
assert await redis_client.pfmerge(string_key, [key3])

@pytest.mark.parametrize("cluster_mode", [True, False])
@pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
async def test_bitcount(self, redis_client: TRedisClient):
key1 = get_random_string(10)
set_key = get_random_string(10)
non_existing_key = get_random_string(10)
value = "foobar"

assert await redis_client.set(key1, value) == OK
assert await redis_client.bitcount(key1) == 26
assert await redis_client.bitcount(key1, OffsetOptions(1, 1)) == 6
assert await redis_client.bitcount(key1, OffsetOptions(0, -5)) == 10
assert await redis_client.bitcount(non_existing_key, OffsetOptions(5, 30)) == 0
assert await redis_client.bitcount(non_existing_key) == 0

# key exists, but it is not a string
assert await redis_client.sadd(set_key, [value]) == 1
with pytest.raises(RequestError):
await redis_client.bitcount(set_key)
with pytest.raises(RequestError):
await redis_client.bitcount(set_key, OffsetOptions(1, 1))

if await check_if_server_version_lt(redis_client, "7.0.0"):
# exception thrown because BIT and BYTE options were implemented after 7.0.0
with pytest.raises(RequestError):
await redis_client.bitcount(
key1, OffsetOptions(2, 5, BitmapIndexType.BYTE)
)
with pytest.raises(RequestError):
await redis_client.bitcount(
key1, OffsetOptions(2, 5, BitmapIndexType.BIT)
)
else:
assert (
await redis_client.bitcount(
key1, OffsetOptions(2, 5, BitmapIndexType.BYTE)
)
== 16
)
assert (
await redis_client.bitcount(
key1, OffsetOptions(5, 30, BitmapIndexType.BIT)
)
== 17
)
assert (
await redis_client.bitcount(
key1, OffsetOptions(5, -5, BitmapIndexType.BIT)
)
== 23
)
assert (
await redis_client.bitcount(
non_existing_key, OffsetOptions(5, 30, BitmapIndexType.BIT)
)
== 0
)

# key exists but it is not a string
with pytest.raises(RequestError):
await redis_client.bitcount(
set_key, OffsetOptions(1, 1, BitmapIndexType.BIT)
)

@pytest.mark.parametrize("cluster_mode", [True, False])
@pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
async def test_setbit(self, redis_client: TRedisClient):
Expand Down
13 changes: 13 additions & 0 deletions python/python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import pytest
from glide import RequestError
from glide.async_commands.bitmap import BitmapIndexType, OffsetOptions
from glide.async_commands.command_args import Limit, ListDirection, OrderBy
from glide.async_commands.core import (
GeospatialData,
Expand Down Expand Up @@ -57,6 +58,7 @@ async def transaction_test(
key17 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # sort
key18 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # sort
key19 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # bitmap
key20 = "{{{}}}:{}".format(keyslot, get_random_string(3)) # bitmap

value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S")
value2 = get_random_string(5)
Expand Down Expand Up @@ -350,6 +352,17 @@ async def transaction_test(
transaction.setbit(key19, 1, 0)
args.append(1)

transaction.set(key20, "foobar")
args.append(OK)
transaction.bitcount(key20)
args.append(26)
transaction.bitcount(key20, OffsetOptions(1, 1))
args.append(6)

if not await check_if_server_version_lt(redis_client, "7.0.0"):
transaction.bitcount(key20, OffsetOptions(5, 30, BitmapIndexType.BIT))
args.append(17)

transaction.geoadd(
key12,
{
Expand Down

0 comments on commit 4da8f52

Please sign in to comment.