Merge pull request #109 from kayqueGovetri/FIX/atexit-chrome-edge

FIX: Adjust atexit in edge and chrome

Author: kayqueGovetri

botcity/web/bot.py

@@ -49,6 +49,10 @@ def _cleanup(bot: ReferenceType[WebBot]):
     if bot() is not None:
         try:
             bot().stop_browser()
+            temp_dir = bot()._botcity_temp_dir
+            if not temp_dir:
+                return None
+            shutil.rmtree(temp_dir, ignore_errors=True)
         except Exception:
             pass
 
@@ -91,6 +95,7 @@ def __init__(self, headless=False):
         self._shift_hold = False
 
         self._download_folder_path = os.getcwd()
+        self._botcity_temp_dir = None
 
         atexit.register(_cleanup, ref(self))
 
@@ -367,6 +372,11 @@ def stop_browser(self):
             self.activate_tab(self.get_tabs()[-1])
         self._driver.close()
         self._driver.quit()
+        if self.options is not None:
+            try:
+                self._botcity_temp_dir = self.options._botcity_temp_dir
+            except Exception:
+                self._botcity_temp_dir = None
         self.options = None
         self.capabilities = None
         self._driver = None

botcity/web/browsers/chrome.py

@@ -1,4 +1,3 @@
-import atexit
 import json
 import os
 import tempfile
@@ -8,7 +7,6 @@
 from selenium.webdriver.chrome.options import Options as ChromeOptions
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
 from selenium.webdriver.chrome.service import Service as ChromeService  # noqa: F401, F403
-from ..util import cleanup_temp_dir
 
 
 def default_options(headless=False, download_folder_path=None, user_data_dir=None,
@@ -72,11 +70,11 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
             chrome_options.add_argument("--no-sandbox")
     except AttributeError:
         pass
-
+    chrome_options._botcity_temp_dir = None
     if not user_data_dir:
         temp_dir = tempfile.TemporaryDirectory(prefix="botcity_")
         user_data_dir = temp_dir.name
-        atexit.register(cleanup_temp_dir, temp_dir)
+        chrome_options._botcity_temp_dir = user_data_dir
 
     chrome_options.add_argument(f"--user-data-dir={user_data_dir}")
 

botcity/web/browsers/edge.py

@@ -1,4 +1,3 @@
-import atexit
 import json
 import os
 import tempfile
@@ -8,7 +7,6 @@
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
 from selenium.webdriver import Edge, EdgeOptions  # noqa: F401, F403
 from selenium.webdriver.edge.service import Service as EdgeService  # noqa: F401, F403
-from ..util import cleanup_temp_dir
 
 
 def default_options(headless=False, download_folder_path=None, user_data_dir=None,
@@ -70,10 +68,11 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     except AttributeError:
         pass
 
+    edge_options._botcity_temp_dir = None
     if not user_data_dir:
         temp_dir = tempfile.TemporaryDirectory(prefix="botcity_")
         user_data_dir = temp_dir.name
-        atexit.register(cleanup_temp_dir, temp_dir)
+        edge_options._botcity_temp_dir = user_data_dir
 
     edge_options.add_argument(f"--user-data-dir={user_data_dir}")
 

botcity/web/browsers/firefox.py

@@ -1,4 +1,3 @@
-import atexit
 import os
 import tempfile
 from typing import Dict
@@ -8,7 +7,6 @@
 from selenium.webdriver.firefox.options import Options as FirefoxOptions
 from selenium.webdriver.firefox.service import Service as FirefoxService  # noqa: F401, F403
 
-from ..util import cleanup_temp_dir
 
 FIREFOX_MIMETYPES_TO_DOWNLOAD = ['application/vnd.hzn-3d-crossword', 'video/3gpp', 'video/3gpp2',
                                  'application/vnd.mseq', 'application/vnd.3m.post-it-notes',
@@ -359,10 +357,11 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     firefox_options.page_load_strategy = page_load_strategy
     if headless:
         firefox_options.add_argument('-headless')
+    firefox_options._botcity_temp_dir = None
     if not user_data_dir:
         temp_dir = tempfile.TemporaryDirectory(prefix="botcity_")
         user_data_dir = temp_dir.name
-        atexit.register(cleanup_temp_dir, temp_dir)
+        firefox_options._botcity_temp_dir = user_data_dir
     if binary_path:
         firefox_options.binary_location = str(binary_path)
     firefox_options.set_preference("profile", user_data_dir)

botcity/web/browsers/undetected_chrome.py

@@ -1,4 +1,3 @@
-import atexit
 import json
 import os
 import platform
@@ -8,7 +7,6 @@
 from undetected_chromedriver import Chrome  # noqa: F401, F403
 from undetected_chromedriver.options import ChromeOptions
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
-from ..util import cleanup_temp_dir
 
 try:
     from undetected_chromedriver import Service as ChromeService  # noqa: F401, F403
@@ -74,10 +72,11 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     except AttributeError:
         pass
 
+    chrome_options._botcity_temp_dir = None
     if not user_data_dir:
         temp_dir = tempfile.TemporaryDirectory(prefix="botcity_")
         user_data_dir = temp_dir.name
-        atexit.register(cleanup_temp_dir, temp_dir)
+        chrome_options._botcity_temp_dir = user_data_dir
 
     chrome_options.add_argument(f"--user-data-dir={user_data_dir}")