diff --git a/fsm_telebot/storage/redis.py b/fsm_telebot/storage/redis.py new file mode 100644 index 0000000..45e9b24 --- /dev/null +++ b/fsm_telebot/storage/redis.py @@ -0,0 +1,90 @@ +# -*- coding:utf-8; -*- +import json +import typing + +from redis import Redis + +from .base import BaseStorage + + +class RedisStorage(BaseStorage): + """ + Storage based on Redis + """ + + def __init__(self, + host: typing.Optional[typing.AnyStr] = 'localhost', + port: typing.Optional[int] = 6379, + db: typing.Optional[int] = 0, + password: typing.Optional[typing.AnyStr] = None, + **kwargs): + self._redis = Redis(host=host, port=port, db=db, password=password, **kwargs) + + def _get_record(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None) -> typing.Dict: + """ + Get record from storage + :param chat: + :param user: + :return: + """ + chat, user = self.check_address(chat=chat, user=user) + addr = f"fsm:{chat}:{user}" + + data = self._redis.get(addr) + if data is None: + return {'state': None, 'data': {}} + return json.loads(data) + + def _set_record(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, + state=None, data=None, bucket=None) -> typing.Dict: + """ + Write record to storage + :param bucket: + :param chat: + :param user: + :param state: + :param data: + :return: + """ + if data is None: + data = {} + if bucket is None: + bucket = {} + + chat, user = self.check_address(chat=chat, user=user) + addr = f"fsm:{chat}:{user}" + + record = {'state': state, 'data': data, 'bucket': bucket} + + self._redis.set(addr, json.dumps(record)) + + def get_state(self, chat: typing.Union[int, str, None] = None, user: typing.Union[int, str, None] = None, + default: typing.Optional[str] = None) -> typing.Union[str]: + record = self._get_record(chat=chat, user=user) + return record['state'] or default + + def get_data(self, chat: typing.Union[int, str, None] = None, user: typing.Union[int, str, None] = None, + default: typing.Optional[str] = None) -> typing.Dict: + record = self._get_record(chat=chat, user=user) + return record['data'] or default + + def set_state(self, chat: typing.Union[int, str, None] = None, user: typing.Union[int, str, None] = None, + state: typing.Optional[typing.AnyStr] = None): + record = self._get_record(chat=chat, user=user) + self._set_record(chat=chat, user=user, state=state, data=record['data']) + + def set_data(self, chat: typing.Union[int, str, None] = None, user: typing.Union[int, str, None] = None, + data: typing.Dict = None): + record = self._get_record(chat=chat, user=user) + self._set_record(chat=chat, user=user, state=record['state'], data=data) + + def update_data(self, chat: typing.Union[int, str, None] = None, user: typing.Union[int, str, None] = None, + data: typing.Dict = None, **kwargs): + if data is None: + data = {} + record = self._get_record(chat=chat, user=user) + record_data = record.get('data', {}) + record_data.update(data, **kwargs) + self._set_record(chat=chat, user=user, state=record['state'], data=record_data)