From 836130a7183fb4b21d4eda00288dd120c02a26b5 Mon Sep 17 00:00:00 2001 From: coder2020official Date: Mon, 2 May 2022 02:08:48 +0500 Subject: [PATCH 1/2] Allow only state objects --- telebot/asyncio_storage/memory_storage.py | 4 +++- telebot/asyncio_storage/pickle_storage.py | 4 ++-- telebot/asyncio_storage/redis_storage.py | 4 +++- telebot/storage/memory_storage.py | 4 +++- telebot/storage/pickle_storage.py | 4 ++-- telebot/storage/redis_storage.py | 3 ++- 6 files changed, 15 insertions(+), 8 deletions(-) diff --git a/telebot/asyncio_storage/memory_storage.py b/telebot/asyncio_storage/memory_storage.py index 58a6e354e..b4d8e16c0 100644 --- a/telebot/asyncio_storage/memory_storage.py +++ b/telebot/asyncio_storage/memory_storage.py @@ -1,5 +1,7 @@ from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext +from telebot.asyncio_handler_backends import State + class StateMemoryStorage(StateStorageBase): def __init__(self) -> None: self.data = {} @@ -8,7 +10,7 @@ def __init__(self) -> None: async def set_state(self, chat_id, user_id, state): - if isinstance(state, object): + if isinstance(state, State): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/asyncio_storage/pickle_storage.py b/telebot/asyncio_storage/pickle_storage.py index dd6419e7e..fab8e81c8 100644 --- a/telebot/asyncio_storage/pickle_storage.py +++ b/telebot/asyncio_storage/pickle_storage.py @@ -1,7 +1,7 @@ from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext import os - +from telebot.asyncio_handler_backends import State import pickle @@ -47,7 +47,7 @@ def update_data(self): file.close() async def set_state(self, chat_id, user_id, state): - if isinstance(state, object): + if isinstance(state, State): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/asyncio_storage/redis_storage.py b/telebot/asyncio_storage/redis_storage.py index f2a26065f..e90c89f6f 100644 --- a/telebot/asyncio_storage/redis_storage.py +++ b/telebot/asyncio_storage/redis_storage.py @@ -1,6 +1,8 @@ from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext import json +from telebot.asyncio_handler_backends import State + redis_installed = True try: import aioredis @@ -65,7 +67,7 @@ async def set_state(self, chat_id, user_id, state): """ response = await self.get_record(chat_id) user_id = str(user_id) - if isinstance(state, object): + if isinstance(state, State): state = state.name if response: if user_id in response: diff --git a/telebot/storage/memory_storage.py b/telebot/storage/memory_storage.py index 45d4da3cc..73f4c4ff1 100644 --- a/telebot/storage/memory_storage.py +++ b/telebot/storage/memory_storage.py @@ -1,5 +1,7 @@ from telebot.storage.base_storage import StateStorageBase, StateContext +from telebot.handler_backends import State + class StateMemoryStorage(StateStorageBase): def __init__(self) -> None: self.data = {} @@ -8,7 +10,7 @@ def __init__(self) -> None: def set_state(self, chat_id, user_id, state): - if isinstance(state, object): + if isinstance(state, State): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/storage/pickle_storage.py b/telebot/storage/pickle_storage.py index a2736903f..0a569072a 100644 --- a/telebot/storage/pickle_storage.py +++ b/telebot/storage/pickle_storage.py @@ -1,6 +1,6 @@ from telebot.storage.base_storage import StateStorageBase, StateContext import os - +from telebot.handler_backends import State import pickle @@ -53,7 +53,7 @@ def update_data(self): file.close() def set_state(self, chat_id, user_id, state): - if isinstance(state, object): + if isinstance(state, State): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/storage/redis_storage.py b/telebot/storage/redis_storage.py index ff21b6e13..e46c6f617 100644 --- a/telebot/storage/redis_storage.py +++ b/telebot/storage/redis_storage.py @@ -1,6 +1,7 @@ from pyclbr import Class from telebot.storage.base_storage import StateStorageBase, StateContext import json +from telebot.handler_backends import State redis_installed = True try: @@ -65,7 +66,7 @@ def set_state(self, chat_id, user_id, state): """ response = self.get_record(chat_id) user_id = str(user_id) - if isinstance(state, object): + if isinstance(state, State): state = state.name if response: From f9cd0d7e081e4a29e0f19b596581931d96681d55 Mon Sep 17 00:00:00 2001 From: coder2020official Date: Mon, 2 May 2022 14:45:43 +0500 Subject: [PATCH 2/2] Avoid circular import --- telebot/asyncio_storage/memory_storage.py | 4 +--- telebot/asyncio_storage/pickle_storage.py | 3 +-- telebot/asyncio_storage/redis_storage.py | 3 +-- telebot/storage/memory_storage.py | 3 +-- telebot/storage/pickle_storage.py | 3 +-- telebot/storage/redis_storage.py | 3 +-- 6 files changed, 6 insertions(+), 13 deletions(-) diff --git a/telebot/asyncio_storage/memory_storage.py b/telebot/asyncio_storage/memory_storage.py index b4d8e16c0..45c2ad914 100644 --- a/telebot/asyncio_storage/memory_storage.py +++ b/telebot/asyncio_storage/memory_storage.py @@ -1,7 +1,5 @@ from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext -from telebot.asyncio_handler_backends import State - class StateMemoryStorage(StateStorageBase): def __init__(self) -> None: self.data = {} @@ -10,7 +8,7 @@ def __init__(self) -> None: async def set_state(self, chat_id, user_id, state): - if isinstance(state, State): + if hasattr(state, 'name'): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/asyncio_storage/pickle_storage.py b/telebot/asyncio_storage/pickle_storage.py index fab8e81c8..4928d4a1b 100644 --- a/telebot/asyncio_storage/pickle_storage.py +++ b/telebot/asyncio_storage/pickle_storage.py @@ -1,7 +1,6 @@ from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext import os -from telebot.asyncio_handler_backends import State import pickle @@ -47,7 +46,7 @@ def update_data(self): file.close() async def set_state(self, chat_id, user_id, state): - if isinstance(state, State): + if hasattr(state, 'name'): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/asyncio_storage/redis_storage.py b/telebot/asyncio_storage/redis_storage.py index e90c89f6f..0cf52b90f 100644 --- a/telebot/asyncio_storage/redis_storage.py +++ b/telebot/asyncio_storage/redis_storage.py @@ -1,7 +1,6 @@ from telebot.asyncio_storage.base_storage import StateStorageBase, StateContext import json -from telebot.asyncio_handler_backends import State redis_installed = True try: @@ -67,7 +66,7 @@ async def set_state(self, chat_id, user_id, state): """ response = await self.get_record(chat_id) user_id = str(user_id) - if isinstance(state, State): + if hasattr(state, 'name'): state = state.name if response: if user_id in response: diff --git a/telebot/storage/memory_storage.py b/telebot/storage/memory_storage.py index 73f4c4ff1..67cd984b5 100644 --- a/telebot/storage/memory_storage.py +++ b/telebot/storage/memory_storage.py @@ -1,6 +1,5 @@ from telebot.storage.base_storage import StateStorageBase, StateContext -from telebot.handler_backends import State class StateMemoryStorage(StateStorageBase): def __init__(self) -> None: @@ -10,7 +9,7 @@ def __init__(self) -> None: def set_state(self, chat_id, user_id, state): - if isinstance(state, State): + if hasattr(state, 'name'): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/storage/pickle_storage.py b/telebot/storage/pickle_storage.py index 0a569072a..39b10a3aa 100644 --- a/telebot/storage/pickle_storage.py +++ b/telebot/storage/pickle_storage.py @@ -1,6 +1,5 @@ from telebot.storage.base_storage import StateStorageBase, StateContext import os -from telebot.handler_backends import State import pickle @@ -53,7 +52,7 @@ def update_data(self): file.close() def set_state(self, chat_id, user_id, state): - if isinstance(state, State): + if hasattr(state, 'name'): state = state.name if chat_id in self.data: if user_id in self.data[chat_id]: diff --git a/telebot/storage/redis_storage.py b/telebot/storage/redis_storage.py index e46c6f617..23a1f4305 100644 --- a/telebot/storage/redis_storage.py +++ b/telebot/storage/redis_storage.py @@ -1,7 +1,6 @@ from pyclbr import Class from telebot.storage.base_storage import StateStorageBase, StateContext import json -from telebot.handler_backends import State redis_installed = True try: @@ -66,7 +65,7 @@ def set_state(self, chat_id, user_id, state): """ response = self.get_record(chat_id) user_id = str(user_id) - if isinstance(state, State): + if hasattr(state, 'name'): state = state.name if response: