Add Pyodide handles (#49)

ryanking13 ran into the problem in pyodide/pyodide#3184 (comment) that it's
inconvenient to send results from one @run_in_pyodide invocation to another.
This adds a mechanism to do so. A SeleniumHandle (perhaps can be named better)
can be returned to the host and when passed back into another @run_in_pyodide
function one can use it to access the object.

This should allow us to make fixtures specifically for use with @run_in_pyodide
functions that return one or more SeleniumHandles.

CHANGELOG.md

@@ -3,6 +3,10 @@
 - `JsException` raise from within pyodide is now unpickled correctly in the host. ([#45](https://github.com/pyodide/pytest-pyodide/issues/45))
 - Improve error messages when unpickling error messages with objects that don't exist in the host environment
    ([#46](https://github.com/pyodide/pytest-pyodide/issues/46))
+- Added the `PyodideHandle` class which allows returning a reference to a Python
+  object in the Pyodide runtime from a `@run_in_pyodide` function. This is
+  useful for fixtures designed to be used with `@run_in_pyodide`.
+  ([#49](https://github.com/pyodide/pytest-pyodide/issues/49))
 
 ## [0.22.2] - 2022.09.08
 

README.md

@@ -21,8 +21,8 @@ You would also one at least one of the following runtimes,
 
 ## Usage
 
-1. First you would need a compatible version of Pyodide. You can download the Pyodide build artifacts from releases with,
-   ```
+1. First you need a compatible version of Pyodide. You can download the Pyodide build artifacts from releases with,
+   ```bash
    wget https://github.com/pyodide/pyodide/releases/download/0.21.0/pyodide-build-0.21.0.tar.bz2
    tar xjf pyodide-build-0.21.0.tar.bz2
    mv pyodide dist/
@@ -77,10 +77,55 @@ def test_type_of_int(selenium, x):
     assert type(x) is int
 ```
 
-These arguments must be picklable. You can also use fixtures as long as the
-return values of the fixtures are picklable (most commonly, if they are `None`).
-As a special case, the function will see the `selenium` fixture as `None` inside
-the test.
+The first argument to a `@run_in_pyodide` function must be a browser runner,
+generally a `selenium` fixture. The remaining arguments and the return value of
+the `@run_in_pyodide` function must be picklable. The arguments will be pickled
+in the host Python and unpickled in the Pyodide Python. The reverse will happen
+to the return value. The first `selenium` argument will be `None` inside the
+body of the function (it is used internally by the fixture). Note that a
+consequence of this is that the received arguments are copies. Changes made to
+an argument will not be reflected in the host Python:
+```py
+@run_in_pyodide
+def mutate_dict(selenium, x):
+    x["a"] = -1
+    return x
+
+def test_mutate_dict():
+    d = {"a" : 9, "b" : 7}
+    assert mutate_dict(d) == { "a" : -1, "b" : 7 }
+    # d is unchanged because it was implicitly copied into the Pyodide runtime!
+    assert d == {"a" : 9, "b" : 7}
+```
+
+You can also use fixtures as long as the return values of the fixtures are
+picklable (most commonly, if they are `None`). As a special case, the function
+will see the `selenium` fixture as `None` inside the test.
+
+If you need to return a persistent reference to a Pyodide Python object, you can
+use the special `PyodideHandle` class:
+```py
+@run_in_pyodide
+def get_pyodide_handle(selenium):
+    from pytest_pyodide.decorator import PyodideHandle
+    d = { "a" : 2 }
+    return PyodideHandle(d)
+
+@run_in_pyodide
+def set_value(selenium, h, key, value):
+    h.obj[key] = value
+
+@run_in_pyodide
+def get_value(selenium, h, key, value):
+    return h.obj[key]
+
+def test_pyodide_handle(selenium):
+    h = get_pyodide_handle(selenium)
+    assert get_value(selenium, h, "a") == 2
+    set_value(selenium, h, "a", 3)
+    assert get_value(selenium, h, "a") == 3
+```
+This can be used to create fixtures for use with `@run_in_pyodide`.
 
 It is possible to use `run_in_pyodide` as an inner function:
 
@@ -92,10 +137,7 @@ def test_inner_function(selenium):
         return 7
     assert inner_function(selenium_mock, 6) == 7
 ```
-
-Again both the arguments and return value must be pickleable.
-
-Also, the function will not see closure variables at all:
+However, the function will not see closure variables at all:
 
 ```py
 def test_inner_function_closure(selenium):
@@ -107,6 +149,9 @@ def test_inner_function_closure(selenium):
     # Raises `NameError: 'x' is not defined`
     assert inner_function(selenium_mock) == 7
 ```
+Thus, the only value of inner `@run_in_pyodide` functions is to limit the scope
+of the function definition. If you need a closure, you will have to wrap it in a
+second function call.
 
 ## Specifying a browser
 

pytest_pyodide/_decorator_in_pyodide.py

@@ -0,0 +1,113 @@
+"""
+This file is not imported normally, it is loaded as a string and then exec'd
+into a module called pytest_pyodide.decorator inside of Pyodide.
+
+We use the name `pytest_pyodide.decorator` for this file for two reasons:
+
+1. so that unpickling works smoothly
+2. so that importing PyodideHandle works smoothly inside Pyodide
+
+We could handle 1. by subclassing Unpickler and overriding find_class. Then we
+could give a different name like `from pytest_pyodide.in_pyodide import
+PyodideHandle` or something. But I think this current approach is the easiest
+for users to make sense of. It is probably still quite confusing.
+
+See also:
+https://github.com/pyodide/pytest-pyodide/issues/43
+"""
+
+import ctypes
+import pickle
+from base64 import b64decode, b64encode
+from io import BytesIO
+from typing import Any
+
+
+def pointer_to_object(ptr: int) -> Any:
+    """Interpret ptr as a PyObject* and convert it to the actual Python object.
+
+    Hopefully we got our reference counting right or this will blow up!
+    """
+    # This was the first way I thought of to convert a pointer into a Python
+    # object: use PyObject_SetItem to assign it to a dictionary.
+    # ctypes doesn't seem to have an API to do this directly.
+    temp: dict[int, Any] = {}
+    ctypes.pythonapi.PyObject_SetItem(id(temp), id(0), ptr)
+    return temp[0]
+
+
+class PyodideHandle:
+    """See documentation for the same-name class in decorator.py
+
+    We pickle this with persistent_id (see below) so there is no need for
+    __getstate__. The only reason we pickle with persistent_id is that on the
+    other side when we unpickle we want to inject a selenium instance so that
+    the reference count can be released by the finalizer. It seems most
+    convenient to do that injection with "persistent_load".
+    """
+
+    def __init__(self, obj: Any):
+        self.obj = obj
+        self.ptr = id(obj)
+
+    def __setstate__(self, state: dict[str, Any]):
+        self.ptr = state["ptr"]
+        self.obj = pointer_to_object(self.ptr)
+
+
+class Pickler(pickle.Pickler):
+    def persistent_id(self, obj: Any) -> Any:
+        if not isinstance(obj, PyodideHandle):
+            return None
+        ctypes.pythonapi.Py_IncRef(obj.ptr)
+        return ("PyodideHandle", obj.ptr)
+
+
+def encode(x: Any) -> str:
+    f = BytesIO()
+    p = Pickler(f)
+    p.dump(x)
+    return b64encode(f.getvalue()).decode()
+
+
+def decode(x: str) -> Any:
+    return pickle.loads(b64decode(x))
+
+
+async def run_in_pyodide_main(
+    mod64: str, args64: str, module_filename: str, func_name: str, async_func: bool
+) -> tuple[int, str]:
+    """
+    This actually runs the code for run_in_pyodide.
+    """
+    __tracebackhide__ = True
+
+    # We've pickled and base 64 encoded the ast module and the arguments so first
+    # we have to decode them.
+    mod = decode(mod64)
+    args: tuple[Any] = decode(args64)
+
+    # Compile and execute the ast
+    co = compile(mod, module_filename, "exec")
+    d: dict[str, Any] = {}
+    exec(co, d)
+
+    try:
+        # Look up the appropriate function on the module and execute it.
+        # The first None fills in the "selenium" argument.
+        result = d[func_name](None, *args)
+        if async_func:
+            result = await result
+        return (0, encode(result))
+    except BaseException as e:
+        try:
+            # If tblib is present, we can show much better tracebacks.
+            from tblib import pickling_support
+
+            pickling_support.install()
+        except ImportError:
+            pass
+        return (1, encode(e))
+
+
+__all__ = ["PyodideHandle", "encode"]