Is it possible to have contextlib.closing() to call an arbitrary cleanup method instead of .close()

Question

I was running into some problems using with statement(PEP 343) in python to automatically manage resource clean up after the context. In particular, with statement always assumes the resource clean up method is .close(). I.E. in the following block of code, browser.close() is automatically getting called when execution runs out of the context, however, browser.close() is not the proper clean up as it only close the current window not the entire browser. What it should have done is to call browser.quit().

with contextlib.closing(webdriver.Firefox()) as browser:
    # do something with browser
# at this point browser.close() has been called.

Unfortunately, contextlib.closing doesn't provide a way to customize the clean-up method name to be called, as evidently seen here:

 def __exit__(self, *exc_info):
      self.thing.close()

However, I do notice that there is an argument exec_info but was not used in that particular method. Anybody knows why?

The bigger question is, as title suggested, if possible, how to have self.thing to call an arbitrary clean up method? If not, what's the best work around? Should I fall back to use try...finally?

Answer 1

Well, it's python, you can make your own closing class, based on contextlib.closing and override __exit__() method:

import contextlib
from selenium import webdriver

class closing(contextlib.closing):
    def __exit__(self, *exc_info):
        self.thing.quit()


with closing(webdriver.Firefox()) as browser:
    browser.get('http://stackoverflow.com')

FYI, there was a proposal to make webdriver a context manager, but it was closed as won't fix since quit() is the right way to close the browser, and it should be called explicitly, see Make webdriver a context manager [python bindings].

Answer 2

In addition to alecxe's answer, which is perfectly valid, you could also do something like this:

webdriver.Firefox.close = webdriver.Firefox.quit

Now the close method exists and is the same as the quit method.

Or perhaps better:

from selenium.webdriver.remote.webdriver import WebDriver as RemoteWebDriver
RemoteWebDriver.close = lambda self: self.quit()

This patches the base class to add a close() method that calls the current instance's quit(), so it will work with all the drivers. (The first just patches the Firefox driver.) The lambda is necessary because the quit() method on the base class is overridden on many (all?) of the drivers, so pointing close to quit on the base class would call the base class's method, not the one on the class actually being used.

Either way, the class is now perfectly suited for use with closing() and the with statement.