Redis(Cache) now (re)supports function as a key_prefix

Author: dbascoules

Diff for: src/cachelib/redis.py

@@ -57,6 +57,11 @@ def __init__(
             self._read_client = self._write_client = host
         self.key_prefix = key_prefix or ""
 
+    def _get_prefix(self):
+        return (
+            self.key_prefix if isinstance(self.key_prefix, str) else self.key_prefix()
+        )
+
     def _normalize_timeout(self, timeout: _t.Optional[int]) -> int:
         """Normalize timeout by setting it to default of 300 if
         not defined (None) or -1 if explicitly set to zero.
@@ -69,11 +74,11 @@ def _normalize_timeout(self, timeout: _t.Optional[int]) -> int:
         return timeout
 
     def get(self, key: str) -> _t.Any:
-        return self.serializer.loads(self._read_client.get(self.key_prefix + key))
+        return self.serializer.loads(self._read_client.get(self._get_prefix() + key))
 
     def get_many(self, *keys: str) -> _t.List[_t.Any]:
         if self.key_prefix:
-            prefixed_keys = [self.key_prefix + key for key in keys]
+            prefixed_keys = [self._get_prefix() + key for key in keys]
         else:
             prefixed_keys = list(keys)
         return [self.serializer.loads(x) for x in self._read_client.mget(prefixed_keys)]
@@ -82,20 +87,20 @@ def set(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> _t.A
         timeout = self._normalize_timeout(timeout)
         dump = self.serializer.dumps(value)
         if timeout == -1:
-            result = self._write_client.set(name=self.key_prefix + key, value=dump)
+            result = self._write_client.set(name=self._get_prefix() + key, value=dump)
         else:
             result = self._write_client.setex(
-                name=self.key_prefix + key, value=dump, time=timeout
+                name=self._get_prefix() + key, value=dump, time=timeout
             )
         return result
 
     def add(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> _t.Any:
         timeout = self._normalize_timeout(timeout)
         dump = self.serializer.dumps(value)
-        created = self._write_client.setnx(name=self.key_prefix + key, value=dump)
+        created = self._write_client.setnx(name=self._get_prefix() + key, value=dump)
         # handle case where timeout is explicitly set to zero
         if created and timeout != -1:
-            self._write_client.expire(name=self.key_prefix + key, time=timeout)
+            self._write_client.expire(name=self._get_prefix() + key, time=timeout)
         return created
 
     def set_many(
@@ -109,41 +114,40 @@ def set_many(
         for key, value in mapping.items():
             dump = self.serializer.dumps(value)
             if timeout == -1:
-                pipe.set(name=self.key_prefix + key, value=dump)
+                pipe.set(name=self._get_prefix() + key, value=dump)
             else:
-                pipe.setex(name=self.key_prefix + key, value=dump, time=timeout)
+                pipe.setex(name=self._get_prefix() + key, value=dump, time=timeout)
         results = pipe.execute()
-        res = zip(mapping.keys(), results)  # noqa: B905
-        return [k for k, was_set in res if was_set]
+        return [k for k, was_set in zip(mapping.keys(), results) if was_set]
 
     def delete(self, key: str) -> bool:
-        return bool(self._write_client.delete(self.key_prefix + key))
+        return bool(self._write_client.delete(self._get_prefix() + key))
 
     def delete_many(self, *keys: str) -> _t.List[_t.Any]:
         if not keys:
             return []
         if self.key_prefix:
-            prefixed_keys = [self.key_prefix + key for key in keys]
+            prefixed_keys = [self._get_prefix() + key for key in keys]
         else:
             prefixed_keys = [k for k in keys]
         self._write_client.delete(*prefixed_keys)
         return [k for k in prefixed_keys if not self.has(k)]
 
     def has(self, key: str) -> bool:
-        return bool(self._read_client.exists(self.key_prefix + key))
+        return bool(self._read_client.exists(self._get_prefix() + key))
 
     def clear(self) -> bool:
         status = 0
         if self.key_prefix:
-            keys = self._read_client.keys(self.key_prefix + "*")
+            keys = self._read_client.keys(self._get_prefix() + "*")
             if keys:
                 status = self._write_client.delete(*keys)
         else:
             status = self._write_client.flushdb()
         return bool(status)
 
     def inc(self, key: str, delta: int = 1) -> _t.Any:
-        return self._write_client.incr(name=self.key_prefix + key, amount=delta)
+        return self._write_client.incr(name=self._get_prefix() + key, amount=delta)
 
     def dec(self, key: str, delta: int = 1) -> _t.Any:
-        return self._write_client.incr(name=self.key_prefix + key, amount=-delta)
+        return self._write_client.incr(name=self._get_prefix() + key, amount=-delta)