<NME> webdriver.py
<BEF> # coding=utf-8
# @license
# Copyright 2020 Google Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Interface for controlling a browser that runs Neuroglancer."""

from typing import Sequence, Optional

import tempfile
import time

class Webdriver(object):

    def __init__(self, viewer=None, headless=True, window_size=(1920, 1080),
                 print_logs=True):
        import selenium.webdriver
        import selenium.webdriver.chrome.options
        import selenium.webdriver.common.desired_capabilities
        try:
            # Use chromedriver_binary package if available
            import chromedriver_binary
        except ImportError:
            # Fallback to system chromedriver
            pass
        if viewer is None:
            from .viewer import Viewer
            viewer = Viewer()
        self.viewer = viewer
        chrome_options = selenium.webdriver.chrome.options.Options()
        if headless:
            chrome_options.add_argument('--headless')
        chrome_options.add_argument('--window_size=%dx%d' % (window_size[0], window_size[1]))
        caps = selenium.webdriver.common.desired_capabilities.DesiredCapabilities.CHROME.copy()
        caps['goog:loggingPrefs'] = {'browser': 'ALL'}
        self.driver = selenium.webdriver.Chrome(options=chrome_options, desired_capabilities=caps)
        self.driver.get(viewer.get_viewer_url())
        self._pending_logs = []
        self._pending_logs_to_print = []
        self._logs_lock = threading.Lock()
        self.debug = debug
        self._logfile = None
        if browser == 'firefox':
            self._logfile = tempfile.NamedTemporaryFile(suffix='neuroglancer-geckodriver.log')
        self._init_driver()
        self._pending_logs = []
        self._pending_logs_to_print = []
        self._logs_lock = threading.Lock()
        self._closed = False

        t.daemon = True
        t.start()

    def __enter__(self):
        return self

        chrome_options = selenium.webdriver.chrome.options.Options()
        if self.headless:
            chrome_options.add_argument('--headless')
        chrome_options.add_experimental_option("excludeSwitches", ['enable-automation'])
        if self.docker:
            # https://www.intricatecloud.io/2019/05/running-webdriverio-tests-using-headless-chrome-inside-a-container/
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--disable-setuid-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument('--window_size=%dx%d' %
                                    (self.window_size[0], self.window_size[1]))
        for arg in self.extra_command_line_args:
            chrome_options.add_argument(arg)
        caps = selenium.webdriver.common.desired_capabilities.DesiredCapabilities.CHROME.copy()
        caps['goog:loggingPrefs'] = {'browser': 'ALL'}
        try:
            orig_init = selenium.webdriver.chrome.service.Service.__init__
            if self.debug:
                selenium.webdriver.chrome.service.Service.__init__ = patched_init
            self.driver = selenium.webdriver.Chrome(executable_path=executable_path,
                                                    options=chrome_options,
                                                    desired_capabilities=caps)
        finally:
            if self.debug:
                selenium.webdriver.chrome.service.Service.__init__ = orig_init

    def _init_firefox(self):
        import selenium.webdriver
        import selenium.webdriver.firefox.firefox_binary
        executable_path = 'geckodriver'
        try:
            # Use webdriver_manager package if available
                return new_state
            time.sleep(0.1)

    @property
    def root_element(self):
        return self.driver.find_element_by_xpath('//body')
    def _init_driver(self):
        if self.browser == 'chrome':
            self._init_chrome()
        elif self.browser == 'firefox':
            self._init_firefox()
        else:
            raise ValueError('unsupported browser: %s, must be "chrome" or "firefox"' %
                             (self.browser, ))
        self.driver.get(self.viewer.get_viewer_url())

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if self._logfile is not None:
            self._logfile.file.close()
            self._logfile = None
        self.driver.quit()
        self._closed = True

    def _get_new_logs(self):
        if self.browser == 'chrome':
            new_logs = self.driver.get_log('browser')
        else:
            cur_offset = self._logfile.file.tell()
            new_data = self._logfile.file.read()
            # rfind may return -1, still works
            end_within_data = new_data.rfind(b'\n') + 1
            new_data = new_data[:end_within_data]
            self._logfile.file.seek(cur_offset + end_within_data)
            new_logs = []
            for msg in new_data.decode().split('\n'):
                msg = msg.strip()
                if not msg: continue
                if (not msg.startswith('console.log: ') and not msg.startswith('JavaScript ')):
                    continue
                new_logs.append({'message': msg})
        self._pending_logs.extend(new_logs)
        self._pending_logs_to_print.extend(new_logs)

    def _get_logs_to_print(self):
        with self._logs_lock:
            self._get_new_logs()
            new_logs = self._pending_logs_to_print
            self._pending_logs_to_print = []
            return new_logs

    def get_log(self):
        with self._logs_lock:
            self._get_new_logs()
            new_logs = self._pending_logs
            self._pending_logs = []
            return new_logs

    def get_log_messages(self):
        return '\n'.join(x['message'] for x in self.get_log())

    def sync(self):
        """Wait until client is ready."""
        while True:
            new_state = self.viewer.screenshot().viewer_state
            # Ensure self.viewer.state has also been updated to the new state.
            # The state sent in the screenshot reply can be newer.
            if new_state == self.viewer.state:
                return new_state
            time.sleep(0.1)

    def reload_browser(self):
        """Reloads the browser (useful if it crashes/becomes unresponsive)."""
        with self._logs_lock:
            try:
                self.driver.quit()
            except:
                pass
            self._init_driver()

    @property
    def root_element(self):
        return self.driver.find_element('xpath', '//body')

    def action_chain(self):
        import selenium.webdriver
        return selenium.webdriver.common.action_chains.ActionChains(self.driver)

<MSG> feat(python): add reload and docker support to webdriver interface

<DFF> @@ -21,29 +21,23 @@ import threading
 
 class Webdriver(object):
 
-    def __init__(self, viewer=None, headless=True, window_size=(1920, 1080),
+    def __init__(self,
+                 viewer=None,
+                 headless=True,
+                 window_size=(1920, 1080),
+                 debug=False,
+                 docker=False,
                  print_logs=True):
-        import selenium.webdriver
-        import selenium.webdriver.chrome.options
-        import selenium.webdriver.common.desired_capabilities
-        try:
-            # Use chromedriver_binary package if available
-            import chromedriver_binary
-        except ImportError:
-            # Fallback to system chromedriver
-            pass
         if viewer is None:
             from .viewer import Viewer
             viewer = Viewer()
         self.viewer = viewer
-        chrome_options = selenium.webdriver.chrome.options.Options()
-        if headless:
-            chrome_options.add_argument('--headless')
-        chrome_options.add_argument('--window_size=%dx%d' % (window_size[0], window_size[1]))
-        caps = selenium.webdriver.common.desired_capabilities.DesiredCapabilities.CHROME.copy()
-        caps['goog:loggingPrefs'] = {'browser': 'ALL'}
-        self.driver = selenium.webdriver.Chrome(options=chrome_options, desired_capabilities=caps)
-        self.driver.get(viewer.get_viewer_url())
+        self.headless = headless
+        self.window_size = window_size
+        self.headless = headless
+        self.docker = docker
+        self.debug = debug
+        self._init_driver()
         self._pending_logs = []
         self._pending_logs_to_print = []
         self._logs_lock = threading.Lock()
@@ -60,6 +54,56 @@ class Webdriver(object):
         t.daemon = True
         t.start()
 
+    def _init_driver(self):
+        import selenium.webdriver
+        import selenium.webdriver.chrome.options
+        import selenium.webdriver.common.service
+        import selenium.webdriver.chrome.service
+
+        def patched_init(self, executable_path, port=0, service_args=None, log_path=None, env=None):
+            log_path = '/dev/stderr'
+            self.service_args = service_args or []
+            if log_path:
+                self.service_args.append('--log-path=%s' % log_path)
+            selenium.webdriver.common.service.Service.__init__(
+                self,
+                executable_path,
+                port=port,
+                env=env,
+                log_file=None,
+                start_error_message=
+                "Please see https://sites.google.com/a/chromium.org/chromedriver/home")
+
+        import selenium.webdriver.common.desired_capabilities
+        try:
+            # Use chromedriver_binary package if available
+            import chromedriver_binary
+        except ImportError:
+            # Fallback to system chromedriver
+            pass
+        chrome_options = selenium.webdriver.chrome.options.Options()
+        if self.headless:
+            chrome_options.add_argument('--headless')
+        if self.docker:
+            # https://www.intricatecloud.io/2019/05/running-webdriverio-tests-using-headless-chrome-inside-a-container/
+            chrome_options.add_argument('--no-sandbox')
+            chrome_options.add_argument('--disable-gpu')
+            chrome_options.add_argument('--disable-setuid-sandbox')
+            chrome_options.add_argument('--disable-dev-shm-usage')
+        chrome_options.add_argument('--window_size=%dx%d' %
+                                    (self.window_size[0], self.window_size[1]))
+        caps = selenium.webdriver.common.desired_capabilities.DesiredCapabilities.CHROME.copy()
+        caps['goog:loggingPrefs'] = {'browser': 'ALL'}
+        try:
+            orig_init = selenium.webdriver.chrome.service.Service.__init__
+            if self.debug:
+                selenium.webdriver.chrome.service.Service.__init__ = patched_init
+            self.driver = selenium.webdriver.Chrome(options=chrome_options, desired_capabilities=caps)
+        finally:
+            if self.debug:
+                selenium.webdriver.chrome.service.Service.__init__ = orig_init
+        self.driver.get(self.viewer.get_viewer_url())
+
     def __enter__(self):
         return self
 
@@ -99,6 +143,15 @@ class Webdriver(object):
                 return new_state
             time.sleep(0.1)
 
+    def reload_browser(self):
+        """Reloads the browser (useful if it crashes/becomes unresponsive)."""
+        with self._logs_lock:
+            try:
+                self.driver.quit()
+            except:
+                pass
+            self._init_driver()
+
     @property
     def root_element(self):
         return self.driver.find_element_by_xpath('//body')

feat(python): add reload and docker support to webdriver interface