Merge branch 'main' into selenium-handle

pyodide · hoodmane · Oct 18, 2022 · Oct 14, 2022 · Oct 14, 2022 · Oct 17, 2022
commit e4313c46e720ab253e24faac5959eb27629308dd
diff --git a/pytest_pyodide/decorator.py b/pytest_pyodide/decorator.py
@@ -5,7 +5,7 @@
 from collections.abc import Callable, Collection
 from copy import deepcopy
 from io import BytesIO
-from typing import Any
+from typing import Any, Protocol
 
 import pytest
 
@@ -18,16 +18,19 @@ def package_is_built(package_name: str):
     return _package_is_built(package_name, pytest.pyodide_dist_dir)  # type: ignore[arg-type]
 
 
-class SeleniumType:
+class SeleniumType(Protocol):
     JavascriptException: type
     browser: str
 
-    def load_package(self, *args, **kwargs):
+    def load_package(self, pkgs: str | list[str]):
         ...
 
     def run_async(self, code: str):
         ...
 
+    def run_js(self, code: str):
+        ...
+
 
 def _encode(obj: Any) -> str:
     """
@@ -37,20 +40,28 @@ def _encode(obj: Any) -> str:
     return b64encode(pickle.dumps(obj)).decode()
 
 
+class _ReadableFileobj(Protocol):
+    def read(self, __n: int) -> bytes:
+        ...
+
+    def readline(self) -> bytes:
+        ...
+
+
 class Unpickler(pickle.Unpickler):
-    def __init__(self, file, selenium):
+    def __init__(self, file: _ReadableFileobj, selenium: SeleniumType):
         super().__init__(file)
         self.selenium = selenium
 
-    def persistent_load(self, pid):
+    def persistent_load(self, pid: Any) -> Any:
         if not isinstance(pid, tuple) or len(pid) != 2 or pid[0] != "SeleniumHandle":
             raise pickle.UnpicklingError("unsupported persistent object")
         ptr = pid[1]
         # the SeleniumHandle needs access to selenium in order to free the
         # reference count.
         return SeleniumHandle(self.selenium, ptr)
 
-    def find_class(self, module, name):
+    def find_class(self, module: str, name: str) -> Any:
         """
         Catch exceptions that only exist in the pyodide environment and
         convert them to exception in the host.
@@ -73,9 +84,9 @@ class SeleniumHandle:
     Because of this, we don't bother implementing __setstate__.
     """
 
-    def __init__(self, selenium, ptr):
+    def __init__(self, selenium: SeleniumType, ptr: int):
         self.selenium = selenium
-        self.ptr = ptr
+        self.ptr: int | None = ptr
 
     def __del__(self):
         if self.ptr is None:
@@ -88,10 +99,21 @@ def __del__(self):
             """
         )
 
-    def __getstate__(self):
+    def __getstate__(self) -> dict[str, Any]:
         return {"ptr": self.ptr}
 
 
+def _decode(result: str, selenium: SeleniumType) -> Any:
+    try:
+        return Unpickler(BytesIO(b64decode(result)), selenium).load()
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            f"There was a problem with unpickling the return value/exception from your pyodide environment. "
+            f"This usually means the type of the return value does not exist in your host environment. "
+            f"The original message is: {exc}. "
+        ) from None
+
+
 def _create_outer_test_function(
     run_test: Callable,
     node: Any,
@@ -306,7 +328,7 @@ def _run_test(self, selenium: SeleniumType, args: tuple):
         r = selenium.run_async(code)
         [status, result] = r
 
-        result = Unpickler(BytesIO(b64decode(result)), selenium).load()
+        result = _decode(result, selenium)
         if status:
             raise result
         else: