Fix States

telebot/storage/memory_storage.py

@@ -8,6 +8,8 @@ class StateMemoryStorage(StateStorageBase):
     
     
     def set_state(self, chat_id, user_id, state):
+        if isinstance(state, object):
+            state = state.name
         if chat_id in self.data:
             if user_id in self.data[chat_id]:
                 self.data[chat_id][user_id]['state'] = state
@@ -31,6 +33,7 @@ class StateMemoryStorage(StateStorageBase):
 
     
     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']

telebot/storage/pickle_storage.py

@@ -53,6 +53,8 @@ class StatePickleStorage(StateStorageBase):
         file.close()
 
     def set_state(self, chat_id, user_id, state):
+        if isinstance(state, object):
+            state = state.name
         if chat_id in self.data:
             if user_id in self.data[chat_id]:
                 self.data[chat_id][user_id]['state'] = state

telebot/storage/redis_storage.py

@@ -1,3 +1,4 @@
+from pyclbr import Class
 from telebot.storage.base_storage import StateStorageBase, StateContext
 import json
 
@@ -64,6 +65,9 @@ class StateRedisStorage(StateStorageBase):
         """
         response = self.get_record(chat_id)
         user_id = str(user_id)
+        if isinstance(state, object):
+            state = state.name
+
         if response:
             if user_id in response:
                 response[user_id]['state'] = state