Use cloudpickle instead of JSON serialization.

Author: tbenthompson

docs/changelog.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+
+- replaced JSON serialization with cloudpickle. This allows extracting a much wider range of objects from the notebook subprocess.
+
 ## 0.4.2
 
 - Documentation and CoC updates to improve developer access (Thank you PyLadies Vancouver!)

docs/usage/index.md

@@ -63,7 +63,7 @@ my_list = ['list', 'from', 'notebook']
 
 Reference objects to functions can be called with,
 
-- explicit JSON serializable values (like `dict`, `list`, `int`, `float`, `str`, `bool`, etc)
+- explicit cloudpickle-serializable values (almost all Python objects)
 - other reference objects
 
 ```{code-block} python

testbook/client.py

@@ -6,6 +6,7 @@
 from nbclient import NotebookClient
 from nbclient.exceptions import CellExecutionError
 from nbformat.v4 import new_code_cell
+import cloudpickle
 
 from .exceptions import (
     TestbookCellTagNotFoundError,
@@ -39,7 +40,7 @@ def ref(self, name: str) -> Union[TestbookObjectReference, Any]:
         # Check if exists
         self.inject(name, pop=True)
         try:
-            self.inject(f"import json; json.dumps({name})", pop=True)
+            self.inject(f"import cloudpickle; cloudpickle.dumps({name})", pop=True)
             return self.value(name)
         except Exception:
             return TestbookObjectReference(self, name)
@@ -248,9 +249,8 @@ def inject(
 
     def value(self, code: str) -> Any:
         """
-        Execute given code in the kernel and return JSON serializeable result.
+        Execute given code in the kernel and returns the serializeable result.
 
-        If the result is not JSON serializeable, it raises `TestbookAttributeError`.
         This error object will also contain an attribute called `save_varname` which
         can be used to create a reference object with :meth:`ref`.
 
@@ -277,34 +277,24 @@ def value(self, code: str) -> Any:
                 'code provided does not produce execute_result'
             )
 
-        save_varname = random_varname()
-
-        inject_code = f"""
-            import json
-            from IPython import get_ipython
-            from IPython.display import JSON
-
-            {save_varname} = get_ipython().last_execution_result.result
-
-            json.dumps({save_varname})
-            JSON({{"value" : {save_varname}}})
-        """
-
-        try:
-            outputs = self.inject(inject_code, pop=True).outputs
-
-            if outputs[0].output_type == "error":
-                # will receive error when `allow_errors` is set to True
-                raise TestbookRuntimeError(
-                    outputs[0].evalue, outputs[0].traceback, outputs[0].ename
-                )
-
-            return outputs[0].data['application/json']['value']
-
-        except TestbookRuntimeError:
-            e = TestbookSerializeError('could not JSON serialize output')
-            e.save_varname = save_varname
-            raise e
+        import tempfile
+        with tempfile.NamedTemporaryFile() as tmp:
+            try: 
+                inject_code = f"""
+                    import cloudpickle
+                    with open('{tmp.name}', 'wb') as f:
+                        cloudpickle.dump(get_ipython().last_execution_result.result, f)
+                """
+                outputs = self.inject(inject_code, pop=True).outputs
+                if len(outputs) > 0 and outputs[0].output_type == "error":
+                    # will receive error when `allow_errors` is set to True
+                    raise TestbookRuntimeError(
+                        outputs[0].evalue, outputs[0].traceback, outputs[0].ename
+                    )
+                with open(tmp.name, 'rb') as f:
+                    return cloudpickle.load(f) 
+            except TestbookRuntimeError:
+                raise TestbookSerializeError('could not serialize output')
 
     @contextmanager
     def patch(self, target, **kwargs):

testbook/exceptions.py

@@ -10,7 +10,7 @@ class TestbookCellTagNotFoundError(TestbookError):
 
 
 class TestbookSerializeError(TestbookError):
-    """Raised when output cannot be JSON serialized"""
+    """Raised when output cannot be serialized"""
 
     pass
 

testbook/tests/test_reference.py

@@ -47,24 +47,8 @@ def test_function_call_with_ref_object(notebook):
     assert double(a) == [2, 4, 6]
 
 
-def test_reference(notebook):
+def test_nontrivial_pickling(notebook):
     Foo = notebook.ref("Foo")
-
-    # Check that when a non-serializeable object is returned, it returns
-    # a reference to that object instead
-    f = Foo('bar')
-
-    assert repr(f) == "\"<Foo value='bar'>\""
-
-    # Valid attribute access
-    assert f.say_hello()
-
-    # Invalid attribute access
-    with pytest.raises(TestbookAttributeError):
-        f.does_not_exist
-
-    assert f.say_hello() == 'Hello bar!'
-
-    # non JSON-serializeable output
-    with pytest.raises(TestbookSerializeError):
-        f.resolve()
+    f = Foo("bar")
+    assert repr(f) == "<Foo value='bar'>"
+    assert(f.say_hello() == "Hello bar!")