1
0
mirror of https://github.com/eternnoir/pyTelegramBotAPI.git synced 2023-08-10 21:12:57 +03:00
pyTelegramBotAPI/telebot/storage/redis_storage.py

180 lines
5.4 KiB
Python
Raw Normal View History

from telebot.storage.base_storage import StateStorageBase, StateContext
import json
redis_installed = True
try:
from redis import Redis, ConnectionPool
except:
redis_installed = False
class StateRedisStorage(StateStorageBase):
"""
This class is for Redis storage.
This will work only for states.
To use it, just pass this class to:
TeleBot(storage=StateRedisStorage())
"""
def __init__(self, host='localhost', port=6379, db=0, password=None, prefix='telebot_'):
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': {}}, ...}, ...}
self.prefix = prefix
if not redis_installed:
raise Exception("Redis is not installed. Install it via 'pip install redis'")
def get_record(self, key):
"""
Function to get record from database.
It has nothing to do with states.
Made for backward compatibility
"""
connection = Redis(connection_pool=self.redis)
result = connection.get(self.prefix+str(key))
connection.close()
if result: return json.loads(result)
return
def set_record(self, key, value):
"""
Function to set record to database.
It has nothing to do with states.
Made for backward compatibility
"""
connection = Redis(connection_pool=self.redis)
connection.set(self.prefix+str(key), json.dumps(value))
connection.close()
return True
def delete_record(self, key):
"""
Function to delete record from database.
It has nothing to do with states.
Made for backward compatibility
"""
connection = Redis(connection_pool=self.redis)
connection.delete(self.prefix+str(key))
connection.close()
return True
def set_state(self, chat_id, user_id, state):
"""
Set state for a particular user in a chat.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
2022-05-02 12:45:43 +03:00
if hasattr(state, 'name'):
2022-02-02 13:44:02 +03:00
state = state.name
if response:
if user_id in response:
response[user_id]['state'] = state
else:
response[user_id] = {'state': state, 'data': {}}
else:
response = {user_id: {'state': state, 'data': {}}}
self.set_record(chat_id, response)
return True
def delete_state(self, chat_id, user_id):
"""
Delete state for a particular user in a chat.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
del response[user_id]
if user_id == str(chat_id):
self.delete_record(chat_id)
return True
else: self.set_record(chat_id, response)
return True
return False
def get_value(self, chat_id, user_id, key):
"""
Get value for a data of a user in a chat.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
if key in response[user_id]['data']:
return response[user_id]['data'][key]
return None
def get_state(self, chat_id, user_id):
"""
Get state of a user in a chat.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
return response[user_id]['state']
return None
def get_data(self, chat_id, user_id):
"""
Get data of particular user in a particular chat.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
return response[user_id]['data']
return None
def reset_data(self, chat_id, user_id):
"""
Reset data of a user in a chat.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
response[user_id]['data'] = {}
self.set_record(chat_id, response)
return True
def set_data(self, chat_id, user_id, key, value):
"""
Set data without interactive data.
"""
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
response[user_id]['data'][key] = value
self.set_record(chat_id, response)
return True
return False
def get_interactive_data(self, chat_id, user_id):
"""
Get Data in interactive way.
You can use with() with this function.
"""
return StateContext(self, chat_id, user_id)
def save(self, chat_id, user_id, data):
response = self.get_record(chat_id)
user_id = str(user_id)
if response:
if user_id in response:
response[user_id]['data'] = dict(data, **response[user_id]['data'])
self.set_record(chat_id, response)
return True
2022-05-04 17:55:43 +03:00