mirror of
https://github.com/eternnoir/pyTelegramBotAPI.git
synced 2023-08-10 21:12:57 +03:00
Updated sync and async. Fixes and new features.
This commit is contained in:
_run
13
telebot/asyncio_storage/__init__.py
Normal file
13
telebot/asyncio_storage/__init__.py
Normal file
@ -0,0 +1,13 @@
|
||||
from telebot.asyncio_storage.memory_storage import StateMemoryStorage
|
||||
from telebot.asyncio_storage.redis_storage import StateRedisStorage
|
||||
from telebot.asyncio_storage.pickle_storage import StatePickleStorage
|
||||
from telebot.asyncio_storage.base_storage import StateContext,StateStorageBase
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
__all__ = [
|
||||
'StateStorageBase', 'StateContext',
|
||||
'StateMemoryStorage', 'StateRedisStorage', 'StatePickleStorage'
|
||||
]
|
||||
69
telebot/asyncio_storage/base_storage.py
Normal file
69
telebot/asyncio_storage/base_storage.py
Normal file
@ -0,0 +1,69 @@
|
||||
import copy
|
||||
|
||||
|
||||
class StateStorageBase:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
async def set_data(self, chat_id, user_id, key, value):
|
||||
"""
|
||||
Set data for a user in a particular chat.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def get_data(self, chat_id, user_id):
|
||||
"""
|
||||
Get data for a user in a particular chat.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def set_state(self, chat_id, user_id, state):
|
||||
"""
|
||||
Set state for a particular user.
|
||||
|
||||
! Note that you should create a
|
||||
record if it does not exist, and
|
||||
if a record with state already exists,
|
||||
you need to update a record.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def delete_state(self, chat_id, user_id):
|
||||
"""
|
||||
Delete state for a particular user.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def reset_data(self, chat_id, user_id):
|
||||
"""
|
||||
Reset data for a particular user in a chat.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def get_state(self, chat_id, user_id):
|
||||
raise NotImplementedError
|
||||
|
||||
async def save(chat_id, user_id, data):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
|
||||
class StateContext:
|
||||
"""
|
||||
Class for data.
|
||||
"""
|
||||
|
||||
def __init__(self, obj, chat_id, user_id):
|
||||
self.obj = obj
|
||||
self.data = None
|
||||
self.chat_id = chat_id
|
||||
self.user_id = user_id
|
||||
|
||||
|
||||
|
||||
async def __aenter__(self):
|
||||
self.data = copy.deepcopy(await self.obj.get_data(self.chat_id, self.user_id))
|
||||
return self.data
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return await self.obj.save(self.chat_id, self.user_id, self.data)
|
||||
64
telebot/asyncio_storage/memory_storage.py
Normal file
64
telebot/asyncio_storage/memory_storage.py
Normal file
@ -0,0 +1,64 @@
|
||||
from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext
|
||||
|
||||
class StateMemoryStorage(StateStorageBase):
|
||||
def __init__(self) -> None:
|
||||
self.data = {}
|
||||
#
|
||||
# {chat_id: {user_id: {'state': None, 'data': {}}, ...}, ...}
|
||||
|
||||
|
||||
async def set_state(self, chat_id, user_id, state):
|
||||
if chat_id in self.data:
|
||||
if user_id in self.data[chat_id]:
|
||||
self.data[chat_id][user_id]['state'] = state
|
||||
return True
|
||||
else:
|
||||
self.data[chat_id][user_id] = {'state': state, 'data': {}}
|
||||
return True
|
||||
self.data[chat_id] = {user_id: {'state': state, 'data': {}}}
|
||||
return True
|
||||
|
||||
async def delete_state(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
del self.data[chat_id][user_id]
|
||||
if chat_id == user_id:
|
||||
del self.data[chat_id]
|
||||
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
async def get_state(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
return self.data[chat_id][user_id]['state']
|
||||
|
||||
return None
|
||||
async def get_data(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
return self.data[chat_id][user_id]['data']
|
||||
|
||||
return None
|
||||
|
||||
async def reset_data(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
self.data[chat_id][user_id]['data'] = {}
|
||||
return True
|
||||
return False
|
||||
|
||||
async def set_data(self, chat_id, user_id, key, value):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
self.data[chat_id][user_id]['data'][key] = value
|
||||
return True
|
||||
raise RuntimeError('chat_id {} and user_id {} does not exist'.format(chat_id, user_id))
|
||||
|
||||
def get_interactive_data(self, chat_id, user_id):
|
||||
return StateContext(self, chat_id, user_id)
|
||||
|
||||
async def save(self, chat_id, user_id, data):
|
||||
self.data[chat_id][user_id]['data'] = data
|
||||
107
telebot/asyncio_storage/pickle_storage.py
Normal file
107
telebot/asyncio_storage/pickle_storage.py
Normal file
@ -0,0 +1,107 @@
|
||||
from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext
|
||||
import os
|
||||
|
||||
|
||||
import pickle
|
||||
|
||||
|
||||
class StatePickleStorage(StateStorageBase):
|
||||
def __init__(self, file_path="./.state-save/states.pkl") -> None:
|
||||
self.file_path = file_path
|
||||
self.create_dir()
|
||||
self.data = self.read()
|
||||
|
||||
async def convert_old_to_new(self):
|
||||
# old looks like:
|
||||
# {1: {'state': 'start', 'data': {'name': 'John'}}
|
||||
# we should update old version pickle to new.
|
||||
# new looks like:
|
||||
# {1: {2: {'state': 'start', 'data': {'name': 'John'}}}}
|
||||
new_data = {}
|
||||
for key, value in self.data.items():
|
||||
# this returns us id and dict with data and state
|
||||
new_data[key] = {key: value} # convert this to new
|
||||
# pass it to global data
|
||||
self.data = new_data
|
||||
self.update_data() # update data in file
|
||||
|
||||
def create_dir(self):
|
||||
"""
|
||||
Create directory .save-handlers.
|
||||
"""
|
||||
dirs = self.file_path.rsplit('/', maxsplit=1)[0]
|
||||
os.makedirs(dirs, exist_ok=True)
|
||||
if not os.path.isfile(self.file_path):
|
||||
with open(self.file_path,'wb') as file:
|
||||
pickle.dump({}, file)
|
||||
|
||||
def read(self):
|
||||
file = open(self.file_path, 'rb')
|
||||
data = pickle.load(file)
|
||||
file.close()
|
||||
return data
|
||||
|
||||
def update_data(self):
|
||||
file = open(self.file_path, 'wb+')
|
||||
pickle.dump(self.data, file, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
file.close()
|
||||
|
||||
async def set_state(self, chat_id, user_id, state):
|
||||
if chat_id in self.data:
|
||||
if user_id in self.data[chat_id]:
|
||||
self.data[chat_id][user_id]['state'] = state
|
||||
return True
|
||||
else:
|
||||
self.data[chat_id][user_id] = {'state': state, 'data': {}}
|
||||
return True
|
||||
self.data[chat_id] = {user_id: {'state': state, 'data': {}}}
|
||||
self.update_data()
|
||||
return True
|
||||
|
||||
async def delete_state(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
del self.data[chat_id][user_id]
|
||||
if chat_id == user_id:
|
||||
del self.data[chat_id]
|
||||
self.update_data()
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
async def get_state(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
return self.data[chat_id][user_id]['state']
|
||||
|
||||
return None
|
||||
async def get_data(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
return self.data[chat_id][user_id]['data']
|
||||
|
||||
return None
|
||||
|
||||
async def reset_data(self, chat_id, user_id):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
self.data[chat_id][user_id]['data'] = {}
|
||||
self.update_data()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def set_data(self, chat_id, user_id, key, value):
|
||||
if self.data.get(chat_id):
|
||||
if self.data[chat_id].get(user_id):
|
||||
self.data[chat_id][user_id]['data'][key] = value
|
||||
self.update_data()
|
||||
return True
|
||||
raise RuntimeError('chat_id {} and user_id {} does not exist'.format(chat_id, user_id))
|
||||
|
||||
def get_interactive_data(self, chat_id, user_id):
|
||||
return StateContext(self, chat_id, user_id)
|
||||
|
||||
async def save(self, chat_id, user_id, data):
|
||||
self.data[chat_id][user_id]['data'] = data
|
||||
self.update_data()
|
||||
178
telebot/asyncio_storage/redis_storage.py
Normal file
178
telebot/asyncio_storage/redis_storage.py
Normal file
@ -0,0 +1,178 @@
|
||||
from pickle import FALSE
|
||||
from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext
|
||||
import json
|
||||
|
||||
redis_installed = True
|
||||
try:
|
||||
import aioredis
|
||||
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_'):
|
||||
if not redis_installed:
|
||||
raise ImportError('AioRedis is not installed. Install it via "pip install aioredis"')
|
||||
|
||||
|
||||
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)
|
||||
|
||||
self.prefix = prefix
|
||||
#self.con = Redis(connection_pool=self.redis) -> use this when necessary
|
||||
#
|
||||
# {chat_id: {user_id: {'state': None, 'data': {}}, ...}, ...}
|
||||
|
||||
async def get_record(self, key):
|
||||
"""
|
||||
Function to get record from database.
|
||||
It has nothing to do with states.
|
||||
Made for backend compatibility
|
||||
"""
|
||||
result = await self.redis.get(self.prefix+str(key))
|
||||
if result: return json.loads(result)
|
||||
return
|
||||
|
||||
async def set_record(self, key, value):
|
||||
"""
|
||||
Function to set record to database.
|
||||
It has nothing to do with states.
|
||||
Made for backend compatibility
|
||||
"""
|
||||
|
||||
await self.redis.set(self.prefix+str(key), json.dumps(value))
|
||||
return True
|
||||
|
||||
async def delete_record(self, key):
|
||||
"""
|
||||
Function to delete record from database.
|
||||
It has nothing to do with states.
|
||||
Made for backend compatibility
|
||||
"""
|
||||
await self.redis.delete(self.prefix+str(key))
|
||||
return True
|
||||
|
||||
async def set_state(self, chat_id, user_id, state):
|
||||
"""
|
||||
Set state for a particular user in a chat.
|
||||
"""
|
||||
response = await self.get_record(chat_id)
|
||||
user_id = str(user_id)
|
||||
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': {}}}
|
||||
await self.set_record(chat_id, response)
|
||||
|
||||
return True
|
||||
|
||||
async def delete_state(self, chat_id, user_id):
|
||||
"""
|
||||
Delete state for a particular user in a chat.
|
||||
"""
|
||||
response = await 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):
|
||||
await self.delete_record(chat_id)
|
||||
return True
|
||||
else: await self.set_record(chat_id, response)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def get_value(self, chat_id, user_id, key):
|
||||
"""
|
||||
Get value for a data of a user in a chat.
|
||||
"""
|
||||
response = await 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
|
||||
|
||||
|
||||
async def get_state(self, chat_id, user_id):
|
||||
"""
|
||||
Get state of a user in a chat.
|
||||
"""
|
||||
response = await self.get_record(chat_id)
|
||||
user_id = str(user_id)
|
||||
if response:
|
||||
if user_id in response:
|
||||
return response[user_id]['state']
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def get_data(self, chat_id, user_id):
|
||||
"""
|
||||
Get data of particular user in a particular chat.
|
||||
"""
|
||||
response = await self.get_record(chat_id)
|
||||
user_id = str(user_id)
|
||||
if response:
|
||||
if user_id in response:
|
||||
return response[user_id]['data']
|
||||
return None
|
||||
|
||||
|
||||
async def reset_data(self, chat_id, user_id):
|
||||
"""
|
||||
Reset data of a user in a chat.
|
||||
"""
|
||||
response = await self.get_record(chat_id)
|
||||
user_id = str(user_id)
|
||||
if response:
|
||||
if user_id in response:
|
||||
response[user_id]['data'] = {}
|
||||
await self.set_record(chat_id, response)
|
||||
return True
|
||||
|
||||
|
||||
|
||||
|
||||
async def set_data(self, chat_id, user_id, key, value):
|
||||
"""
|
||||
Set data without interactive data.
|
||||
"""
|
||||
response = await self.get_record(chat_id)
|
||||
user_id = str(user_id)
|
||||
if response:
|
||||
if user_id in response:
|
||||
response[user_id]['data'][key] = value
|
||||
await 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)
|
||||
|
||||
async def save(self, chat_id, user_id, data):
|
||||
response = await 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'])
|
||||
await self.set_record(chat_id, response)
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user