Добавить собственный метод к модулю selenium

Я хочу добавить метод в модуль selenium (без правки исходного кода selenium)

Код метода:

def wait_find(self, delay=3, located_by=None, elem_name=None, find_all=False):
    WebDriverWait(self, delay).until(EC.presence_of_element_located((located_by, elem_name)))
    return self.find_elements(located_by, elem_name) if find_all else self.find_element(located_by, elem_name)

Код модуля полностью (изъян в том, что driver создается необычным путем):

from selenium import webdriver
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC


class webdriver_with_my_class(webdriver.Firefox):

    def __init__(self, firefox_profile=None, firefox_binary=None,
                 timeout=30, capabilities=None, proxy=None,
                 executable_path="geckodriver", options=None,
                 service_log_path="geckodriver.log", firefox_options=None,
                 service_args=None, desired_capabilities=None, log_path=None,
                 keep_alive=True):

        super().__init__(firefox_profile, firefox_binary,
                         timeout, capabilities, proxy,
                         executable_path, options,
                         service_log_path, firefox_options,
                         service_args, desired_capabilities, log_path,
                         keep_alive)

    def wait_find(self, delay=3, located_by=None, elem_name=None, find_all=False):
        WebDriverWait(self, delay).until(EC.presence_of_element_located((located_by, elem_name)))
        return self.find_elements(located_by, elem_name) if find_all else self.find_element(located_by, elem_name)

Создание драйвера текущим способом:

driver = webdriver_with_my_class(executable_path='geckodriver')

Проблема: driver создается без метода Firefox, что свидетельствует о кривом коде

Желаемое создание драйвера (традиционный способ)

driver = webdriver.Firefox(executable_path='geckodriver)

Ответы (1 шт):

Автор решения: Namerek

Поделюсь чем сам пользуюсь:

from typing import Union

from selenium.common.exceptions import TimeoutException
from selenium.webdriver import Firefox, FirefoxOptions, FirefoxProfile
from selenium.webdriver.common.by import By
from selenium.webdriver.remote.webelement import WebElement
from selenium.webdriver.support import expected_conditions
from selenium.webdriver.support.wait import WebDriverWait

from common import storage


class Browser(Firefox):

    def __init__(self, **params):
        self.logged_in = False
        opt = FirefoxOptions()
        opt.headless = params.get('headless', False)

        img_permission = int(params.get('hide_images', False)) + 1

        prof = FirefoxProfile()
        prof.set_preference('permissions.default.image', img_permission)

        self.__close_delay = params.get('close_delay', 0.0)

        settings = dict(
            [
                *filter(
                    lambda x: bool(x[1]),
                    [
                        ('timeout', params.get('timeout')),
                        ('service_log_path', params.get('log_path', storage.log_dir / 'geckodriver.log')),
                        ('options', opt),
                        ('firefox_profile', prof),
                        ('executable_path', params.get('executable_path', 'geckodriver'))
                    ]
                )
            ]
        )

        super().__init__(**settings)
        if params.get('maximize', True):
            self.maximize_window()

    def wait_for_element(self, by: By, ident: str, timeout: int = 5) -> Union[WebElement, None]:
        try:
            WebDriverWait(self, timeout).until(
                expected_conditions.presence_of_all_elements_located(
                    (by, ident)
                )
            )
            return self.find_element(by, ident)
        except TimeoutException:
            return None

    def wait_element_for_click(self, by: By, ident: str, timeout: int = 5):
        try:
            WebDriverWait(self, timeout).until(
                expected_conditions.element_to_be_clickable(
                    (by, ident)
                )
            )
            return self.find_element(by, ident)
        except TimeoutException:
            return None
→ Ссылка