From cd810eaf6ae6def2b2a12c2554244357beb6d07a Mon Sep 17 00:00:00 2001 From: Kristian Klemon Date: Thu, 6 Jun 2024 19:10:05 +0200 Subject: [PATCH] Allow setting headlessness for chrome web driver creation --- src/penai/utils/web_drivers.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/penai/utils/web_drivers.py b/src/penai/utils/web_drivers.py index 9389c42..13fdc73 100644 --- a/src/penai/utils/web_drivers.py +++ b/src/penai/utils/web_drivers.py @@ -10,13 +10,16 @@ @contextmanager -def create_chrome_web_driver() -> Generator[WebDriver, Any, Any]: +def create_chrome_web_driver(headless: bool = True) -> Generator[WebDriver, Any, Any]: """Helper function to instantiate a Chrome WebDriver instance with all options we need.""" driver = None try: chrome_options = Options() - chrome_options.add_argument("--headless") + + if headless: + chrome_options.add_argument("--headless") + chrome_options.add_argument("--disable-gpu") # The screenshot size might deviate from the actual SVG size on high-dpi devices with a device scale factor != 1.0