Merge pull request #1434 from coder2020official/master

Fix States

telebot/__init__.py

@@ -214,7 +214,7 @@ class TeleBot:
 
         self.threaded = threaded
         if self.threaded:
-            self.worker_pool = util.ThreadPool(num_threads=num_threads)
+            self.worker_pool = util.ThreadPool(self, num_threads=num_threads)
     
     @property
     def user(self) -> types.User:
@@ -781,7 +781,15 @@ class TeleBot:
         if self.threaded:
             self.worker_pool.put(task, *args, **kwargs)
         else:
-            task(*args, **kwargs)
+            try:
+                task(*args, **kwargs)
+            except Exception as e:
+                if self.exception_handler is not None:
+                    handled = self.exception_handler.handle(e)
+                else:
+                    handled = False
+                if not handled:
+                    raise e
 
     def stop_polling(self):
         self.__stop_polling.set()

telebot/asyncio_filters.py

@@ -160,6 +160,13 @@ class StateFilter(AdvancedCustomFilter):
 
     async def check(self, message, text):
         if text == '*': return True
+
+        if isinstance(text, list):
+            new_text = [i.name for i in text]
+            text = new_text
+        elif isinstance(text, object):
+            text = text.name
+            
         if message.chat.type == 'group':
             group_state = await self.bot.current_states.get_state(message.chat.id, message.from_user.id)
             if group_state == text:

telebot/asyncio_storage/memory_storage.py

@@ -8,6 +8,8 @@ class StateMemoryStorage(StateStorageBase):
     
     
     async 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/asyncio_storage/pickle_storage.py

@@ -47,6 +47,8 @@ class StatePickleStorage(StateStorageBase):
         file.close()
 
     async 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/asyncio_storage/redis_storage.py

@@ -65,6 +65,8 @@ class StateRedisStorage(StateStorageBase):
         """
         response = await 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

telebot/custom_filters.py

@@ -1,5 +1,4 @@
 from abc import ABC
-
 class SimpleCustomFilter(ABC):
     """
     Simple Custom Filter base class.
@@ -159,6 +158,12 @@ class StateFilter(AdvancedCustomFilter):
 
     def check(self, message, text):
         if text == '*': return True
+
+        if isinstance(text, list):
+            new_text = [i.name for i in text]
+            text = new_text
+        elif isinstance(text, object):
+            text = text.name
         if message.chat.type == 'group':
             group_state = self.bot.current_states.get_state(message.chat.id, message.from_user.id)
             if group_state == text:

telebot/handler_backends.py

@@ -154,6 +154,7 @@ class State:
         self.name = None
     def __str__(self) -> str:
         return self.name
+
     
 
 class StatesGroup:

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

telebot/util.py

@@ -114,7 +114,8 @@ class WorkerThread(threading.Thread):
 
 class ThreadPool:
 
-    def __init__(self, num_threads=2):
+    def __init__(self, telebot, num_threads=2):
+        self.telebot = telebot
         self.tasks = Queue.Queue()
         self.workers = [WorkerThread(self.on_exception, self.tasks) for _ in range(num_threads)]
         self.num_threads = num_threads
@@ -126,8 +127,13 @@ class ThreadPool:
         self.tasks.put((func, args, kwargs))
 
     def on_exception(self, worker_thread, exc_info):
-        self.exception_info = exc_info
-        self.exception_event.set()
+        if self.telebot.exception_handler is not None:
+            handled = self.telebot.exception_handler.handle(exc_info)
+        else:
+            handled = False
+        if not handled:
+            self.exception_info = exc_info
+            self.exception_event.set()
         worker_thread.continue_event.set()
 
     def raise_exceptions(self):