Merge pull request #5 from hhslepicka/print-to-pdf

ENH: Print to PDF

.gitignore

@@ -1,5 +1,8 @@
 .DS_Store
 
+# VSCode IDE files
+.vscode
+
 # PyCharm IDE files
 .idea
 

botcity/web/bot.py

@@ -1,6 +1,7 @@
 import base64
 import functools
 import io
+import json
 import logging
 import multiprocessing
 import os
@@ -230,6 +231,31 @@ def set_screen_resolution(self, width=None, height=None):
                 """, width, height)
         self._driver.set_window_size(*window_size)
 
+    def _webdriver_command(self, command, params=None, req_type="POST"):
+        """
+        Execute a webdriver command.
+
+        Args:
+            command (str): The command URL after the session part
+            params (dict): The payload to be serialized and sent to the webdriver. Defaults to None.
+            req_type (str, optional): The type of request to be made. Defaults to "POST".
+
+        Returns:
+            str: The value of the response
+        """
+        if not params:
+            params = {}
+
+        resource = f"/session/{self.driver.session_id}/{command}"
+        url = self.driver.command_executor._url + resource
+        body = json.dumps(params)
+        response = self.driver.command_executor._request(req_type, url, body)
+
+        if not response:
+            raise Exception(response.get('value'))
+
+        return response.get('value')
+
     ##########
     # Display
     ##########
@@ -747,7 +773,6 @@ def create_tab(self, url):
             self.navigate_to(url)
 
     def create_window(self, url):
-
         try:
             # Refactor this when Selenium 4 is released
             self.execute_javascript(f"window.open('{url}', '_blank', 'location=0');")
@@ -769,10 +794,47 @@ def close_page(self):
     def activate_tab(self, handle):
         self._driver.switch_to.window(handle)
 
+    def print_pdf(self, path=None, print_options=None):
+        """Print the current page as a PDF file.
+
+        Args:
+            path (str, optional): The path for the file to be saved. Defaults to None.
+            print_options (dict, optional): Print options as defined at. Defaults to None.
+
+        Returns:
+            str: the saved file path
+        """
+        title = self.page_title() or "document"
+        default_path = os.path.expanduser(os.path.join("~", "Desktop", f"{title}.pdf"))
+
+        if self.browser in [Browser.CHROME, Browser.EDGE] and not self.headless:
+            # Chrome still does not support headless webdriver print
+            # but Firefox does.
+            self.execute_javascript("window.print();")
+            # We need to wait for the file to be available in this case.
+            self.wait_for_file(default_path)
+            return default_path
+
+        if print_options is None:
+            print_options = {
+                'landscape': False,
+                'displayHeaderFooter': False,
+                'printBackground': True,
+                'preferCSSPageSize': True,
+                'marginTop': 0,
+                'marginBottom': 0
+            }
+        data = self._webdriver_command("print", print_options)
+        bytes_file = base64.b64decode(data)
+        if not path:
+            path = default_path
+        with open(path, "wb") as f:
+            f.write(bytes_file)
+        return path
+
     #######
     # Mouse
     #######
-
     @only_if_element
     def click_on(self, label):
         """
@@ -1491,14 +1553,14 @@ def sleep(self, interval):
         """
         self.wait(interval)
 
-    def wait_for_file(self, path, timeout=10000):
+    def wait_for_file(self, path, timeout=60000):
         """
-        Invoke the system handler to open the given file.
+        Wait for a file to be available on disk.
 
         Args:
             path (str): The path for the file to be executed
             timeout (int, optional): Maximum wait time (ms) to search for a hit.
-                Defaults to 10000ms (10s).
+                Defaults to 60000ms (60s).
 
         Returns:
             status (bool): Whether or not the file was available before the timeout

botcity/web/browsers/chrome.py

@@ -1,3 +1,4 @@
+import json
 import atexit
 import os
 import tempfile
@@ -10,7 +11,6 @@
 
 def default_options(headless=False, download_folder_path=None, user_data_dir=None):
     chrome_options = ChromeOptions()
-    chrome_options.add_argument("--disable-extensions")
     chrome_options.add_argument("--remote-debugging-port=0")
     chrome_options.add_argument("--no-first-run")
     chrome_options.add_argument("--no-default-browser-check")
@@ -48,8 +48,33 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     if not download_folder_path:
         download_folder_path = os.path.join(os.path.expanduser("~"), "Desktop")
 
+    app_state = {
+        'recentDestinations': [{
+            'id': 'Save as PDF',
+            'origin': 'local'
+        }],
+        'selectedDestinationId': 'Save as PDF',
+        'version': 2
+    }
+
     # Set the Downloads default folder
-    prefs = {"download.default_directory": download_folder_path}
+    prefs = {
+        "printing.print_preview_sticky_settings.appState": json.dumps(app_state),
+        "download.default_directory": download_folder_path,
+        "savefile.default_directory": download_folder_path,
+        "printing.default_destination_selection_rules": {
+            "kind": "local",
+            "namePattern": "Save as PDF",
+        },
+        "safebrowsing.enabled": True
+    }
+
     chrome_options.add_experimental_option("prefs", prefs)
+    chrome_options.add_argument(
+        "user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36"
+    )
+
+    chrome_options.add_argument("--kiosk-printing")
 
     return chrome_options

botcity/web/browsers/edge.py

@@ -1,4 +1,5 @@
 import atexit
+import json
 import os
 import tempfile
 
@@ -9,8 +10,7 @@
 
 def default_options(headless=False, download_folder_path=None, user_data_dir=None):
     edge_options = EdgeOptions()
-    edge_options.use_chrome = True
-    edge_options.add_argument("--disable-extensions")
+    edge_options.use_chromium = True
     edge_options.add_argument("--remote-debugging-port=0")
     edge_options.add_argument("--no-first-run")
     edge_options.add_argument("--no-default-browser-check")
@@ -48,8 +48,38 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     if not download_folder_path:
         download_folder_path = os.path.join(os.path.expanduser("~"), "Desktop")
 
+    app_state = {
+        "recentDestinations": [{
+            "id": "Save as PDF",
+            "origin": "local",
+            "account": ""
+        }],
+        "selectedDestinationId": "Save as PDF",
+        "version": 2,
+        "isHeaderFooterEnabled": False,
+        "marginsType": 2,
+        "isCssBackgroundEnabled": True
+    }
+
     # Set the Downloads default folder
-    prefs = {"download.default_directory": download_folder_path}
+    prefs = {
+        "printing.print_preview_sticky_settings.appState": json.dumps(app_state),
+        "download.default_directory": download_folder_path,
+        "savefile.default_directory": download_folder_path,
+        "printing.default_destination_selection_rules": {
+            "kind": "local",
+            "namePattern": "Save as PDF",
+        },
+        "safebrowsing.enabled": True
+    }
+
     edge_options.add_experimental_option("prefs", prefs)
 
+    edge_options.add_argument(
+        "user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36"
+    )
+
+    edge_options.add_argument("--kiosk-printing")
+
     return edge_options