From 8f2a695ebaaccdd4fab0df9dc6d87fb8edcf1e69 Mon Sep 17 00:00:00 2001 From: nho45 Date: Thu, 21 Dec 2023 22:56:32 +0800 Subject: [PATCH] Add support for LPUSH RPUSH and LRANGE --- pyredis/commands.py | 71 ++++++++++++++-- pyredis/datastore.py | 51 ++++++++++++ tests/test_commands.py | 181 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 294 insertions(+), 9 deletions(-) diff --git a/pyredis/commands.py b/pyredis/commands.py index 25d0dac..7aa0e2c 100644 --- a/pyredis/commands.py +++ b/pyredis/commands.py @@ -86,10 +86,9 @@ def _handle_incr(command, datastore): if len(command) == 2: key = command[1].data.decode() try: - value = datastore[key] + 1 - except KeyError: - value = 1 # first time increase - datastore[key] = value + value = datastore.incr(key) + except TypeError: + return Error("ERR value is not an integer or out of range") return Integer(value) return Error("ERR wrong number of arguments for 'incr' command") @@ -98,16 +97,64 @@ def _handle_decr(command, datastore): if len(command) == 2: key = command[1].data.decode() try: - value = datastore[key] - 1 - except KeyError: - value = -1 + value = datastore.decr(key) except TypeError: - return Error("value is not an integer or out of range") - datastore[key] = value + return Error("ERR value is not an integer or out of range") return Integer(value) return Error("ERR wrong number of arguments for 'decr' command") +def _handle_lpush(command, datastore): + if len(command) >= 2: + count = 0 + key = command[1].data.decode() + + try: + for c in command[2:]: + item = c.data.decode() + count = datastore.prepend(key, item) + return Integer(count) + except TypeError: + return Error( + "WRONGTYPE Operation against a key holding the wrong kind of value" + ) + return Error("ERR wrong number of arguments for 'lpush' command") + + +def _handle_lrange(command, datastore): + if len(command) == 4: + key = command[1].data.decode() + start = int(command[2].data.decode()) + stop = int(command[3].data.decode()) + + try: + items = datastore.lrange(key, start, stop) + return Array([BulkString(i) for i in items]) + except TypeError: + return Error( + "WRONGTYPE Operatino against a key holding the wrong kind of value" + ) + + return Error("ERR wrong number of arguments for 'lrange' command") + + +def _handle_rpush(command, datastore): + if len(command) >= 2: + count = 0 + key = command[1].data.decode() + + try: + for c in command[2:]: + item = c.data.decode() + count = datastore.append(key, item) + return Integer(count) + except TypeError: + return Error( + "WRONGTYPE Operation against a key holding the wrong kind of value" + ) + return Error("ERR wrong number of arguments for 'rpush' command") + + def _handle_unrecognised_command(command, *args): args = " ".join((f"'{c.data.decode()}'" for c in command[1:])) return Error( @@ -135,4 +182,10 @@ def handle_command(command, datastore): return _handle_incr(command, datastore) case "DECR": return _handle_decr(command, datastore) + case "LPUSH": + return _handle_lpush(command, datastore) + case "RPUSH": + return _handle_rpush(command, datastore) + case "LRANGE": + return _handle_lrange(command, datastore) return _handle_unrecognised_command(command) diff --git a/pyredis/datastore.py b/pyredis/datastore.py index 2af90bb..ab66ffc 100644 --- a/pyredis/datastore.py +++ b/pyredis/datastore.py @@ -2,6 +2,8 @@ from dataclasses import dataclass from typing import Any from time import time +from itertools import islice +from collections import deque import random import logging @@ -54,6 +56,26 @@ def __contains__(self, key): with self._lock: return key in self._data + def incr(self, key): + with self._lock: + item = self._data.get(key, DataEntry(0)) + try: + value = int(item.value) + 1 + except ValueError: + raise TypeError + item.value = str(value) + self._data[key] = item + return value + + def decr(self, key): + with self._lock: + try: + value = int(self._data.get(key, DataEntry(0)).value) - 1 + except ValueError: + raise TypeError + self._data[key].value = str(value) + return value + def set_with_expiry(self, key, value, expiry: int): with self._lock: calculated_expiry = int(time() * 1000) + expiry # in miliseconds @@ -82,3 +104,32 @@ def remove_expired_keys(self): # if more than if expired_count > EXPIRY_TEST_SAMPLE_SIZE * 0.25: self.remove_expired_keys() + + def append(self, key, value): + with self._lock: + item = self._data.get(key, DataEntry(deque())) + if not isinstance(item.value, deque): + raise TypeError + item.value.append(value) + self._data[key] = item + return len(item.value) + + def lrange(self, key, start, stop): + with self._lock: + item = self._data.get(key, DataEntry(deque())) + if not isinstance(item.value, deque): + raise TypeError + + return list(islice(item.value, start, stop)) + + def prepend(self, key, value): + with self._lock: + item = self._data.get(key, DataEntry(deque())) + print("HERE") + if not isinstance(item.value, deque): + print(item.value) + raise TypeError + print("HERE 2") + item.value.insert(0, value) + self._data[key] = item + return len(item.value) diff --git a/tests/test_commands.py b/tests/test_commands.py index 953c38d..4c6d504 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -5,6 +5,8 @@ from pyredis.datastore import DataStore from pyredis.types import Array, BulkString, Error, Integer, SimpleString +from collections import deque + @pytest.fixture(scope="module") def datastore(): @@ -117,6 +119,30 @@ def datastore(): ), Integer(2), ), + # Incr Tests + ( + Array([BulkString(b"incr")]), + Error("ERR wrong number of arguments for 'incr' command"), + ), + ( + Array([BulkString(b"incr"), SimpleString(b"key")]), + Error("ERR value is not an integer or out of range"), + ), + # Decr Tests + ( + Array([BulkString(b"decr")]), + Error("ERR wrong number of arguments for 'decr' command"), + ), + # Lpush Tests + ( + Array([BulkString(b"lpush")]), + Error("ERR wrong number of arguments for 'lpush' command"), + ), + # Rpush Tests + ( + Array([BulkString(b"rpush")]), + Error("ERR wrong number of arguments for 'rpush' command"), + ), ], ) def test_handle_command(command, expected, datastore): @@ -241,3 +267,158 @@ def test_handle_decr(): Array([BulkString(b"decr"), SimpleString(b"kd")]), datastore ) assert result == Integer(0) + + +# Lpush Tests +def test_handle_lpush_lrange(): + datastore = DataStore() + result = handle_command( + Array([BulkString(b"lpush"), SimpleString(b"klp"), SimpleString(b"second")]), + datastore, + ) + assert result == Integer(1) + result = handle_command( + Array([BulkString(b"lpush"), SimpleString(b"klp"), SimpleString(b"first")]), + datastore, + ) + assert result == Integer(2) + result = handle_command( + Array( + [ + BulkString(b"lrange"), + SimpleString(b"klp"), + BulkString(b"0"), + BulkString(b"2"), + ] + ), + datastore, + ) + assert result == Array(data=[BulkString("first"), BulkString("second")]) + + +# Rpush Tests +def test_handle_rpush_lrange(): + datastore = DataStore() + result = handle_command( + Array([BulkString(b"rpush"), SimpleString(b"krp"), SimpleString(b"first")]), + datastore, + ) + assert result == Integer(1) + result = handle_command( + Array([BulkString(b"rpush"), SimpleString(b"krp"), SimpleString(b"second")]), + datastore, + ) + assert result == Integer(2) + result = handle_command( + Array( + [ + BulkString(b"lrange"), + SimpleString(b"krp"), + BulkString(b"0"), + BulkString(b"2"), + ] + ), + datastore, + ) + assert result == Array(data=[BulkString("first"), BulkString("second")]) + + +@pytest.fixture +def ds(): + return DataStore() + + +def test_initial_data_invalid_type(): + with pytest.raises(TypeError): + ds = DataStore("string") + + +def test_initial_data(): + ds = DataStore({"k1": 1, "k2": "v2"}) + assert ds["k1"] == 1 + assert ds["k2"] == "v2" + + +def test_in(ds): + ds["key"] = 1 + + assert "key" in ds + assert "key2" not in ds + + +def test_get_item(ds): + ds["key"] = 1 + assert ds["key"] == 1 + + +def test_set_item(ds): + l = ds.append("key", 1) + assert l == 1 + assert ds["key"] == deque([1]) + + +def test_incr(ds): + ds["k"] = "1" + res = ds.incr("k") + assert res == 2 + res = ds.incr("k") + assert res == 3 + + +def test_decr(ds): + ds["k"] = "1" + ds.incr("k") + ds.incr("k") + res = ds.incr("k") + assert res == 4 + res = ds.decr("k") + assert res == 3 + res = ds.decr("k") + assert res == 2 + + +def test_append(ds): + num_entries = ds.append("key", 1) + assert num_entries == 1 + assert ds["key"] == deque([1]) + + +def test_preppend(ds): + ds.append("key", 1) + ds.prepend("key", 2) + assert ds["key"] == deque([2, 1]) + + +def test_expire_on_read(ds): + ds.set_with_expiry("key", "value", 0.01) + sleep(0.15) + with pytest.raises(KeyError): + ds["key"] + + +def test_remove_expired_keys_empty(): + ds = DataStore() + ds.remove_expired_keys() + + +def _fill_ds(ds, size, percent_expired): + num_expired = int(size * (percent_expired / 100)) + + # items without expiry + for i in range(size - num_expired): + ds[f"{i}"] = i + + # items with expiry and that will have expired + for i in range(num_expired): + ds.set_with_expiry(f"e_{i}", i, -1) + + +@pytest.mark.parametrize("size, percent_expired", [(20, 10), (200, 100)]) +def test_remove_expired_keys(size, percent_expired): + expected_len_after_expiry = size - (size * (percent_expired / 100)) + + ds = DataStore() + _fill_ds(ds, size, percent_expired) + + ds.remove_expired_keys() + assert len(ds._data) == expected_len_after_expiry