如何实现Python函数根据调用场景执行不同Selenium操作？

我有一个用于Selenium自动化的my_function函数，以及两个预定义列表locationList和locationList2，需要分别通过for循环调用该函数。函数代码如下：

def my_function(pickup_d_str, return_d_str, loc, loc2):
        driver.find_element(By.XPATH,'//button[@aria-label="Edit pick-up & return"]').click()
        driver.find_element(By.XPATH,'//div[@class="chicklet location-chicklet-clear"]').click()
        driver.find_element(By.XPATH,'//input[@id="pickupLocationTextBox"]').send_keys(loc)
        driver.find_element(By.XPATH,'//li[@id="location-'+ loc2 +'"]').click()
        driver.find_element(By.XPATH,'//button[@id="continueButton"]').click()
        # 后续还有20行代码

遍历locationList的循环：

for l in locationList:
    my_function(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)

遍历locationList2时，需要在调用my_function后额外执行以下代码（因为网页使用另一搜索表单）：

for l in locationList2:
    my_function(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)
    # 额外执行的代码
    WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//button[@id="check-availability"]'))).click()
    driver.find_element(By.XPATH,'//div[@id="select-item-0"]//button[@aria-label="Select"]').click()
    WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//div[@id="search"]'))).click()

之前尝试用try-except添加这些操作，但导致遍历locationList的循环额外增加30秒执行时间，请问如何高效解决该问题？

解决方案 #

1. 给函数添加分支控制参数 #

直接在my_function中新增一个布尔参数，通过参数值明确控制是否执行额外代码，完全避免try-except带来的等待损耗：

def my_function(pickup_d_str, return_d_str, loc, loc2, is_alternative_form=False):
        driver.find_element(By.XPATH,'//button[@aria-label="Edit pick-up & return"]').click()
        driver.find_element(By.XPATH,'//div[@class="chicklet location-chicklet-clear"]').click()
        driver.find_element(By.XPATH,'//input[@id="pickupLocationTextBox"]').send_keys(loc)
        driver.find_element(By.XPATH,'//li[@id="location-'+ loc2 +'"]').click()
        driver.find_element(By.XPATH,'//button[@id="continueButton"]').click()
        # 原有的20行代码
        
        # 根据参数判断是否执行额外操作
        if is_alternative_form:
            WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//button[@id="check-availability"]'))).click()
            driver.find_element(By.XPATH,'//div[@id="select-item-0"]//button[@aria-label="Select"]').click()
            WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//div[@id="search"]'))).click()

调用时分别传入对应参数：

# 遍历locationList，无需额外操作
for l in locationList:
    my_function(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)

# 遍历locationList2，触发额外操作
for l in locationList2:
    my_function(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby, is_alternative_form=True)

2. 拆分函数复用公共逻辑 #

如果额外代码逻辑独立，可以把公共逻辑抽成基础函数，再编写两个包装函数对应不同场景：

def base_function(pickup_d_str, return_d_str, loc, loc2):
        driver.find_element(By.XPATH,'//button[@aria-label="Edit pick-up & return"]').click()
        driver.find_element(By.XPATH,'//div[@class="chicklet location-chicklet-clear"]').click()
        driver.find_element(By.XPATH,'//input[@id="pickupLocationTextBox"]').send_keys(loc)
        driver.find_element(By.XPATH,'//li[@id="location-'+ loc2 +'"]').click()
        driver.find_element(By.XPATH,'//button[@id="continueButton"]').click()
        # 原有的20行代码

def my_function_normal(pickup_d_str, return_d_str, loc, loc2):
    base_function(pickup_d_str, return_d_str, loc, loc2)

def my_function_alternative(pickup_d_str, return_d_str, loc, loc2):
    base_function(pickup_d_str, return_d_str, loc, loc2)
    # 额外执行的代码
    WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//button[@id="check-availability"]'))).click()
    driver.find_element(By.XPATH,'//div[@id="select-item-0"]//button[@aria-label="Select"]').click()
    WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//div[@id="search"]'))).click()

调用时使用对应函数：

for l in locationList:
    my_function_normal(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)

for l in locationList2:
    my_function_alternative(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)

3. 循环内直接处理分支 #

保持原函数不变，在遍历locationList2的循环中直接追加额外代码，逻辑最直观：

# 遍历locationList
for l in locationList:
    my_function(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)

# 遍历locationList2
for l in locationList2:
    my_function(pickup_d_str, return_d_str, l.loc_pri, l.loc_nearby)
    # 直接执行额外代码
    WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//button[@id="check-availability"]'))).click()
    driver.find_element(By.XPATH,'//div[@id="select-item-0"]//button[@aria-label="Select"]').click()
    WebDriverWait(driver, 20).until( EC.visibility_of_element_located((By.XPATH, '//div[@id="search"]'))).click()

内容的提问来源于stack exchange，提问作者Singh