From ce56a035b5c8aabdbdf5116130f91abf7907e9e3 Mon Sep 17 00:00:00 2001
From: Badiboy <badialex@gmail.com>
Date: Tue, 1 Feb 2022 23:58:57 +0300
Subject: [PATCH] Extend custom exception_handler behaviour

---
 telebot/__init__.py | 12 ++++++++++--
 telebot/util.py     | 12 +++++++++---
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/telebot/__init__.py b/telebot/__init__.py
index 2bc716e..1eb496f 100644
--- a/telebot/__init__.py
+++ b/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()
diff --git a/telebot/util.py b/telebot/util.py
index a10019d..6c09ca3 100644
--- a/telebot/util.py
+++ b/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):