Repository: bokeh/bokeh

Issue:
Task: expose webdriver as kwarg to export_png and export_svgs 
Task: expose webdriver as kwarg to export_png and export_svgs

This would allow users to manually share a webdriver session among multiple calls to the export functions.


Discussion:

Resolving patch:
diff --git a/bokeh/io.py b/bokeh/io.py
--- a/bokeh/io.py
+++ b/bokeh/io.py
@@ -624,7 +624,7 @@ def _crop_image(image, left=0, top=0, right=0, bottom=0, **kwargs):
 
     return cropped_image
 
-def _get_screenshot_as_png(obj, **kwargs):
+def _get_screenshot_as_png(obj, driver=None, **kwargs):
     webdriver = import_required('selenium.webdriver',
                                 'To use bokeh.io.export_png you need selenium ' +
                                 '("conda install -c bokeh selenium" or "pip install selenium")')
@@ -637,7 +637,7 @@ def _get_screenshot_as_png(obj, **kwargs):
 
     html_path = _save_layout_html(obj, **kwargs)
 
-    web_driver = kwargs.get('driver', webdriver.PhantomJS(service_log_path=os.path.devnull))
+    web_driver = driver if driver is not None else webdriver.PhantomJS(service_log_path=os.path.devnull)
 
     web_driver.get("file:///" + html_path)
     web_driver.maximize_window()
@@ -652,7 +652,7 @@ def _get_screenshot_as_png(obj, **kwargs):
     bounding_rect_script = "return document.getElementsByClassName('bk-root')[0].children[0].getBoundingClientRect()"
     b_rect = web_driver.execute_script(bounding_rect_script)
 
-    if kwargs.get('driver') is None: # only quit webdriver if not passed in as arg
+    if driver is None: # only quit webdriver if not passed in as arg
         web_driver.quit()
 
     image = Image.open(io.BytesIO(png))
@@ -660,7 +660,7 @@ def _get_screenshot_as_png(obj, **kwargs):
 
     return cropped_image
 
-def export_png(obj, filename=None, height=None, width=None):
+def export_png(obj, filename=None, height=None, width=None, webdriver=None):
     ''' Export the LayoutDOM object or document as a PNG.
 
     If the filename is not given, it is derived from the script name
@@ -679,6 +679,9 @@ def export_png(obj, filename=None, height=None, width=None):
         width (int) : the desired width of the exported layout obj only if
             it's a Plot instance. Otherwise the width kwarg is ignored.
 
+        webdriver (selenium.webdriver) : a selenium webdriver instance to use
+            to export the image.
+
     Returns:
         filename (str) : the filename where the static file is saved.
 
@@ -691,7 +694,7 @@ def export_png(obj, filename=None, height=None, width=None):
 
     '''
 
-    image = _get_screenshot_as_png(obj, height=height, width=width)
+    image = _get_screenshot_as_png(obj, height=height, width=width, driver=webdriver)
 
     if filename is None:
         filename = _detect_filename("png")
@@ -700,7 +703,7 @@ def export_png(obj, filename=None, height=None, width=None):
 
     return os.path.abspath(filename)
 
-def _get_svgs(obj, **kwargs):
+def _get_svgs(obj, driver=None, **kwargs):
     webdriver = import_required('selenium.webdriver',
                                 'To use bokeh.io.export_svgs you need selenium ' +
                                 '("conda install -c bokeh selenium" or "pip install selenium")')
@@ -709,7 +712,7 @@ def _get_svgs(obj, **kwargs):
 
     html_path = _save_layout_html(obj, **kwargs)
 
-    web_driver = kwargs.get('driver', webdriver.PhantomJS(service_log_path=os.path.devnull))
+    web_driver = driver if driver is not None else webdriver.PhantomJS(service_log_path=os.path.devnull)
     web_driver.get("file:///" + html_path)
 
     _wait_until_render_complete(web_driver)
@@ -726,12 +729,12 @@ def _get_svgs(obj, **kwargs):
 
     svgs = web_driver.execute_script(svg_script)
 
-    if kwargs.get('driver') is None: # only quit webdriver if not passed in as arg
+    if driver is None: # only quit webdriver if not passed in as arg
         web_driver.quit()
 
     return svgs
 
-def export_svgs(obj, filename=None, height=None, width=None):
+def export_svgs(obj, filename=None, height=None, width=None, webdriver=None):
     ''' Export the SVG-enabled plots within a layout. Each plot will result
     in a distinct SVG file.
 
@@ -750,6 +753,9 @@ def export_svgs(obj, filename=None, height=None, width=None):
         width (int) : the desired width of the exported layout obj only if
             it's a Plot instance. Otherwise the width kwarg is ignored.
 
+        webdriver (selenium.webdriver) : a selenium webdriver instance to use
+            to export the image.
+
     Returns:
         filenames (list(str)) : the list of filenames where the SVGs files
             are saved.
@@ -759,7 +765,7 @@ def export_svgs(obj, filename=None, height=None, width=None):
         aspect ratios. It is recommended to use the default ``fixed`` sizing mode.
 
     '''
-    svgs = _get_svgs(obj, height=height, width=width)
+    svgs = _get_svgs(obj, height=height, width=width, driver=webdriver)
 
     if len(svgs) == 0:
         logger.warn("No SVG Plots were found.")


Test patch: