From 39d6b1f1e0106bff8a8dea1c73571f4f936c2728 Mon Sep 17 00:00:00 2001
From: yvsvarma <yvsvarma@gmail.com>
Date: Sun, 12 Jan 2025 08:51:05 -0600
Subject: [PATCH] Enhance PrintOptions to support default, predefined, and
 custom page sizes in Python (#15052)

---
 .../webdriver/common/print_page_options.py    | 29 ++++++++++++++++++-
 .../common/print_page_options_tests.py        | 16 ++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/py/selenium/webdriver/common/print_page_options.py b/py/selenium/webdriver/common/print_page_options.py
index 1ffde95c88e67..e0ae9b1b1c600 100644
--- a/py/selenium/webdriver/common/print_page_options.py
+++ b/py/selenium/webdriver/common/print_page_options.py
@@ -402,16 +402,43 @@ class PrintOptions:
     - Set
         - `None`
     """
+    # Predefined page sizes (in centimeters)
+    A4 = {"height": 29.7, "width": 21.0}  # size in cm
+    LEGAL = {"height": 35.56, "width": 21.59}  # size in cm
+    LETTER = {"height": 27.94, "width": 21.59}  # size in cm
+    TABLOID = {"height": 43.18, "width": 27.94}  # size in cm
 
     def __init__(self) -> None:
         self._print_options: _PrintOpts = {}
-        self._page: _PageOpts = {}
+        self._page: _PageOpts = {"height": PrintOptions.A4["height"], "width": PrintOptions.A4["width"]}  # Default page size set to A4
         self._margin: _MarginOpts = {}
 
     def to_dict(self) -> _PrintOpts:
         """:Returns: A hash of print options configured."""
         return self._print_options
 
+    def set_page_size(self, page_size: dict) -> None:
+        """
+        Sets the page size to predefined or custom dimensions.
+
+        Parameters
+        ----------
+        page_size: dict
+        A dictionary containing `height` and `width` as keys with respective values.
+
+        Example
+        -------
+        self.set_page_size(PageSize.A4)  # A4 predefined size
+        self.set_page_size({"height": 15.0, "width": 20.0})  # Custom size in cm
+
+        """
+        self._validate_num_property("height", page_size["height"])
+        self._validate_num_property("width", page_size["width"])
+        self._page["height"] = page_size["height"]
+        self._page["width"] = page_size["width"]
+        self._print_options["page"] = self._page
+
+
     def _validate_num_property(self, property_name: str, value: float) -> None:
         """Helper function to validate some of the properties."""
         if not isinstance(value, (int, float)):
diff --git a/py/test/unit/selenium/webdriver/common/print_page_options_tests.py b/py/test/unit/selenium/webdriver/common/print_page_options_tests.py
index 76666de94f746..1c276bb165c4f 100644
--- a/py/test/unit/selenium/webdriver/common/print_page_options_tests.py
+++ b/py/test/unit/selenium/webdriver/common/print_page_options_tests.py
@@ -44,6 +44,22 @@ def test_raises_exception_if_scale_is_outside_range(print_options):
     with pytest.raises(ValueError):
         print_options.scale = 3
 
+def test_set_page_size(print_options):
+    # Example of setting a default (A4)
+    assert print_options.page_width == PrintOptions.A4["width"]
+    assert print_options.page_height == PrintOptions.A4["height"]
+
+    # Example of setting a predefined page size
+    print_options.set_page_size(PrintOptions.LEGAL)
+    assert print_options.page_width == PrintOptions.LEGAL["width"]
+    assert print_options.page_height == PrintOptions.LEGAL["height"]
+
+    # Test custom page size
+    custom_size = {"height": 25.0, "width": 15.0}
+    print_options.set_page_size(custom_size)
+    assert print_options.page_width == custom_size["width"]
+    assert print_options.page_height == custom_size["height"]
+
 
 def test_raises_exception_if_scale_is_not_an_integer(print_options):
     with pytest.raises(ValueError):