🎨 (store) Tweak prefix/key handling

anystore/store/base.py

@@ -30,31 +30,31 @@ def __init__(self, **data):
         data["scheme"] = urlparse(str(uri)).scheme
         super().__init__(**data)
 
-    def _write(self, key: Uri, value: Value, **kwargs) -> None:
+    def _write(self, key: str, value: Value, **kwargs) -> None:
         """
         Write value with key to acutal backend
         """
         raise NotImplementedError
 
-    def _read(self, key: Uri, raise_on_nonexist: bool | None = True, **kwargs) -> Any:
+    def _read(self, key: str, raise_on_nonexist: bool | None = True, **kwargs) -> Any:
         """
         Read key from actual backend
         """
         raise NotImplementedError
 
-    def _delete(self, key: Uri) -> None:
+    def _delete(self, key: str) -> None:
         """
         Delete key from actual backend
         """
         raise NotImplementedError
 
-    def _stream(self, key: Uri, raise_on_nonexist: bool | None = True, **kwargs) -> Any:
+    def _stream(self, key: str, raise_on_nonexist: bool | None = True, **kwargs) -> Any:
         """
         Stream key line by line from actual backend (for file-like powered backend)
         """
         raise NotImplementedError
 
-    def _exists(self, key: Uri) -> bool:
+    def _exists(self, key: str) -> bool:
         """
         Check if the given key exists
         """
@@ -102,7 +102,7 @@ def get(
 
     def pop(self, key: Uri, *args, **kwargs) -> Any:
         value = self.get(key, *args, **kwargs)
-        self._delete(key)
+        self._delete(self.get_key(key))
         return value
 
     def stream(
@@ -158,7 +158,7 @@ def put(
         )
 
     def exists(self, key: Uri) -> bool:
-        return self._exists(key)
+        return self._exists(self.get_key(key))
 
     def ensure_kwargs(self, **kwargs) -> dict[str, Any]:
         config = clean_dict(self.backend_config)

anystore/store/fs.py

@@ -10,40 +10,40 @@
 from anystore.io import smart_read, smart_stream, smart_write
 from anystore.exceptions import DoesNotExist
 from anystore.store.base import BaseStore
-from anystore.types import Uri, Value
+from anystore.types import Value
 
 
 class Store(BaseStore):
-    def _write(self, key: Uri, value: Value, **kwargs) -> None:
+    def _write(self, key: str, value: Value, **kwargs) -> None:
         kwargs.pop("ttl", None)
-        smart_write(str(key), value, **kwargs)
+        smart_write(key, value, **kwargs)
 
     def _read(
-        self, key: Uri, raise_on_nonexist: bool | None = True, **kwargs
+        self, key: str, raise_on_nonexist: bool | None = True, **kwargs
     ) -> Value | None:
         try:
-            return smart_read(str(key), **kwargs)
+            return smart_read(key, **kwargs)
         except FileNotFoundError:
             if raise_on_nonexist:
                 raise DoesNotExist(f"Key does not exist: `{key}`")
             return None
 
     def _stream(
-        self, key: Uri, raise_on_nonexist: bool | None = True, **kwargs
+        self, key: str, raise_on_nonexist: bool | None = True, **kwargs
     ) -> Generator[Value, None, None]:
         try:
-            yield from smart_stream(str(key), **kwargs)
+            yield from smart_stream(key, **kwargs)
         except FileNotFoundError:
             if raise_on_nonexist:
                 raise DoesNotExist(f"Key does not exist: `{key}`")
 
-    def _exists(self, key: Uri) -> bool:
+    def _exists(self, key: str) -> bool:
         fs = fsspec.filesystem(self.scheme)
-        return fs.exists(self.get_key(key))
+        return fs.exists(key)
 
-    def _delete(self, key: Uri) -> None:
+    def _delete(self, key: str) -> None:
         fs = fsspec.filesystem(self.scheme)
-        fs.delete(self.get_key(key))
+        fs.delete(key)
 
     def _get_key_prefix(self) -> str:
         return str(self.uri).rstrip("/")

anystore/store/redis.py

@@ -2,30 +2,26 @@
 Store backend using redis-like stores such as Redis, Fakeredis or Apache Kvrocks
 """
 
-from typing import Any, Generator, TYPE_CHECKING
 from functools import cache
+from typing import Any, Generator
 
+import fakeredis
 import redis
 
 from anystore.exceptions import DoesNotExist
 from anystore.logging import get_logger
 from anystore.settings import Settings
 from anystore.store.base import BaseStore
-from anystore.types import Uri, Value
-
-if TYPE_CHECKING:
-    import fakeredis
+from anystore.types import Value
 
 
 log = get_logger(__name__)
 
 
 @cache
-def get_redis(uri: str) -> "fakeredis.FakeStrictRedis | redis.Redis":
+def get_redis(uri: str) -> fakeredis.FakeStrictRedis | redis.Redis:
     settings = Settings()
     if settings.redis_debug:
-        import fakeredis
-
         con = fakeredis.FakeStrictRedis()
         con.ping()
         log.info("Redis connected: `fakeredis`")
@@ -37,32 +33,31 @@ def get_redis(uri: str) -> "fakeredis.FakeStrictRedis | redis.Redis":
 
 
 class RedisStore(BaseStore):
-    def _write(self, key: Uri, value: Value, **kwargs) -> None:
+    def _write(self, key: str, value: Value, **kwargs) -> None:
         ttl = kwargs.pop("ttl", None) or None
         con = get_redis(self.uri)
-        con.set(str(key), value, ex=ttl, **kwargs)
+        con.set(key, value, ex=ttl, **kwargs)
 
-    def _read(self, key: Uri, raise_on_nonexist: bool | None = True, **kwargs) -> Any:
+    def _read(self, key: str, raise_on_nonexist: bool | None = True, **kwargs) -> Any:
         con = get_redis(self.uri)
-        key = str(key)
         # `None` could be stored as an actual value, to implement `raise_on_nonexist`
         # we need to check this first:
         if raise_on_nonexist and not con.exists(key):
             raise DoesNotExist
-        res = con.get(str(key))
+        res = con.get(key)
         # mimic fs read mode:
         if kwargs.get("mode") == "r" and isinstance(res, bytes):
             res = res.decode()
         return res
 
-    def _exists(self, key: Uri) -> bool:
+    def _exists(self, key: str) -> bool:
         con = get_redis(self.uri)
-        res = con.exists(self.get_key(key))
+        res = con.exists(key)
         return bool(res)
 
-    def _delete(self, key: Uri) -> None:
+    def _delete(self, key: str) -> None:
         con = get_redis(self.uri)
-        con.delete(self.get_key(key))
+        con.delete(key)
 
     def _get_key_prefix(self) -> str:
         if self.backend_config is not None:

anystore/store/sql.py

@@ -26,7 +26,7 @@
 from anystore.exceptions import DoesNotExist
 from anystore.settings import SqlSettings
 from anystore.store.base import BaseStore
-from anystore.types import Uri, Value
+from anystore.types import Value
 
 
 settings = SqlSettings()
@@ -97,8 +97,7 @@ def __init__(self, **data):
         self._conn = engine.connect()
         self._sqlite = "sqlite" in engine.name.lower()
 
-    def _write(self, key: Uri, value: Value, **kwargs) -> None:
-        key = str(key)
+    def _write(self, key: str, value: Value, **kwargs) -> None:
         ttl = kwargs.pop("ttl", None) or None
         # FIXME on conflict / on duplicate key
         exists = select(self._table).where(self._table.c.key == key)
@@ -114,9 +113,8 @@ def _write(self, key: Uri, value: Value, **kwargs) -> None:
         self._conn.commit()
 
     def _read(
-        self, key: Uri, raise_on_nonexist: bool | None = True, **kwargs
+        self, key: str, raise_on_nonexist: bool | None = True, **kwargs
     ) -> Value | None:
-        key = str(key)
         stmt = select(self._table).where(self._table.c.key == key)
         res = self._conn.execute(stmt).first()
         if res:
@@ -133,16 +131,14 @@ def _read(
         if raise_on_nonexist:
             raise DoesNotExist
 
-    def _exists(self, key: Uri) -> bool:
-        key = str(key)
+    def _exists(self, key: str) -> bool:
         stmt = select(self._table).where(self._table.c.key == key)
         stmt = select(stmt.exists())
         for res in self._conn.execute(stmt).first():
             return bool(res)
         return False
 
-    def _delete(self, key: Uri) -> None:
-        key = str(key)
+    def _delete(self, key: str) -> None:
         stmt = delete(self._table).where(self._table.c.key == key)
         self._conn.execute(stmt)