Merge pull request #24 from joao-voltarelli/fix/master/print_pdf_fix

BUG: fix related to saving files in print_pdf method - Issue #20

Author: hhslepicka

botcity/web/bot.py

@@ -1,5 +1,6 @@
 import base64
 import functools
+import glob
 import io
 import json
 import logging
@@ -973,16 +974,27 @@ def print_pdf(self, path=None, print_options=None):
         """
         title = self.page_title() or "document"
         timeout = 60000
-        if not self.page_title():
-            timeout = 1000
         default_path = os.path.expanduser(os.path.join(self.download_folder_path, f"{title}.pdf"))
 
         if self.browser in [Browser.CHROME, Browser.EDGE] and not self.headless:
+            pdf_current_count = self.get_file_count(file_extension=".pdf")
             # 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, timeout=timeout)
+            if self.page_title():
+                self.wait_for_file(default_path, timeout=timeout)
+            else:
+                # Waiting when the file don't have the page title in path
+                self.wait_for_new_file(file_extension=".pdf", current_count=pdf_current_count)
+
+            # Move the downloaded pdf file if the path is not None
+            if path:
+                last_downloaded_pdf = self.get_last_created_file(self.download_folder_path, ".pdf")
+                os.rename(last_downloaded_pdf, path)
+                return path
+            self.wait(2000)
             return default_path
 
         if print_options is None:
@@ -1858,3 +1870,63 @@ def wait_for_file(self, path, timeout=60000):
                     continue
                 return True
             self.sleep(config.DEFAULT_SLEEP_AFTER_ACTION)
+
+    def get_last_created_file(self, path=None, file_extension=""):
+        """Returns the last created file in a specific folder path.
+
+        Args:
+            path (str, optional): The path of the folder where the file is expected. Defaults to None.
+            file_extension (str, optional): The extension of the file to be searched for (e.g., .pdf, .txt).
+
+        Returns:
+            str: the path of the last created file
+        """
+        if not path:
+            path = self.download_folder_path
+
+        files_path = glob.glob(os.path.expanduser(os.path.join(path, f"*{file_extension}")))
+        last_created_file = max(files_path, key=os.path.getctime)
+        return last_created_file
+
+    def get_file_count(self, path=None, file_extension=""):
+        """Get the total number of files of the same type.
+
+        Args:
+            path (str, optional): The path of the folder where the files are saved.
+            file_extension (str, optional): The extension of the files to be searched for (e.g., .pdf, .txt).
+
+        Returns:
+            int: the number of files of the given type
+        """
+        if not path:
+            path = self.download_folder_path
+
+        files_path = glob.glob(os.path.expanduser(os.path.join(path, f"*{file_extension}")))
+        return len(files_path)
+
+    def wait_for_new_file(self, path=None, file_extension="", current_count=0, timeout=60000):
+        """
+        Wait for a new file to be available on disk without the file path.
+
+        Args:
+            path (str, optional): The path of the folder where the file is expected. Defaults to None.
+            file_extension (str, optional): The extension of the file to be searched for (e.g., .pdf, .txt).
+            current_count (int): The current number of files in the folder of the given type. Defaults to 0 files
+            timeout (int, optional): Maximum wait time (ms) to search for a hit.
+                Defaults to 60000ms (60s).
+
+        Returns:
+            str: the path of the last created file of the given type
+        """
+        if not path:
+            path = self.download_folder_path
+
+        start_time = time.time()
+        while True:
+            elapsed_time = (time.time() - start_time) * 1000
+            if elapsed_time > timeout:
+                return None
+            pdf_count = self.get_file_count(path, f"*{file_extension}")
+            if pdf_count == current_count + 1:
+                return self.get_last_created_file(path, f"*{file_extension}")
+            self.sleep(config.DEFAULT_SLEEP_AFTER_ACTION)

botcity/web/browsers/chrome.py

@@ -83,7 +83,8 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
         },
         "safebrowsing.enabled": True,
         "credentials_enable_service": False,
-        "profile.password_manager_enabled": False
+        "profile.password_manager_enabled": False,
+        "plugins.always_open_pdf_externally": True
     }
 
     chrome_options.add_experimental_option("prefs", prefs)

botcity/web/browsers/edge.py

@@ -86,7 +86,8 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
         },
         "safebrowsing.enabled": True,
         "credentials_enable_service": False,
-        "profile.password_manager_enabled": False
+        "profile.password_manager_enabled": False,
+        "plugins.always_open_pdf_externally": True
     }
 
     edge_options.add_experimental_option("prefs", prefs)