Fix lint

apache · Feb 6, 2025 · ff3bab3 · ff3bab3
1 parent 6aca02e
commit ff3bab3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py
@@ -20,7 +20,7 @@
 from datetime import datetime
 from enum import Enum
 from io import BytesIO
-from typing import TYPE_CHECKING, TypedDict
+from typing import cast, TYPE_CHECKING, TypedDict
 
 from flask import current_app
 
@@ -73,7 +73,7 @@ class ScreenshotCachePayload:
     def __init__(
         self,
         image: bytes | None = None,
-        status: str = StatusValues.PENDING,
+        status: StatusValues = StatusValues.PENDING,
         timestamp: str = "",
     ):
         self._image = image
@@ -209,6 +209,7 @@ def get_from_cache_key(cls, cache_key: str) -> ScreenshotCachePayload | None:
             elif isinstance(payload, ScreenshotCachePayload):
                 pass
             elif isinstance(payload, dict):
+                payload = cast(ScreenshotCachePayloadType, payload)
                 payload = ScreenshotCachePayload.from_dict(payload)
             return payload
         logger.info("Failed at getting from cache: %s", cache_key)

diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py
@@ -26,6 +26,7 @@
 from superset.utils.screenshots import (
     BaseScreenshot,
     ScreenshotCachePayload,
+    ScreenshotCachePayloadType,
 )
 
 BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
@@ -120,23 +121,23 @@ def _setup_compute_and_cache(self, mocker: MockerFixture, screenshot_obj):
     def test_happy_path(self, mocker: MockerFixture, screenshot_obj):
         self._setup_compute_and_cache(mocker, screenshot_obj)
         screenshot_obj.compute_and_cache(force=False)
-        cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key")
+        cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key")
         assert cache_payload["status"] == "Updated"
 
     def test_screenshot_error(self, mocker: MockerFixture, screenshot_obj):
         mocks = self._setup_compute_and_cache(mocker, screenshot_obj)
         get_screenshot: MagicMock = mocks.get("get_screenshot")
         get_screenshot.side_effect = Exception
         screenshot_obj.compute_and_cache(force=False)
-        cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key")
+        cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key")
         assert cache_payload["status"] == "Error"
 
     def test_resize_error(self, mocker: MockerFixture, screenshot_obj):
         mocks = self._setup_compute_and_cache(mocker, screenshot_obj)
         resize_image: MagicMock = mocks.get("resize_image")
         resize_image.side_effect = Exception
         screenshot_obj.compute_and_cache(force=False)
-        cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key")
+        cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key")
         assert cache_payload["status"] == "Error"
 
     def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj):
@@ -154,7 +155,7 @@ def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj):
         # Ensure that it processes when force = True
         screenshot_obj.compute_and_cache(force=True)
         get_screenshot.assert_called_once()
-        cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key")
+        cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key")
         assert cache_payload["status"] == "Updated"
 
     def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj):
@@ -176,7 +177,7 @@ def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj):
             force=True, window_size=window_size, thumb_size=thumb_size
         )
         get_screenshot.assert_called_once()
-        cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key")
+        cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key")
         assert cache_payload["image"] != b"initial_value"
 
     def test_resize(self, mocker: MockerFixture, screenshot_obj):