From 24ec1254a38ee523148f58c053487c210b14346e Mon Sep 17 00:00:00 2001 From: Cub11k Date: Fri, 8 Sep 2023 20:13:34 +0300 Subject: [PATCH] Add redis_url to StateRedisStorage to allow initialization by URL Useful for rediss:// and unix:// schemes --- telebot/asyncio_storage/redis_storage.py | 7 +++++-- telebot/storage/redis_storage.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/telebot/asyncio_storage/redis_storage.py b/telebot/asyncio_storage/redis_storage.py index 8d2519b71..84db253e5 100644 --- a/telebot/asyncio_storage/redis_storage.py +++ b/telebot/asyncio_storage/redis_storage.py @@ -20,7 +20,7 @@ class StateRedisStorage(StateStorageBase): To use it, just pass this class to: TeleBot(storage=StateRedisStorage()) """ - def __init__(self, host='localhost', port=6379, db=0, password=None, prefix='telebot_'): + def __init__(self, host='localhost', port=6379, db=0, password=None, prefix='telebot_', redis_url=None): if not redis_installed: raise ImportError('AioRedis is not installed. Install it via "pip install aioredis"') @@ -28,7 +28,10 @@ def __init__(self, host='localhost', port=6379, db=0, password=None, prefix='tel aioredis_version = tuple(map(int, aioredis.__version__.split(".")[0])) if aioredis_version < (2,): raise ImportError('Invalid aioredis version. Aioredis version should be >= 2.0.0') - self.redis = aioredis.Redis(host=host, port=port, db=db, password=password) + if redis_url: + self.redis = aioredis.Redis.from_url(redis_url) + else: + self.redis = aioredis.Redis(host=host, port=port, db=db, password=password) self.prefix = prefix #self.con = Redis(connection_pool=self.redis) -> use this when necessary diff --git a/telebot/storage/redis_storage.py b/telebot/storage/redis_storage.py index 453d6ae33..3fac57c46 100644 --- a/telebot/storage/redis_storage.py +++ b/telebot/storage/redis_storage.py @@ -15,9 +15,12 @@ class StateRedisStorage(StateStorageBase): To use it, just pass this class to: TeleBot(storage=StateRedisStorage()) """ - def __init__(self, host='localhost', port=6379, db=0, password=None, prefix='telebot_'): + def __init__(self, host='localhost', port=6379, db=0, password=None, prefix='telebot_', redis_url=None): super().__init__() - self.redis = ConnectionPool(host=host, port=port, db=db, password=password) + if redis_url: + self.redis = ConnectionPool.from_url(redis_url) + else: + self.redis = ConnectionPool(host=host, port=port, db=db, password=password) #self.con = Redis(connection_pool=self.redis) -> use this when necessary # # {chat_id: {user_id: {'state': None, 'data': {}}, ...}, ...}