From d818f410c4098c34990bf2a82412f97f08e7668c Mon Sep 17 00:00:00 2001
From: Hood Chatham <roberthoodchatham@gmail.com>
Date: Fri, 27 May 2022 12:28:46 -0700
Subject: [PATCH] Use tblib to pickle errors in `run_in_pyodide` (#2619)

---
 .circleci/config.yml                          |  2 +-
 packages/tblib/meta.yaml                      | 14 ++++
 pyodide-build/pyodide_build/common.py         |  1 +
 .../pyodide_test_runner/decorator.py          | 44 ++++++-------
 .../tests/test_decorator.py                   | 64 +++++--------------
 requirements.txt                              |  1 +
 run_docker                                    |  2 +-
 7 files changed, 53 insertions(+), 75 deletions(-)
 create mode 100644 packages/tblib/meta.yaml

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 17fbcb1b6..41aecc9f9 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -6,7 +6,7 @@ defaults: &defaults
     # Note: when updating the docker image version,
     #       make sure there are no extra old versions lying around.
     #       (e.g. `rg -F --hidden <old_tag>`)
-    - image: pyodide/pyodide-env:20220504-py310-chrome101-firefox100
+    - image: pyodide/pyodide-env:20220525-py310-chrome102-firefox100
   environment:
     - EMSDK_NUM_CORES: 3
       EMCC_CORES: 3
diff --git a/packages/tblib/meta.yaml b/packages/tblib/meta.yaml
new file mode 100644
index 000000000..2fe70c51d
--- /dev/null
+++ b/packages/tblib/meta.yaml
@@ -0,0 +1,14 @@
+package:
+  name: tblib
+  version: 1.7.0
+source:
+  url: https://files.pythonhosted.org/packages/f8/cd/2fad4add11c8837e72f50a30e2bda30e67a10d70462f826b291443a55c7d/tblib-1.7.0-py2.py3-none-any.whl
+  sha256: 289fa7359e580950e7d9743eab36b0691f0310fce64dee7d9c31065b8f723e23
+test:
+  imports:
+    - tblib
+about:
+  home: https://github.com/ionelmc/python-tblib
+  PyPI: https://pypi.org/project/tblib
+  summary: Traceback serialization library.
+  license: BSD-2-Clause
diff --git a/pyodide-build/pyodide_build/common.py b/pyodide-build/pyodide_build/common.py
index bdd5658a4..e7ca01037 100644
--- a/pyodide-build/pyodide_build/common.py
+++ b/pyodide-build/pyodide_build/common.py
@@ -76,6 +76,7 @@ CORE_PACKAGES = {
     "cpp-exceptions-test",
     "ssl",
     "pytest",
+    "tblib",
 }
 
 CORE_SCIPY_PACKAGES = {
diff --git a/pyodide-test-runner/pyodide_test_runner/decorator.py b/pyodide-test-runner/pyodide_test_runner/decorator.py
index 6afdba2bc..6d1c9a59c 100644
--- a/pyodide-test-runner/pyodide_test_runner/decorator.py
+++ b/pyodide-test-runner/pyodide_test_runner/decorator.py
@@ -3,7 +3,6 @@ import pickle
 import sys
 from base64 import b64decode, b64encode
 from copy import deepcopy
-from traceback import TracebackException
 from typing import Any, Callable, Collection
 
 from pyodide_test_runner.utils import package_is_built as _package_is_built
@@ -169,6 +168,9 @@ class run_in_pyodide:
             REWRITTEN_MODULE_ASTS if pytest_assert_rewrites else ORIGINAL_MODULE_ASTS
         )
 
+        if package_is_built("tblib"):
+            self._pkgs.append("tblib")
+
         self._pytest_assert_rewrites = pytest_assert_rewrites
 
     def _code_template(self, args: tuple) -> str:
@@ -186,15 +188,20 @@ class run_in_pyodide:
             co = compile(mod, {self._module_filename!r}, "exec")
             d = {{}}
             exec(co, d)
+            def encode(x):
+                return b64encode(pickle.dumps(x)).decode()
             try:
                 result = d[{self._func_name!r}](None, *args)
                 if {self._async_func}:
                     result = await result
+                return [0, encode(result)]
             except BaseException as e:
-                import traceback
-                tb = traceback.TracebackException(type(e), e, e.__traceback__)
-                serialized_err = pickle.dumps(tb)
-                return b64encode(serialized_err).decode()
+                try:
+                    from tblib import pickling_support
+                    pickling_support.install()
+                except ImportError:
+                    pass
+                return [1, encode(e)]
 
         try:
             result = await __tmp()
@@ -210,27 +217,14 @@ class run_in_pyodide:
         if self._pkgs:
             selenium.load_package(self._pkgs)
 
-        result = selenium.run_async(code)
+        r = selenium.run_async(code)
+        [status, result] = r
 
-        if result:
-            err: TracebackException = pickle.loads(b64decode(result))
-            err.stack.pop(0)  # Get rid of __tmp in traceback
-            self._fail(err)
-
-    def _fail(self, err: TracebackException):
-        """
-        Fail the test with a helpful message.
-
-        Separated out for test mock purposes.
-        """
-        msg = "Error running function in pyodide\n\n" + "".join(err.format(chain=True))
-        if self._pytest_not_built:
-            msg += (
-                "\n"
-                "Note: pytest not available in Pyodide. We could generate a"
-                "better traceback if pytest were available."
-            )
-        pytest.fail(msg, pytrace=False)
+        result = pickle.loads(b64decode(result))
+        if status:
+            raise result
+        else:
+            return result
 
     def _generate_pyodide_ast(
         self, module_ast: ast.Module, funcname: str, func_line_no: int
diff --git a/pyodide-test-runner/pyodide_test_runner/tests/test_decorator.py b/pyodide-test-runner/pyodide_test_runner/tests/test_decorator.py
index a14e437d5..6af630404 100644
--- a/pyodide-test-runner/pyodide_test_runner/tests/test_decorator.py
+++ b/pyodide-test-runner/pyodide_test_runner/tests/test_decorator.py
@@ -1,6 +1,4 @@
 import asyncio
-from dataclasses import dataclass, field
-from typing import Any
 
 import pytest
 from pyodide_test_runner.decorator import run_in_pyodide
@@ -53,43 +51,19 @@ class selenium_mock:
         return asyncio.new_event_loop().run_until_complete(eval_code_async(code))
 
 
-@dataclass
-class local_mocks_cls:
-    exc_list: list[Any] = field(default_factory=list)
-
-    def check_err(self, ty, msg):
-        try:
-            assert self.exc_list
-            err = self.exc_list[0]
-            assert err
-            assert "".join(err.format_exception_only()) == msg
-        finally:
-            del self.exc_list[0]
-
-    def _patched_fail(self, exc):
-        self.exc_list.append(exc)
+def test_local1():
+    with pytest.raises(AssertionError, match="assert 6 == 7"):
+        example_func1(selenium_mock)
 
 
-@pytest.fixture
-def local_mocks(monkeypatch):
-    mocks = local_mocks_cls()
-    monkeypatch.setattr(run_in_pyodide, "_fail", mocks._patched_fail)
-    return mocks
+def test_local2():
+    with pytest.raises(AssertionError, match="assert 6 == 7"):
+        example_func2(selenium_mock)
 
 
-def test_local1(local_mocks):
-    example_func1(selenium_mock)
-    local_mocks.check_err(AssertionError, "AssertionError: assert 6 == 7\n")
-
-
-def test_local2(local_mocks):
-    example_func1(selenium_mock)
-    local_mocks.check_err(AssertionError, "AssertionError: assert 6 == 7\n")
-
-
-def test_local3(local_mocks):
-    async_example_func(selenium_mock)
-    local_mocks.check_err(AssertionError, "AssertionError: assert 6 == 7\n")
+def test_local3():
+    with pytest.raises(AssertionError, match="assert 6 == 7"):
+        async_example_func(selenium_mock)
 
 
 def test_local_inner_function():
@@ -129,7 +103,7 @@ def example_decorator_func(selenium):
     pass
 
 
-def test_local4(local_mocks):
+def test_local4():
     example_decorator_func(selenium_mock)
     assert example_decorator_func.dec_info == [
         ("testdec1", "a"),
@@ -138,18 +112,13 @@ def test_local4(local_mocks):
     ]
 
 
-def test_local5(local_mocks):
-    example_func1(selenium_mock)
-    local_mocks.check_err(AssertionError, "AssertionError: assert 6 == 7\n")
-
-
 class selenium_mock_fail_load_package(selenium_mock):
     @staticmethod
     def load_package(*args, **kwargs):
         raise OSError("STOP!")
 
 
-def test_local_fail_load_package(local_mocks):
+def test_local_fail_load_package():
     exc = None
     try:
         example_func1(selenium_mock_fail_load_package)
@@ -169,13 +138,12 @@ def test_local_fail_load_package(local_mocks):
         )
 
 
-def test_selenium(selenium, local_mocks):
-    example_func1(selenium)
+def test_selenium(selenium):
+    with pytest.raises(AssertionError, match="assert 6 == 7"):
+        example_func1(selenium)
 
-    local_mocks.check_err(AssertionError, "AssertionError: assert 6 == 7\n")
-
-    example_func2(selenium)
-    local_mocks.check_err(AssertionError, "AssertionError: assert 6 == 7\n")
+    with pytest.raises(AssertionError, match="assert 6 == 7"):
+        example_func2(selenium)
 
 
 @run_in_pyodide
diff --git a/requirements.txt b/requirements.txt
index 9c92fa740..bced41197 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,5 +15,6 @@
   pytest-rerunfailures
   pytest-xdist
   selenium==4.1.0
+  tblib
   # maintenance
   bump2version
diff --git a/run_docker b/run_docker
index 0a06af1f5..98bf2c61c 100755
--- a/run_docker
+++ b/run_docker
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 
 PYODIDE_IMAGE_REPO="pyodide"
-PYODIDE_IMAGE_TAG="20220504-py310-chrome101-firefox100"
+PYODIDE_IMAGE_TAG="20220525-py310-chrome102-firefox100"
 PYODIDE_PREBUILT_IMAGE_TAG="0.20.0"
 DEFAULT_PYODIDE_DOCKER_IMAGE="${PYODIDE_IMAGE_REPO}/pyodide-env:${PYODIDE_IMAGE_TAG}"
 DEFAULT_PYODIDE_SYSTEM_PORT="none"