Allow only state objects

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 @@ class StateMemoryStorage(StateStorageBase):
     
     
     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]:

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 @@ class StatePickleStorage(StateStorageBase):
         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]:

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 @@ class StateRedisStorage(StateStorageBase):
         """
         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:

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 @@ class StateMemoryStorage(StateStorageBase):
     
     
     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]:

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 @@ class StatePickleStorage(StateStorageBase):
         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]:

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 @@ class StateRedisStorage(StateStorageBase):
         """
         response = self.get_record(chat_id)
         user_id = str(user_id)
-        if isinstance(state, object):
+        if isinstance(state, State):
             state = state.name
 
         if response: