Add step 5 get set with expiry

pyredis/commands.py

@@ -22,11 +22,29 @@ def _handle_ping(command):
 
 
 def _handle_set(command, datastore):
-    if len(command) >= 3:
+    length = len(command)
+    if length >= 3:
         key = command[1].data.decode()
         value = command[2].data.decode()
-        datastore[key] = value
-        return SimpleString("OK")
+
+        if length == 3:
+            datastore[key] = value
+            return SimpleString("OK")
+        elif length == 5:
+            expiry_mode = command[3].data.decode()
+            try:
+                expiry = int(command[4].data.decode())
+            except ValueError:
+                return Error("ERR value is not an integer or out of range")
+
+            if expiry_mode == "ex":
+                datastore.set_with_expiry(key, value, expiry * 1000)
+                return SimpleString("OK")
+            elif expiry_mode == "px":
+                datastore.set_with_expiry(key, value, expiry)
+                return SimpleString("OK")
+        return Error("ERR syntax error")
+
     return Error("ERR wrong number of arguments for 'set' command")
 
 

pyredis/datastore.py

@@ -1,16 +1,52 @@
 from threading import Lock
+from dataclasses import dataclass
+from typing import Any
+from time import time
+
+
+@dataclass
+class DataEntry:
+    """Class to represent a data entry. Contains the data and the expiry in milisecond."""
+
+    value: Any
+    expiry: int = 0
 
 
 class DataStore:
-    def __init__(self):
-        self._data = dict()
+    """
+    The core data store, provides a thread safe dictionary extended with
+    the interface needed to support Redis functionality.
+    """
+
+    def __init__(self, initial_data=None):
+        self._data: dict[str, DataEntry] = dict()
         self._lock = Lock()
+        if initial_data:
+            if not isinstance(initial_data, dict):
+                raise TypeError("Initial Data should be of type dict")
+
+            for key, value in initial_data.items():
+                self._data[key] = DataEntry(value)
 
     def __getitem__(self, key):
         with self._lock:
             item = self._data[key]
-        return item
+
+            # if key expired
+            if item.expiry and item.expiry < int(time() * 1000):
+                del self._data[key]
+                raise KeyError
+
+            return item.value
 
     def __setitem__(self, key, value):
         with self._lock:
-            self._data[key] = value
+            self._data[key] = DataEntry(value)
+
+    def set_with_expiry(self, key, value, expiry: int):
+        with self._lock:
+            calculated_expiry = int(time() * 1000) + expiry  # in miliseconds
+            self._data[key] = DataEntry(value, calculated_expiry)
+
+    def check_expiry(datastore):
+        pass

tests/test_commands.py

@@ -1,4 +1,5 @@
 import pytest
+from time import sleep, time_ns
 
 from pyredis.commands import handle_command
 from pyredis.datastore import DataStore
@@ -45,8 +46,127 @@ def datastore():
             Array([BulkString(b"get"), SimpleString(b"always_exist_key")]),
             BulkString("default"),
         ),
+        # Set with Expire Errors
+        (
+            Array(
+                [
+                    BulkString(b"set"),
+                    SimpleString(b"key"),
+                    SimpleString(b"value"),
+                    SimpleString(b"ex"),
+                ]
+            ),
+            Error("ERR syntax error"),
+        ),
+        (
+            Array(
+                [
+                    BulkString(b"set"),
+                    SimpleString(b"key"),
+                    SimpleString(b"value"),
+                    SimpleString(b"px"),
+                ]
+            ),
+            Error("ERR syntax error"),
+        ),
+        (
+            Array(
+                [
+                    BulkString(b"set"),
+                    SimpleString(b"key"),
+                    SimpleString(b"value"),
+                    SimpleString(b"foo"),
+                ]
+            ),
+            Error("ERR syntax error"),
+        ),
     ],
 )
 def test_handle_command(command, expected, datastore):
     result = handle_command(command, datastore)
     assert result == expected
+
+
+def test_get_with_expiry(datastore):
+    key = "key"
+    value = "value"
+    px = 100
+
+    command = [
+        BulkString(b"set"),
+        SimpleString(b"key"),
+        SimpleString(b"value"),
+        BulkString(b"px"),
+        BulkString(f"{px}".encode()),
+    ]
+    result = handle_command(command, datastore)
+    assert result == SimpleString("OK")
+    sleep((px + 100) / 1000)
+    command = [BulkString(b"get"), SimpleString(b"key")]
+    result = handle_command(command, datastore)
+    assert result == BulkString(None)
+
+
+def test_set_with_expiry():
+    datastore = DataStore()
+    key = "key"
+    value = "value"
+    ex = 1
+    px = 100
+
+    base_command = [BulkString(b"set"), SimpleString(b"key"), SimpleString(b"value")]
+
+    # seconds
+    command = base_command[:]
+    command.extend([BulkString(b"ex"), BulkString(f"{ex}".encode())])
+    expected_expiry = time_ns() + (ex * 10**9)
+    result = handle_command(command, datastore)
+    assert result == SimpleString("OK")
+    stored = datastore._data[key]
+    assert stored.value == value
+    diff = -expected_expiry - stored.expiry
+    assert diff < 10000
+
+    # milliseconds
+    command = base_command[:]
+    command.extend([BulkString(b"px"), BulkString(f"{px}".encode())])
+    expected_expiry = time_ns() + (ex * 10**6)
+    result = handle_command(command, datastore)
+    assert result == SimpleString("OK")
+    stored = datastore._data[key]
+    assert stored.value == value
+    diff = -expected_expiry - stored.expiry
+    assert diff < 10000
+
+
+def test_get_with_expiry():
+    datastore = DataStore()
+    key = "key"
+    value = "value"
+    px = 100
+
+    command = [
+        BulkString(b"set"),
+        SimpleString(b"key"),
+        SimpleString(b"value"),
+        BulkString(b"px"),
+        BulkString(f"{px}".encode()),
+    ]
+    result = handle_command(command, datastore)
+    assert result == SimpleString("OK")
+    sleep((px + 100) / 1000)
+    command = [BulkString(b"get"), SimpleString(b"key")]
+    result = handle_command(command, datastore)
+    assert result == BulkString(None)
+
+
+def test_set_and_get_item(datastore):
+    datastore["key"] = 1
+    assert datastore["key"] == 1
+
+
+def test_expire_on_read(datastore):
+    datastore.set_with_expiry("key", "value", 0.01)
+    sleep(0.15)
+    with pytest.raises(KeyError):
+        datastore["key"]

tests/test_protocol.py

@@ -1,4 +1,5 @@
 import pytest
+from time import sleep
 
 from pyredis.protocol import extract_frame_from_buffer, encode_message
 from pyredis.datastore import DataStore
@@ -84,9 +85,3 @@ def test_read_frame(buffer, expected):
 def test_encode_message(message, expected):
     encoded_message = encode_message(message)
     assert encoded_message == expected
-
-
-def test_set_and_get_item():
-    ds = DataStore()
-    ds["key"] = 1
-    assert ds["key"] == 1