mirror/pyTelegramBotAPI
1
0
Fork 0
mirror of https://github.com/eternnoir/pyTelegramBotAPI.git synced 2023-08-10 21:12:57 +03:00
Code Releases Activity

Updated sync and async. Fixes and new features.

Browse Source
This commit is contained in:
_run
2022-01-24 17:15:04 +04:00
parent 56fbf491bc
commit a3cda2e0ff
18 changed files with 1160 additions and 525 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
telebot/storage/__init__.py Normal file
View File

@ -0,0 +1,13 @@
from telebot.storage.memory_storage import StateMemoryStorage
from telebot.storage.redis_storage import StateRedisStorage
from telebot.storage.pickle_storage import StatePickleStorage
from telebot.storage.base_storage import StateContext,StateStorageBase
__all__ = [
'StateStorageBase', 'StateContext',
'StateMemoryStorage', 'StateRedisStorage', 'StatePickleStorage'
]

65
telebot/storage/base_storage.py Normal file
View File

@ -0,0 +1,65 @@
import copy
class StateStorageBase:
def __init__(self) -> None:
pass
def set_data(self, chat_id, user_id, key, value):
"""
Set data for a user in a particular chat.
"""
raise NotImplementedError
def get_data(self, chat_id, user_id):
"""
Get data for a user in a particular chat.
"""
raise NotImplementedError
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
def delete_state(self, chat_id, user_id):
"""
Delete state for a particular user.
"""
raise NotImplementedError
def reset_data(self, chat_id, user_id):
"""
Reset data for a particular user in a chat.
"""
raise NotImplementedError
def get_state(self, chat_id, user_id):
raise NotImplementedError
def save(chat_id, user_id, data):
raise NotImplementedError
class StateContext:
"""
Class for data.
"""
def __init__(self , obj, chat_id, user_id) -> None:
self.obj = obj
self.data = copy.deepcopy(obj.get_data(chat_id, user_id))
self.chat_id = chat_id
self.user_id = user_id
def __enter__(self):
return self.data
def __exit__(self, exc_type, exc_val, exc_tb):
return self.obj.save(self.chat_id, self.user_id, self.data)

64
telebot/storage/memory_storage.py Normal file
View File

@ -0,0 +1,64 @@
from telebot.storage.base_storage import StateStorageBase, StateContext
class StateMemoryStorage(StateStorageBase):
def __init__(self) -> None:
self.data = {}
#
# {chat_id: {user_id: {'state': None, 'data': {}}, ...}, ...}
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
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
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
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
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
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)
def save(self, chat_id, user_id, data):
self.data[chat_id][user_id]['data'] = data

112
telebot/storage/pickle_storage.py Normal file
View File

@ -0,0 +1,112 @@
from telebot.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()
def convert_old_to_new(self):
"""
Use this function to convert old storage to new storage.
This function is for people who was using pickle storage
that was in version <=4.3.1.
"""
# 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()
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
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
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
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
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
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)
def save(self, chat_id, user_id, data):
self.data[chat_id][user_id]['data'] = data
self.update_data()

176
telebot/storage/redis_storage.py Normal file
View File

@ -0,0 +1,176 @@
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 backend 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 backend 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 backend 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)
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