【Python】Seleniumをwith構文でまとめる

browser.py

from selenium import webdriver
from selenium.webdriver.firefox.firefox_binary import FirefoxBinary

class Firefox(object):
    def __init__(self, binary=None):
        if binary is None:
          self.driver = webdriver.Firefox()
        else:
          binary = FirefoxBinary(binary)
          self.driver = webdriver.Firefox(firefox_binary=binary)

    def __enter__(self):
        return self.driver

    def __exit__(self, exception_type, exception_value, traceback):
        self.driver.close()