FIX: Adjust atexit in edge and chrome

Author: kayqueGovetri

botcity/web/bot.py

@@ -46,9 +46,42 @@
 
 
 def _cleanup(bot: ReferenceType[WebBot]):
+
+    def get_user_data_dir(opt, browser_name):
+        browsers = {
+            Browser.CHROME: ("--user-data-dir=", opt.arguments if opt else []),
+            Browser.EDGE: ("--user-data-dir=", opt.arguments if opt else []),
+            Browser.UNDETECTED_CHROME: ("--user-data-dir=", opt.arguments if opt else []),
+            Browser.FIREFOX: None,
+            Browser.IE: None
+        }
+
+        items = browsers.get(browser_name, None)
+        if not items:
+            return None
+
+        result = next(filter(lambda x: x.startswith(items[0]), items[1]), None)
+        if not result:
+            return None
+
+        match_result = re.search(rf'{items[0]}([^ ]+)', result)
+
+        if not match_result:
+            return None
+
+        return match_result.group(1)
+
     if bot() is not None:
         try:
+            options = bot().options
+            name = bot().browser
             bot().stop_browser()
+            if not bot():
+                return None
+            user_data_dir = get_user_data_dir(options, name)
+            if not user_data_dir or "botcity_" not in user_data_dir:
+                return None
+            shutil.rmtree(user_data_dir, ignore_errors=True)
         except Exception:
             pass
 

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,
@@ -76,7 +74,6 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     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.add_argument(f"--user-data-dir={user_data_dir}")
 

botcity/web/browsers/edge.py

@@ -8,7 +8,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,
@@ -73,7 +72,6 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     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.add_argument(f"--user-data-dir={user_data_dir}")
 

botcity/web/browsers/firefox.py

@@ -8,7 +8,7 @@
 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
+from botcity.web.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',