Parameter type definition in function of attribute Python

Question

i am working as a TEA, i use Python + Selenium Webdriver and Behave to run my tests.

For this i define the webdriver inside the context as follows

    driver = webdriver.Chrome()
    driver.implicitly_wait(5)
    context.driver = driver
    context.driver.get('https://www.ebay.com/')

This is defined in one function, but my problem is that when I want to use the driver in another step function i don't get the snippets for the the functions of driver. This makes me lose a lot of time looking what's the exact name of the function I am trying to use.

For example

@then('Click the "Search" button')
def step_impl(context: behave.runner.Context):
    """
    :type context: behave.runner.Context
    """
    search_btn = context.driver.find_element_by_xpath('//*[@id="gh-btn"]')
    search_btn.click()

I know I can define the type of the parameter received directly like

var_name: str (def step_impl(context: behave.runner.Context) , but the case I am looking for is more specific. like context.driver: 'webdriver object'

Is there a way to define the context.driver type so my ide will interpret that this is a webdriver object and in that way get the methods of the driver?

I am using PyCharm Professional

Answer 1

PyCharm should be able to work with this:

@then('Click the "Search" button')
def step_impl(context: behave.runner.Context):
    """
    :type context: behave.runner.Context
    """
    driver: webdriver.Chrome = context.driver
    search_btn = driver.find_element_by_xpath('//*[@id="gh-btn"]')
    search_btn.click()

Answer 2

def step_impl(context: WebDriver):
    search_btn = context.driver.find_element_by_xpath('//*[@id="gh-btn"]')
    search_btn.click()

you don't have to define type two times.

The import you need:

from selenium.webdriver.firefox.webdriver import WebDriver