Merge pull request #110 from kayqueGovetri/FIX/atexit-chrome-base

FIX: Remove atexit in browsers base chrome and firefox

Author: kayqueGovetri

botcity/web/bot.py

@@ -47,6 +47,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
 
@@ -88,6 +92,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))
 
@@ -291,6 +296,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,8 +7,6 @@
 from selenium.webdriver.chrome.options import Options as ChromeOptions
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
 
-from ..util import cleanup_temp_dir
-
 
 def default_options(headless=False, download_folder_path=None, user_data_dir=None,
                     page_load_strategy="normal") -> ChromeOptions:
@@ -70,10 +67,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}")
 

botcity/web/browsers/edge.py

@@ -1,4 +1,3 @@
-import atexit
 import json
 import os
 import tempfile
@@ -8,8 +7,6 @@
 from msedge.selenium_tools import Edge, EdgeOptions  # noqa: F401, F403
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
 
-from ..util import cleanup_temp_dir
-
 
 def default_options(headless=False, download_folder_path=None, user_data_dir=None,
                     page_load_strategy="normal") -> EdgeOptions:
@@ -66,10 +63,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.common.desired_capabilities import DesiredCapabilities
 from selenium.webdriver.firefox.options import Options as FirefoxOptions
 
-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',
@@ -358,10 +356,11 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
         page_load_strategy = page_load_strategy
     firefox_options.page_load_strategy = page_load_strategy
     firefox_options.headless = 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
     firefox_profile = webdriver.FirefoxProfile(user_data_dir)
     firefox_profile.set_preference("security.default_personal_cert", "Select Automatically")
     firefox_profile.set_preference('browser.download.folderList', 2)