Fix typing errors.

Author: ayjayt

src/py/kaleido/_kaleido_tab/__init__.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import base64
 from typing import TYPE_CHECKING
 
 import logistro
 
+from kaleido._utils import to_thread
+
 from . import _devtools_utils as _dtools
 from . import _js_logger
 from ._errors import _raise_error
-from ._utils import to_thread
 
 if TYPE_CHECKING:
     import asyncio
@@ -17,6 +19,9 @@
 
     from kaleido import _fig_tools
 
+
+_TEXT_FORMATS = ("svg", "json")  # eps
+
 _logger = logistro.getLogger(__name__)
 
 
@@ -50,7 +55,7 @@ def __init__(self, tab, *, _stepper=False):
 
         """
         self.tab = tab
-        self.js_logger = _js_logger.JavascriptLogger()
+        self.js_logger = _js_logger.JavascriptLogger(self.tab)
         self._stepper = _stepper
 
     async def navigate(self, url: str | Path = ""):
@@ -141,9 +146,28 @@ async def _calc_fig(
         )
         _raise_error(result)
 
+        # TODO(AJP): better define these error mechanics, is this a devtools
+        # function or what
+        # upon implementation of error, might not be necessary
+        # to do these go-lang/c style returns
+        # but we have to collect and associate
+        # with the gather + profile
+        # None-non return values are a problem
+        # In general, need to better understand stuff here
+
         _logger.debug2(f"Result of function call: {result}")
+        js_response, error = _dtools.check_kaleido_js_response(result)
+        if error:
+            raise error
 
-        img_bytes, error = await _dtools.get_bytes(result)
+        if (response_format := js_response.get("format")) == "pdf":
+            img_raw, error = await _dtools.print_pdf(self.tab)
+        else:
+            img_raw = js_response.get("result")  # type: ignore[assignment]
         if error:
             raise error
-        return img_bytes
+
+        if response_format not in _TEXT_FORMATS:
+            return base64.b64decode(img_raw), None
+        else:
+            return str.encode(img_raw), None

src/py/kaleido/_kaleido_tab/_devtools_utils.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import base64
 import json
 from typing import TYPE_CHECKING
 
@@ -74,14 +73,13 @@ async def exec_js_fn(
     return await tab.send_command("Runtime.callFunctionOn", params=params)
 
 
-_TEXT_FORMATS = ("svg", "json")  # eps
-
-
-async def get_bytes(self, response) -> tuple[None, bytes, None | Exception]:
+def check_kaleido_js_response(
+    response,
+) -> tuple[
+    dict,
+    Exception | None,
+]:
     # TODO(AJP) provoke a js error and return js error
-
-    # check_js could be another function, the next two sections should be
-    # factored out
     js_response = json.loads(
         response.get(
             "result",
@@ -95,34 +93,37 @@ async def get_bytes(self, response) -> tuple[None, bytes, None | Exception]:
             "value",
         ),
     )
+    if not js_response:  # not loved, neither {}
+        return {}, RuntimeError(
+            f"JS Response not understood: {response}",
+        )
 
     if js_response["code"] != 0:
-        return None, KaleidoError(js_response["code"], js_response["message"])
-
-    # this shouldn't be in here
-    if (response_format := js_response.get("format")) == "pdf":
-        pdf_params = {
-            "printBackground": True,
-            "marginTop": 0.1,
-            "marginBottom": 0.1,
-            "marginLeft": 0.1,
-            "marginRight": 0.1,
-            "preferCSSPageSize": True,
-            "pageRanges": "1",
-        }
-        pdf_response = await self.tab.send_command(
-            "Page.printToPDF",
-            params=pdf_params,
-        )
-        e = _get_error(pdf_response)
-        if e:
-            return e
-        img_raw = pdf_response.get("result").get("data")
-    else:
-        img_raw = js_response.get("result")
-
-    # Base64 decode binary types
-    if response_format not in _TEXT_FORMATS:
-        return base64.b64decode(img_raw), None
-    else:
-        return str.encode(img_raw), None
+        return {}, KaleidoError(js_response["code"], js_response["message"])
+
+    return js_response
+
+
+async def print_pdf(
+    tab: choreographer.Tab,
+) -> tuple[
+    str,
+    Exception | None,
+]:
+    pdf_params = {
+        "printBackground": True,
+        "marginTop": 0.1,
+        "marginBottom": 0.1,
+        "marginLeft": 0.1,
+        "marginRight": 0.1,
+        "preferCSSPageSize": True,
+        "pageRanges": "1",
+    }
+    pdf_response = await tab.send_command(
+        "Page.printToPDF",
+        params=pdf_params,
+    )
+    e = _get_error(pdf_response)
+    if e:
+        return "", e
+    return pdf_response.get("result", {}).get("data"), None

src/py/kaleido/_kaleido_tab/_js_logger.py

@@ -26,20 +26,21 @@ class JavascriptLogger:
     log: list[Any]
     """A list of console outputs from the tab."""
 
-    def __init__(self) -> None:
+    def __init__(self, tab: choreographer.Tab) -> None:
         self.log = []
+        self.tab = tab
 
-    def activate(self, tab: choreographer.Tab):
-        tab.unsubscribe("Runtime.consoleAPICalled")
-        tab.subscribe(
+    def activate(self):
+        self.tab.unsubscribe("Runtime.consoleAPICalled")
+        self.tab.subscribe(
             "Runtime.consoleAPICalled",
-            _make_console_logger("tab js console", self.javascript_log),
+            _make_console_logger("tab js console", self.log),
         )
 
-    def reset(self, tab: choreographer.Tab):
-        tab.unsubscribe("Runtime.consoleAPICalled")
+    def reset(self):
+        self.tab.unsubscribe("Runtime.consoleAPICalled")
         self.log = []
-        tab.subscribe(
+        self.tab.subscribe(
             "Runtime.consoleAPICalled",
-            _make_console_logger("tab js console", self.javascript_log),
+            _make_console_logger("tab js console", self.log),
         )

src/py/kaleido/_page_generator.py

@@ -8,7 +8,11 @@
 from . import _utils
 
 if TYPE_CHECKING:
-    UrlAndCharset = tuple[str | Path, str]  # type: TypeAlias
+    from typing import Tuple, Union
+
+    from typing_extensions import TypeAlias
+
+    UrlAndCharset: TypeAlias = Tuple[Union[str, Path], str]
     """A tuple to explicitly set charset= in the <script> tag."""
 
 _logger = logistro.getLogger(__name__)