Address code review comments

altair/utils/headless.py

@@ -8,9 +8,9 @@
 import tempfile
 
 try:
-    from selenium import webdriver
+    import selenium.webdriver
 except ImportError:
-    webdriver = None
+    selenium = None
 
 
 @contextlib.contextmanager
@@ -98,7 +98,7 @@ def spec_to_image_mimebundle(spec, format, mode,
                              vegaembed_version,
                              vegalite_version=None,
                              driver_timeout=10,
-                             webdriver_class=None):
+                             webdriver='chrome'):
     """Conver a vega/vega-lite specification to a PNG/SVG image
 
     Parameters
@@ -118,8 +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).
+    webdriver : string {'chrome' | 'firefox'}
+        Webdriver to use.
 
     Returns
     -------
@@ -128,8 +128,9 @@ def spec_to_image_mimebundle(spec, format, mode,
 
     Note
     ----
-    This requires the pillow, selenium, and chromedriver or geckodriver
-    packages to be installed.
+    This requires the pillow and selenium Python modules to be installed.
+    Additionally it requires either chromedriver (if webdriver=='chrome') or
+    geckodriver (if webdriver=='firefox')
     """
     # TODO: allow package versions to be specified
     # TODO: detect & use local Jupyter caches of JS packages?
@@ -141,17 +142,17 @@ def spec_to_image_mimebundle(spec, format, mode,
     if mode == 'vega-lite' and vegalite_version is None:
         raise ValueError("must specify vega-lite version")
 
-    if webdriver is None:
+    if selenium is None:
         raise ImportError("selenium package 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
+    if webdriver == 'chrome':
+        webdriver_class = selenium.webdriver.Chrome
+        webdriver_options_class = selenium.webdriver.chrome.options.Options
+    elif webdriver == 'firefox':
+        webdriver_class = selenium.webdriver.Firefox
+        webdriver_options_class = selenium.webdriver.firefox.options.Options
     else:
-        raise ValueError("Only Chrome and Firefox webdrivers are supported")
+        raise ValueError("webdriver can only be 'chrome' or 'firefox'")
 
     html = HTML_TEMPLATE.format(vega_version=vega_version,
                                 vegalite_version=vegalite_version,
@@ -160,7 +161,7 @@ def spec_to_image_mimebundle(spec, format, mode,
     webdriver_options = webdriver_options_class()
     webdriver_options.add_argument("--headless")
 
-    if issubclass(webdriver_class, webdriver.Chrome):
+    if issubclass(webdriver_class, selenium.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):

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, webdriver_class=None):
+         opt=None, json_kwds=None, webdriver='chrome'):
     """Save a chart to file in a variety of formats
 
     Supported formats are [json, html, png, svg]
@@ -39,8 +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).
+    webdriver : string {'chrome' | 'firefox'}
+        Webdriver to use.
     """
     if json_kwds is None:
         json_kwds = {}
@@ -87,7 +87,7 @@ def save(chart, fp, vega_version, vegaembed_version,
                                               vega_version=vega_version,
                                               vegalite_version=vegalite_version,
                                               vegaembed_version=vegaembed_version,
-                                              webdriver_class=webdriver_class)
+                                              webdriver=webdriver)
         if format == 'png':
             write_file_or_filename(fp, mimebundle['image/png'], mode='wb')
         else: