feat: add a config to enable retina quality images in screenshots (#1…

…7409) * add feature flag for retina screenshot support * use config for pixel density * run black
apache · Nov 15, 2021 · 3ee9e11 · 3ee9e11
1 parent 9f1bf1c
commit 3ee9e11
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/superset/config.py b/superset/config.py
@@ -1085,7 +1085,11 @@ def SQL_QUERY_MUTATOR(  # pylint: disable=invalid-name,unused-argument
 WEBDRIVER_TYPE = "firefox"
 
 # Window size - this will impact the rendering of the data
-WEBDRIVER_WINDOW = {"dashboard": (1600, 2000), "slice": (3000, 1200)}
+WEBDRIVER_WINDOW = {
+    "dashboard": (1600, 2000),
+    "slice": (3000, 1200),
+    "pixel_density": 1,
+}
 
 # An optional override to the default auth hook used to provide auth to the
 # offline webdriver

diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
@@ -27,7 +27,7 @@
     TimeoutException,
     WebDriverException,
 )
-from selenium.webdriver import chrome, firefox
+from selenium.webdriver import chrome, firefox, FirefoxProfile
 from selenium.webdriver.common.by import By
 from selenium.webdriver.remote.webdriver import WebDriver
 from selenium.webdriver.support import expected_conditions as EC
@@ -51,22 +51,26 @@ class DashboardStandaloneMode(Enum):
 
 
 class WebDriverProxy:
-    def __init__(
-        self, driver_type: str, window: Optional[WindowSize] = None,
-    ):
+    def __init__(self, driver_type: str, window: Optional[WindowSize] = None):
         self._driver_type = driver_type
         self._window: WindowSize = window or (800, 600)
         self._screenshot_locate_wait = current_app.config["SCREENSHOT_LOCATE_WAIT"]
         self._screenshot_load_wait = current_app.config["SCREENSHOT_LOAD_WAIT"]
 
     def create(self) -> WebDriver:
+        pixel_density = current_app.config["WEBDRIVER_WINDOW"].get("pixel_density", 1)
         if self._driver_type == "firefox":
             driver_class = firefox.webdriver.WebDriver
             options = firefox.options.Options()
+            profile = FirefoxProfile()
+            profile.set_preference("layout.css.devPixelsPerPx", str(pixel_density))
+            kwargs: Dict[Any, Any] = dict(options=options, firefox_profile=profile)
         elif self._driver_type == "chrome":
             driver_class = chrome.webdriver.WebDriver
             options = chrome.options.Options()
+            options.add_argument(f"--force-device-scale-factor={pixel_density}")
             options.add_argument(f"--window-size={self._window[0]},{self._window[1]}")
+            kwargs = dict(options=options)
         else:
             raise Exception(f"Webdriver name ({self._driver_type}) not supported")
         # Prepare args for the webdriver init
@@ -75,7 +79,6 @@ def create(self) -> WebDriver:
         for arg in current_app.config["WEBDRIVER_OPTION_ARGS"]:
             options.add_argument(arg)
 
-        kwargs: Dict[Any, Any] = dict(options=options)
         kwargs.update(current_app.config["WEBDRIVER_CONFIGURATION"])
         logger.info("Init selenium driver")
 
@@ -102,7 +105,7 @@ def destroy(driver: WebDriver, tries: int = 2) -> None:
             pass
 
     def get_screenshot(
-        self, url: str, element_name: str, user: "User",
+        self, url: str, element_name: str, user: "User"
     ) -> Optional[bytes]:
         params = {"standalone": DashboardStandaloneMode.REPORT.value}
         req = PreparedRequest()