一日一技：使用上下文管理器来强制关闭 Chromedriver

Posted 2020-12-21

一日一技：使用上下文管理器来强制关闭 Chromedriver

摄影：产品经理
与产品经理环游世界
当我们使用 Selenium 通过 Chromedriver 启动 Chrome 浏览网页时，可能会由于某些异常情况导致程序崩溃，但 Chromedriver 进程不会退出。

例如，我们编写一段显然有问题的代码：

from selenium.webdriver import Chrome

driver = Chrome(‘./chromedriver‘)
driver.get(‘https://www.kingname.info‘)

1 + ‘a‘ # 这一行代码必定导致程序崩溃

代码报错以后，弹出的 Chrome 窗口不会自动关闭。并且 chromedriver 的进程也不会自动结束，如下图所示：

这就会导致系统中出现越来越多的 chromedriver 进程，从而占用大量的内存。

为了防止这种情况，我们必须想办法，在任何情况下都需要保证退出 chromedriver。

你可能会使用一个超大型的 try ... except...把所有与 selenium 相关的代码都包起来：

from selenium.webdriver import Chrome

driver = Chrome(‘./chromedriver‘)
try:
    driver.get(‘https://www.kingname.info‘)
 #第2行
 #第3行
 #第4行
 #。。。
 #第 n 行
except Exception:
    driver.quit()

当然你也可以把具体的操作步骤放在函数里面，然后 try...except... 把函数包住。但本质上是一样的。

但这种超大型的 try...except...一是会导致程序速度减慢，二是程序出现了其他异常的时候，真正的报错信息无法正常打印出来：

>>> a = {}
>>> try:
>>>    a[‘k‘]
>>> except Exception as e:
>>>    print(e)

当你看到这个没头没尾的‘k‘，你不知道是哪一行有问题，也不知道具体有什么问题。

那么，我们有没有办法，既不使用 try ... except ...，但是又能在程序崩溃的时候自动退出 chromedriver 呢？

这个时候我们就可以使用上下文管理器。

我们先来包装一下 Selenium，实现一个带有上下文管理器的类。创建一个SafeDriver.py文件：

from selenium.webdriver import Chrome

class SafeDriver:
    def __init__(self):
        self.driver = Chrome(‘./chromedriver‘)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.driver:
            self.driver.quit()

然后，我们在另一个程序里面调用它：

from SafeDriver import SafeDriver

safe_driver = SafeDriver()
with safe_driver as driver:
    driver.driver.get(‘https://www.kingname.info‘)
    a = {}
    a[‘k‘]

程序进入到with safe_driver as driver的缩进里面，我们会得到一个driver变量，它可以用来操作浏览器。

我们只需要在缩进里面正常写代码即可。一旦由于某种原因导致缩进里面的代码报错，Python 自动会进入SafeDriver类的exit方法中，执行里面的代码。在这个方法里面，我们就可以关闭 chromedriver。从而保证只要程序异常退出，浏览器一定会被关闭，不会遗留进程。

运行效果如下图所示：

报错信息和出错的行数都能正常打印出来了。

我们来看看如何实现一个包含上下文管理器的程序：

随意定义一个类，里面写好你需要执行的逻辑
增加enter(self)方法，定义进入上下文管理器时返回的内容
增加exit(self, exc_type, exc_val, exc_tb)方法，定义退出上下文管理器时需要执行的代码
需要注意的是，enter和exit需要成对使用，不能单独使用其中一个。

在上面的代码中，enter方法仅仅返回了self，于是，下面两段代码：

safe_driver = SafeDriver()

with safe_driver as instance:
    pass

仅仅从功能上来说，instance 变量与safe_driver变量完全一样，都可以使用safe_driver.driver和instance.driver。

所不同的是，使用with启用上下文管理器以后，在退出缩进的时候会执行exit中的内容。

为了简便起见，我们可以使用with safe_driver.driver as driver，直接拿到对象中的self.driver属性，这样可以直接使用类似于driver.get(‘https://www.kingname.info‘)访问网站，而不是instance.driver.get(‘https://www.kingname.info‘)。少敲几次键盘。