ENH: Add Firefox option for save

Using `options` instead of `chrome_options` follows the deprecation
advice in Selenium.

According to its docs, Firefox takes single dash options, but I verified
that it understands `--headless`.

I tested as root and Firefox does not need `--no-sandbox`, which is
convenient because it does not have such a flag.

Fixes vega#714.

altair/utils/headless.py

@@ -12,11 +12,6 @@
 except ImportError:
     webdriver = None
 
-try:
-    from selenium.webdriver.chrome.options import Options as ChromeOptions
-except ImportError:
-    ChromeOptions = None
-
 
 @contextlib.contextmanager
 def temporary_filename(**kwargs):
@@ -102,7 +97,8 @@ def spec_to_image_mimebundle(spec, format, mode,
                              vega_version,
                              vegaembed_version,
                              vegalite_version=None,
-                             driver_timeout=10):
+                             driver_timeout=10,
+                             webdriver_class=None):
     """Conver a vega/vega-lite specification to a PNG/SVG image
 
     Parameters
@@ -122,6 +118,8 @@ def spec_to_image_mimebundle(spec, format, mode,
     driver_timeout : int (optional)
         the number of seconds to wait for page load before raising an
         error (default: 10)
+    webdriver_class : Type[Union[webdriver.Chrome, webdriver.Firefox]]
+        Webdriver to use (default: webdriver.Chrome).
 
     Returns
     -------
@@ -130,8 +128,8 @@ def spec_to_image_mimebundle(spec, format, mode,
 
     Note
     ----
-    This requires the pillow, selenium, and chrome headless packages to be
-    installed.
+    This requires the pillow, selenium, and chromedriver or geckodriver
+    packages to be installed.
     """
     # TODO: allow package versions to be specified
     # TODO: detect & use local Jupyter caches of JS packages?
@@ -146,23 +144,29 @@ def spec_to_image_mimebundle(spec, format, mode,
     if webdriver is None:
         raise ImportError("selenium package is required "
                           "for saving chart as {0}".format(format))
-    if ChromeOptions is None:
-        raise ImportError("chromedriver is required "
-                          "for saving chart as {0}".format(format))
+    if webdriver_class is None:
+        webdriver_class = webdriver.Chrome
+    if issubclass(webdriver_class, webdriver.Chrome):
+        webdriver_options_class = webdriver.chrome.options.Options
+    elif issubclass(webdriver_class, webdriver.Firefox):
+        webdriver_options_class = webdriver.firefox.options.Options
+    else:
+        raise ValueError("Only Chrome and Firefox webdrivers are supported")
 
     html = HTML_TEMPLATE.format(vega_version=vega_version,
                                 vegalite_version=vegalite_version,
                                 vegaembed_version=vegaembed_version)
 
-    chrome_options = ChromeOptions()
-    chrome_options.add_argument("--headless")
+    webdriver_options = webdriver_options_class()
+    webdriver_options.add_argument("--headless")
 
-    # for linux/osx root user, need to add --no-sandbox option.
-    # since geteuid doesn't exist on windows, we don't check it
-    if hasattr(os, 'geteuid') and (os.geteuid() == 0):
-        chrome_options.add_argument('--no-sandbox')
+    if issubclass(webdriver_class, webdriver.Chrome):
+        # for linux/osx root user, need to add --no-sandbox option.
+        # since geteuid doesn't exist on windows, we don't check it
+        if hasattr(os, 'geteuid') and (os.geteuid() == 0):
+            webdriver_options.add_argument('--no-sandbox')
 
-    driver = webdriver.Chrome(chrome_options=chrome_options)
+    driver = webdriver_class(options=webdriver_options)
 
     try:
         driver.set_page_load_timeout(driver_timeout)

altair/utils/save.py

@@ -9,7 +9,7 @@
 
 def save(chart, fp, vega_version, vegaembed_version,
          format=None, mode=None, vegalite_version=None,
-         opt=None, json_kwds=None):
+         opt=None, json_kwds=None, webdriver_class=None):
     """Save a chart to file in a variety of formats
 
     Supported formats are [json, html, png, svg]
@@ -39,6 +39,8 @@ def save(chart, fp, vega_version, vegaembed_version,
     json_kwds : dict
         Additional keyword arguments are passed to the output method
         associated with the specified format.
+    webdriver_class : Type[Union[webdriver.Chrome, webdriver.Firefox]]
+        Webdriver to use (default: webdriver.Chrome).
     """
     if json_kwds is None:
         json_kwds = {}
@@ -84,7 +86,8 @@ def save(chart, fp, vega_version, vegaembed_version,
         mimebundle = spec_to_image_mimebundle(spec=spec, format=format, mode=mode,
                                               vega_version=vega_version,
                                               vegalite_version=vegalite_version,
-                                              vegaembed_version=vegaembed_version)
+                                              vegaembed_version=vegaembed_version,
+                                              webdriver_class=webdriver_class)
         if format == 'png':
             write_file_or_filename(fp, mimebundle['image/png'], mode='wb')
         else:

doc/user_guide/saving_charts.rst

@@ -22,8 +22,8 @@ However, saving these images requires some additional dependencies to run the
 javascript code necessary to interpret the Vega-Lite specification and output
 it in the form of an image.
 
-Altair is set up to do this conversion using selenium and headless chrome, which
-requires the following:
+Altair is set up to do this conversion using selenium and headless Chrome or
+Firefox, which requires the following:
 
 - the Selenium_ python package. This can be installed using::
 
@@ -33,16 +33,20 @@ requires the following:
 
       $ pip install selenium
 
-- a recent version of `Google Chrome`_. Please see the Chrome installation page
-  for installation details for your own operating system.
+- a recent version of `Google Chrome`_ or `Mozilla Firefox`_. Please see the
+  Chrome or Firefox installation page for installation details for your own
+  operating system.
 
-- `Chrome Driver`_, which allows Google Chrome to be run in a *headless* state
-  (i.e. to execute Javascript code without opening an actual browser window).
-  If you use homebrew on OSX, this can be installed with::
+- `Chrome Driver`_ or `Gecko Driver`_, which allows Chrome or Firefox to be run
+  in a *headless* state (i.e. to execute Javascript code without opening an
+  actual browser window).  If you use homebrew on OSX, this can be installed
+  with::
 
       $ brew install chromedriver
+      $ brew install geckodriver
 
-  See the ``chromedriver`` documentation for details on installation.
+  See the ``chromedriver`` or ``geckodriver`` documentation for details on
+  installation.
 
 Once those dependencies are installed, you should be able to save charts as
 ``png`` or ``svg``.
@@ -171,5 +175,7 @@ You can view the result here: `chart.html </_static/chart.html>`_.
 
 .. _Selenium: http://selenium-python.readthedocs.io/
 .. _Google Chrome: https://www.google.com/chrome/
+.. _Mozilla Firefox: https://www.mozilla.org/firefox/
 .. _Chrome Driver: https://sites.google.com/a/chromium.org/chromedriver/
+.. _Gecko Driver: https://github.com/mozilla/geckodriver/releases
 .. _vegaEmbed: https://github.com/vega/vega-embed