diff --git a/botcity/web/bot.py b/botcity/web/bot.py index 0ddbf4a..9b018d5 100644 --- a/botcity/web/bot.py +++ b/botcity/web/bot.py @@ -1,3 +1,6 @@ +from __future__ import annotations + +import atexit import base64 import functools import glob @@ -27,6 +30,7 @@ from selenium.webdriver.support.wait import WebDriverWait, TimeoutException, NoSuchElementException from selenium.webdriver.support import expected_conditions as EC from selenium.webdriver.common.print_page_options import PrintOptions +from weakref import ReferenceType, ref from . import config, cv2find, compat from .browsers import BROWSER_CONFIGS, Browser, PageLoadStrategy @@ -41,6 +45,14 @@ logger = logging.getLogger(__name__) +def _cleanup(bot: ReferenceType[WebBot]): + if bot() is not None: + try: + bot().stop_browser() + except Exception: + pass + + class WebBot(BaseBot): KEYS = Keys DEFAULT_DIMENSIONS = (1600, 900) @@ -80,6 +92,8 @@ def __init__(self, headless=False): self._download_folder_path = os.getcwd() + atexit.register(_cleanup, ref(self)) + def __enter__(self): pass