Add support for LPUSH RPUSH and LRANGE

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)

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)

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