1

I'm trying to subclass the Chrome WebDriver to include some initialization and cleanup code, but then Python complains that the created object is set to None:

import glob
import selenium
import subprocess
from selenium.webdriver.common.by import By


class WebDriver(selenium.webdriver.Chrome):
    def __init__(self, url, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.url = url

    def __enter__(self):
        self.get(self.url)
        self.implicitly_wait(15)

    def __exit__(self, type, value, traceback):
        self.quit()
        for path in glob.glob('/tmp/.org.chromium.Chromium.*'):
            subprocess.run(['rm', '-rf', path], check=True)


with WebDriver('https://google.com') as driver:
    driver.find_element(By.ID, 'lst-ib').send_keys('Search')

Running the code with Python 3:

$ python3 test.py
Traceback (most recent call last):
  File "test.py", line 43, in <module>
    driver.find_element(By.ID, 'lst-ib').send_keys('Search')
AttributeError: 'NoneType' object has no attribute 'find_element'
admirabilis
  • 2,290
  • 2
  • 18
  • 33

1 Answers1

2

Your __enter__() magic method should return self for the driver variable to be pointed to the instance of the WebDriver class:

def __enter__(self):
    self.get(self.url)
    self.implicitly_wait(15)
    return self

To get more information about why and how this works, please see:

Community
  • 1
  • 1
alecxe
  • 462,703
  • 120
  • 1,088
  • 1,195