support provied storage class with speciftic client

z44d · Sep 20, 2024 · b9303bc · b9303bc
1 parent d892de4
commit b9303bc
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 29 deletions.
diff --git a/tgram/storage/base.py b/tgram/storage/base.py
@@ -5,17 +5,17 @@
 
 
 class StorageBase(ABC):
-    def __init__(self, bot: "tgram.TgBot", type: "str") -> None:
+    def __init__(self, bot: "tgram.TgBot", type: "str", client) -> None:
         if type == "kvsqlite":
             from kvsqlite import Client
 
-            self.client = Client(
+            self.client = client or Client(
                 "tgram-" + str(bot.me.id), workers=bot.workers, loop=bot.loop
             )
         elif type == "redis":
             import redis.asyncio as redis
 
-            self.client = redis.Redis(decode_responses=True)
+            self.client = client or redis.Redis(decode_responses=True)
         self.bot = bot
 
     @abstractmethod

diff --git a/tgram/storage/kvsqlite_storage.py b/tgram/storage/kvsqlite_storage.py
@@ -5,8 +5,8 @@
 
 
 class KvsqliteStorage(StorageBase):
-    def __init__(self, bot: "tgram.TgBot") -> None:
-        super().__init__(bot, "kvsqlite")
+    def __init__(self, bot: "tgram.TgBot", client=None) -> None:
+        super().__init__(bot, "kvsqlite", client)
 
     async def set(self, key: str, value: Any) -> bool:
         return await self.client.set(key, value)

diff --git a/tgram/storage/redis_storage.py b/tgram/storage/redis_storage.py
@@ -6,8 +6,8 @@
 
 
 class RedisStorage(StorageBase):
-    def __init__(self, bot: "tgram.TgBot") -> None:
-        super().__init__(bot, "redis")
+    def __init__(self, bot: "tgram.TgBot", client=None) -> None:
+        super().__init__(bot, "redis", client)
 
     async def set(self, key: str, value: Any) -> bool:
         return await self.client.hset("tgram-" + str(self.bot.me.id), key, value)

diff --git a/tgram/tgbot.py b/tgram/tgbot.py
@@ -254,7 +254,9 @@ def __init__(
         retry_after: Union[int, bool] = None,
         plugins: Union[Path, str] = None,
         skip_updates: bool = True,
-        storage: Literal["kvsqlite", "redis"] = None,
+        storage: Union[
+            KvsqliteStorage, RedisStorage, Literal["kvsqlite", "redis"]
+        ] = None,
     ) -> None:
         self.bot_token = bot_token
         self.api_url = api_url
@@ -289,30 +291,33 @@ def __init__(
         self._api_url: str = f"{api_url}bot{bot_token}/"
 
         if storage:
-            if storage.lower() == "kvsqlite":
-                try:
-                    __import__("kvsqlite")
-                except ModuleNotFoundError:
-                    raise ValueError(
-                        "Please install kvsqlite module before using storage, see more https://pypi.org/project/Kvsqlite/"
-                    )
+            if isinstance(storage, (KvsqliteStorage, RedisStorage)):
+                self.storage = storage
+            else:
+                if storage.lower() == "kvsqlite":
+                    try:
+                        __import__("kvsqlite")
+                    except ModuleNotFoundError:
+                        raise ValueError(
+                            "Please install kvsqlite module before using storage, see more https://pypi.org/project/Kvsqlite/"
+                        )
+                    else:
+                        self.storage = KvsqliteStorage(self)
+                elif storage.lower() == "redis":
+                    try:
+                        __import__("redis")
+                    except ModuleNotFoundError:
+                        raise ValueError(
+                            "Please install redis module before using storage, see more https://pypi.org/project/redis/"
+                        )
+                    else:
+                        self.storage = RedisStorage(self)
                 else:
-                    self.storage = KvsqliteStorage(self)
-            elif storage.lower() == "redis":
-                try:
-                    __import__("redis")
-                except ModuleNotFoundError:
                     raise ValueError(
-                        "Please install redis module before using storage, see more https://pypi.org/project/redis/"
-                    )
-                else:
-                    self.storage = RedisStorage(self)
-            else:
-                raise ValueError(
-                    "Unsupported storage engine {}, only {} are supported for now.".format(
-                        storage, " ,".join(i for i in ["redis", "kvsqlite"])
+                        "Unsupported storage engine {}, only {} are supported for now.".format(
+                            storage, " ,".join(i for i in ["redis", "kvsqlite"])
+                        )
                     )
-                )
 
     def add_handler(self, handler: "tgram.handlers.Handler", group: int = 0) -> None:
         if handler.type == "all":