Как использовать набор текста с возвратом заводского метода

Question

У меня проблема с правильными комментариями (или я бы хотел, чтобы это было слишком просто). У меня есть фабричный метод:

def make_driver(driver: str, starting_page: str, headless: bool) -> Union[FirefoxDriver, ChromeDriver]:
    drivers = {'firefox': FirefoxDriver,
               'chrome': ChromeDriver}
    try:
        return drivers[driver](starting_page, headless)
    except KeyError:
        print(f"Driver [{driver}] is not supported")

и позже у меня есть такая функция:

def get_email_input(driver: ???) -> WebElement:
    return driver.find_element_by_css_selector("#email")

Проблема в том, что я не хочу просто повторять Union[FirefoxDriver, ChromeDriver] все время поэтому PyCharm знал бы, что это какой-то WebDriver и дает правильное приглашение с помощью методов этого класса.
Я мог бы сделать что-то вроде MyWebDriver = Union[all, supported, drivers] в каком-то файле настроек и импортировать его, но я не знаю, является ли он лучшим
Я попытался поменять функцию в классе с помощью __new__:

class Driver:
    def __new__(cls, *args, **kwargs) -> Union[FirefoxDriver, ChromeDriver]:
        try/except as in func above

, затем я попытался get_email_input(driver: Driver):, но это не работает - PyCharm считает, что это "чистый" класс Driver, не то, что возвращает __new__

Кто-нибудь имеет какое-то представление об этом?