From 34b114ff3c0a1094b76fbecbafe6815946724d55 Mon Sep 17 00:00:00 2001
From: Michael Mintz <mdmintz@gmail.com>
Date: Sun, 10 Nov 2019 23:56:39 -0500
Subject: [PATCH] Add the save_page_source(file_name) method

---
 seleniumbase/fixtures/base_case.py    |  9 +++++++++
 seleniumbase/fixtures/page_actions.py | 27 +++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/seleniumbase/fixtures/base_case.py b/seleniumbase/fixtures/base_case.py
index 7effd520..3280c4c5 100755
--- a/seleniumbase/fixtures/base_case.py
+++ b/seleniumbase/fixtures/base_case.py
@@ -1468,6 +1468,15 @@ class BaseCase(unittest.TestCase):
         """ The screenshot will be in PNG format. """
         return page_actions.save_screenshot(self.driver, name, folder)
 
+    def save_page_source(self, name, folder=None):
+        """ Saves the page HTML to the current directory (or given subfolder).
+            If the folder specified doesn't exist, it will get created.
+            @Params
+            name - The file name to save the current page's HTML to.
+            folder - The folder to save the file to. (Default = current folder)
+        """
+        return page_actions.save_page_source(self.driver, name, folder)
+
     def wait_for_ready_state_complete(self, timeout=None):
         try:
             # If there's an alert, skip
diff --git a/seleniumbase/fixtures/page_actions.py b/seleniumbase/fixtures/page_actions.py
index 94fff0da..28517bf7 100755
--- a/seleniumbase/fixtures/page_actions.py
+++ b/seleniumbase/fixtures/page_actions.py
@@ -32,6 +32,7 @@ from selenium.webdriver.remote.errorhandler import NoAlertPresentException
 from selenium.webdriver.remote.errorhandler import NoSuchFrameException
 from selenium.webdriver.remote.errorhandler import NoSuchWindowException
 from seleniumbase.config import settings
+from seleniumbase.core import log_helper
 
 
 def is_element_present(driver, selector, by=By.CSS_SELECTOR):
@@ -488,6 +489,32 @@ def save_screenshot(driver, name, folder=None):
             pass
 
 
+def save_page_source(driver, name, folder=None):
+    """
+    Saves the page HTML to the current directory (or given subfolder).
+    If the folder specified doesn't exist, it will get created.
+    @Params
+    name - The file name to save the current page's HTML to.
+    folder - The folder to save the file to. (Default = current folder)
+    """
+    if "." not in name:
+        name = name + ".html"
+    if folder:
+        abs_path = os.path.abspath('.')
+        file_path = abs_path + "/%s" % folder
+        if not os.path.exists(file_path):
+            os.makedirs(file_path)
+        html_file_path = "%s/%s" % (file_path, name)
+    else:
+        html_file_path = name
+    page_source = driver.page_source
+    html_file = codecs.open(html_file_path, "w+", "utf-8")
+    rendered_source = log_helper.get_html_source_with_base_href(
+        driver, page_source)
+    html_file.write(rendered_source)
+    html_file.close()
+
+
 def _get_last_page(driver):
     try:
         last_page = driver.current_url