Merge pull request #33 from vgalin/V2.0.1

vgalin · web-flow · commit 8a62c321826c · 2021-07-06T01:53:22.000+02:00
V2.0.1
diff --git a/README.md b/README.md
@@ -49,7 +49,7 @@ hti = Html2Image()
 <summary> Multiple arguments can be passed to the constructor (click to expand):</summary>
 
 -   `browser` :  Browser that will be used, set by default to `'chrome'` (the only browser supported by HTML2Image at the moment)
--   `browser_path` : The path or the command that can be used to find the executable of a specific browser.
+-   `browser_executable` : The path or the command that can be used to find the executable of a specific browser.
 -   `output_path` : Path to the folder to which taken screenshots will be outputed. Default is the current working directory of your python program.
 -   `size` : 2-Tuple reprensenting the size of the screenshots that will be taken. Default value is `(1920, 1080)`.
 -   `temp_path` : Path that will be used to put together different resources when screenshotting strings of files. Default value is `%TEMP%/html2image` on Windows, and `/tmp/html2image` on Linux and MacOS.
diff --git a/html2image/browsers/browser.py b/html2image/browsers/browser.py
@@ -9,12 +9,12 @@ def __init__(self, flags):
 
     @property
     @abstractmethod
-    def executable_path(self):
+    def executable(self):
         pass
 
-    @executable_path.setter
+    @executable.setter
     @abstractmethod
-    def executable_path(self, value):
+    def executable(self, value):
         pass
 
     @abstractmethod
diff --git a/html2image/browsers/chrome.py b/html2image/browsers/chrome.py
@@ -1,17 +1,36 @@
 from .browser import Browser
+from .search_utils import get_command_origin, find_first_defined_env_var
 
 import subprocess
 import platform
 import os
 import shutil
 
+ENV_VAR_LOOKUP_TOGGLE = 'HTML2IMAGE_TOGGLE_ENV_VAR_LOOKUP'
 
-def _find_chrome(user_given_path=None):
+CHROME_EXECUTABLE_ENV_VAR_CANDIDATES = [
+    'HTML2IMAGE_CHROME_BIN',
+    'HTML2IMAGE_CHROME_EXE',
+    'CHROME_BIN',
+    'CHROME_EXE',
+]
+
+
+def _find_chrome(user_given_executable=None):
     """ Finds a Chrome executable.
 
     Search Chrome on a given path. If no path given,
     try to find Chrome or Chromium-browser on a Windows or Unix system.
 
+    Parameters
+    ----------
+    - `user_given_executable`: str (optional)
+        + A filepath leading to a Chrome/ Chromium executable
+        + Or a filename found in the current working directory
+        + Or a keyword that executes Chrome/ Chromium, ex:
+            - 'chromium' on linux systems
+            - 'chrome' on windows (if typing `start chrome` in a cmd works)
+
     Raises
     ------
     - `FileNotFoundError`
@@ -23,15 +42,58 @@ def _find_chrome(user_given_path=None):
         + Path of the chrome executable on the current machine.
     """
 
-    # TODO when other browsers will be available:
-    # Ensure that the given executable is a chrome one.
+    # try to find a chrome bin/exe in ENV
+    path_from_env = find_first_defined_env_var(
+        env_var_list=CHROME_EXECUTABLE_ENV_VAR_CANDIDATES,
+        toggle=ENV_VAR_LOOKUP_TOGGLE
+    )
 
-    if user_given_path is not None:
-        if os.path.isfile(user_given_path):
-            return user_given_path
+    if path_from_env:
+        print(
+            f'Found a potential chrome executable in the {path_from_env} '
+            f'environment variable:\n{path_from_env}\n'
+        )
+        return path_from_env
+
+    # if an executable is given, try to use it
+    if user_given_executable is not None:
+
+        # On Windows, we cannot "safely" validate that user_given_executable
+        # seems to be a chrome executable, as we cannot run it with
+        # the --version flag.
+        # https://bugs.chromium.org/p/chromium/issues/detail?id=158372
+        #
+        # We thus do the "bare minimum" and check if user_given_executable
+        # is a file, a filepath, or corresponds to a keyword that can be used
+        # with the start command, like so: `start user_given_executable`
+        if platform.system() == 'Windows':
+            command_origin = get_command_origin(user_given_executable)
+            if command_origin:
+                return command_origin
+
+            # cannot validate user_given_executable
+            raise FileNotFoundError()
+
+        # On a non-Windows OS, we can validate in a basic way that
+        # user_given_executable leads to a Chrome / Chromium executable,
+        # or is a command, using the --version flag
         else:
-            raise FileNotFoundError('Could not find chrome in the given path.')
+            try:
+                if 'chrom' in subprocess.check_output(
+                    [user_given_executable, '--version']
+                ).decode('utf-8').lower():
+                    return user_given_executable
+            except Exception:
+                pass
+
+        # We got a user_given_executable but couldn't validate it
+        raise FileNotFoundError(
+            'Failed to find a seemingly valid chrome executable '
+            'in the given path.'
+        )
 
+    # Executable not in ENV or given by the user, try to find it
+    # Search for executable on a Windows OS
     if platform.system() == 'Windows':
         prefixes = [
             os.getenv('PROGRAMFILES(X86)'),
@@ -46,45 +108,54 @@ def _find_chrome(user_given_path=None):
             if os.path.isfile(path_candidate):
                 return path_candidate
 
+    # Search for executable on a Linux OS
     elif platform.system() == "Linux":
 
-        # search google-chrome
-        version_result = subprocess.check_output(
-            ["google-chrome", "--version"]
-        )
-
-        if 'Google Chrome' in str(version_result):
-            return "google-chrome"
+        chrome_commands = [
+            'chromium',
+            'chromium-browser',
+            'chrome',
+            'google-chrome'
+        ]
 
-        # else search chromium-browser
+        for chrome_command in chrome_commands:
+            if shutil.which(chrome_command):
+                # check the --version for "chrom" ?
+                return chrome_command
 
         # snap seems to be a special case?
         # see https://stackoverflow.com/q/63375327/12182226
-        version_result = subprocess.check_output(
-            ["chromium-browser", "--version"]
-        )
-        if 'snap' in str(version_result):
-            chrome_snap = (
-                '/snap/chromium/current/usr/lib/chromium-browser/chrome'
-            )
-            if os.path.isfile(chrome_snap):
-                return chrome_snap
-        else:
-            which_result = shutil.which('chromium-browser')
-            if which_result is not None and os.path.isfile(which_result):
-                return which_result
 
+        try:
+            version_result = subprocess.check_output(
+                ["chromium-browser", "--version"]
+            )
+            if 'snap' in str(version_result):
+                chrome_snap = (
+                    '/snap/chromium/current/usr/lib/chromium-browser/chrome'
+                )
+                if os.path.isfile(chrome_snap):
+                    return chrome_snap
+        except Exception:
+            pass
+
+    # Search for executable on MacOS
     elif platform.system() == "Darwin":
         # MacOS system
         chrome_app = (
             '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome'
         )
-        version_result = subprocess.check_output(
-            [chrome_app, "--version"]
-        )
-        if "Google Chrome" in str(version_result):
-            return chrome_app
 
+        try:
+            version_result = subprocess.check_output(
+                [chrome_app, "--version"]
+            )
+            if "Google Chrome" in str(version_result):
+                return chrome_app
+        except Exception:
+            pass
+
+    # Couldn't find an executable (or OS not in Windows, Linux or Mac)
     raise FileNotFoundError(
         'Could not find a Chrome executable on this '
         'machine, please specify it yourself.'
@@ -97,7 +168,7 @@ class ChromeHeadless(Browser):
 
         Parameters
         ----------
-        - `executable_path` : str, optional
+        - `executable` : str, optional
             + Path to a chrome executable.
 
         - `flags` : list of str
@@ -109,8 +180,8 @@ class ChromeHeadless(Browser):
             + Whether or not to print the command used to take a screenshot.
     """
 
-    def __init__(self, executable_path=None, flags=None, print_command=False):
-        self.executable_path = executable_path
+    def __init__(self, executable=None, flags=None, print_command=False):
+        self.executable = executable
         if not flags:
             self.flags = [
                 '--default-background-color=0',
@@ -122,12 +193,12 @@ def __init__(self, executable_path=None, flags=None, print_command=False):
         self.print_command = print_command
 
     @property
-    def executable_path(self):
-        return self._executable_path
+    def executable(self):
+        return self._executable
 
-    @executable_path.setter
-    def executable_path(self, value):
-        self._executable_path = _find_chrome(value)
+    @executable.setter
+    def executable(self, value):
+        self._executable = _find_chrome(value)
 
     def screenshot(
         self,
@@ -171,7 +242,7 @@ def screenshot(
         # command used to launch chrome in
         # headless mode and take a screenshot
         command = [
-            f'{self.executable_path}',
+            f'{self.executable}',
             '--headless',
             f'--screenshot={os.path.join(output_path, output_file)}',
             f'--window-size={size[0]},{size[1]}',
diff --git a/html2image/browsers/firefox.py b/html2image/browsers/firefox.py
@@ -10,11 +10,11 @@ def __init__(self):
         )
 
     @property
-    def executable_path(self):
+    def executable(self):
         pass
 
-    @executable_path.setter
-    def executable_path(self, value):
+    @executable.setter
+    def executable(self, value):
         pass
 
     def render(self, **kwargs):
diff --git a/html2image/browsers/search_utils.py b/html2image/browsers/search_utils.py
@@ -0,0 +1,95 @@
+import shutil
+import os
+try:
+    from winreg import ConnectRegistry, OpenKey, QueryValueEx,\
+            HKEY_LOCAL_MACHINE, HKEY_CURRENT_USER, KEY_READ
+except ImportError:
+    # os is not Windows, and there is no need for winreg
+    pass
+
+
+def get_command_origin(command):
+    ''' Finds the path of a given command (windows only).
+
+    This function is inspired by the `start` command on windows.
+    It will search if the given command corresponds :
+    - To an executable in the current directory
+    - To an executable in the PATH environment variable
+    - To an executable mentionned in the registry in the
+      HKEY_LOCAL_MACHINE or HKEY_LOCAL_MACHINE hkeys in
+      SOFTWARE\\[Wow6432Node\\]Microsoft\\Windows\\CurrentVersion\\App Paths\\
+
+    It also "indirectly" validates a path to a file.
+
+    Parameters
+    ----------
+    - `command`: str
+        + command that will be searched.
+
+    Returns
+    -------
+    - str or None
+        + the path corresponding to `command`
+        + None, if no path was found
+    '''
+
+    # `start "command"` itself could be used to run the command but we want to
+    # assess that the command is a valid one.
+    command = command.replace('start ', '')
+
+    # which search in current directory + path
+    # and file extention can be ommited
+    which_result = shutil.which(command)
+    if which_result:
+        return which_result
+
+    # search for command in registry
+
+    command = command.replace('.exe', '').strip()
+
+    # Once combined, these hkeys and keys are where the `start`
+    # command seach for executable.
+    hkeys = [HKEY_LOCAL_MACHINE, HKEY_CURRENT_USER]
+    keys = [
+        "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\"
+        f"App Paths\\{command}.exe",
+
+        "SOFTWARE\\Wow6432Node\\Microsoft\\Windows\\"
+        "CurrentVersion\\App Paths\\"
+        f"{command}.exe",
+    ]
+
+    for hkey in hkeys:
+        with ConnectRegistry(None, hkey) as registry:
+            for key in keys:
+                try:
+                    with OpenKey(registry, key, 0, KEY_READ) as k:
+                        # if no exception: we found the exe
+                        return QueryValueEx(k, '')[0]
+                except Exception:
+                    # cannot open key, do nothing and proceed to the next one
+                    pass
+    return None
+
+
+def find_first_defined_env_var(env_var_list, toggle):
+    '''
+    Returns the value of the first defined environment variable
+    encountered in the `env_var_list` list, but only if the
+    the `toggle` environment variableif defined.
+
+    Parameters
+    ----------
+    - `env_var_list`: list[str]
+        + list of environment variable names
+    - `toggle`: str
+        + environment variable name
+
+
+    '''
+    if toggle in os.environ:
+        for env_var in env_var_list:
+            value = os.environ.get(env_var)
+            if value:
+                return value
+    return None
diff --git a/html2image/html2image.py b/html2image/html2image.py
@@ -37,7 +37,7 @@ class Html2Image():
             + Type of the browser that will be used to take screenshots.
             + Default is Chrome.
 
-        - `browser_path` : str, optional
+        - `browser_executable` : str, optional
             + Path to a browser executable.
 
         - `output_path` : str, optional
@@ -64,11 +64,11 @@ class Html2Image():
     def __init__(
         self,
         browser='chrome',
-        browser_path=None,
+        browser_executable=None,
         output_path=os.getcwd(),
         size=(1920, 1080),
         temp_path=None,
-        custom_flags=[],
+        custom_flags=None,
     ):
 
         if browser.lower() not in browser_map:
@@ -82,7 +82,7 @@ def __init__(
 
         browser_class = browser_map[browser.lower()]
         self.browser = browser_class(
-            executable_path=browser_path,
+            executable=browser_executable,
             flags=custom_flags,
         )
 
diff --git a/pyproject.toml b/pyproject.toml
