web_driver_thread_pool_executor.py

class WebDriverThreadPoolExecutor(ThreadPoolExecutor):
    def __init__(self, *args, driver_class, driver_options=None, logger=LOGGER_NAME, **kwargs):
        self._logger = logging.getLogger(logger)
        self._drivers = SimpleQueue()

        for number in range(kwargs["max_workers"]):
            self._logger.debug(
                f"Spawning driver '{driver_class.__name__}' number {number}/{kwargs['max_workers']}")
            self._drivers.put(driver_class(options=driver_options))

        super().__init__(*args, **kwargs)

    def submit(self, fn, /, *args, **kwargs):
        self._logger.debug(
            f"Submitting task '{fn.__qualname__}' with args {args} and kwargs {kwargs}")

        def _fn(*args, **kwargs):
            driver = self._drivers.get(block=True)
            result = fn(driver, *args, **kwargs)

            if self._shutdown:
                driver.close()

            self._drivers.put(driver)

            return result

        return super().submit(_fn, *args, **kwargs)

    def shutdown(self, *args, **kwargs):
        self._logger.debug("Shutting down thread pool")
        super().shutdown(*args, **kwargs)

        while not self._drivers.empty():
            self._drivers.get().close()