Bugfix/pickle errors (#629)

* Add tests for showing unpickling error

* Add fix so exceptions can be loaded from pickle

This fixes an issue where the PapermillExecutionError can be pickled but will
not be unpicklable.

Pythons BaseException class has a handling to make sure that exceptions are able
to be loaded once they are pickled. This works by ensuring the args used to generate the instance
are saved on the exception. However by only passing the message to the superclass, the pickle
module was unable to load a new instance of PapermillExecutionError from a file.

A new __str__ method was added to preserve the functionality of calling str() on the exception and
only getting the error message back, while ensuring the superclass had access to all instantiating
arguments.

Co-authored-by: Tyler Calder <calder-ty@protonmail.com>
Co-authored-by: Matthew Seal <mseal007@gmail.com>

papermill/exceptions.py

@@ -23,18 +23,26 @@ class PapermillExecutionError(PapermillException):
     """Raised when an exception is encountered in a notebook."""
 
     def __init__(self, cell_index, exec_count, source, ename, evalue, traceback):
+        args = cell_index, exec_count, source, ename, evalue, traceback
         self.cell_index = cell_index
         self.exec_count = exec_count
         self.source = source
         self.ename = ename
         self.evalue = evalue
         self.traceback = traceback
+
+        super(PapermillExecutionError, self).__init__(*args)
+
+    def __str__(self):
+        # Standard Behavior of an exception is to produce a string representation of its arguments
+        # when called with str(). In order to maintain compatability with previous versions which
+        # passed only the message to the superclass constructor, __str__ method is implemented to
+        # provide the same result as was produced in the past.
         message = "\n" + 75 * "-" + "\n"
-        message += 'Exception encountered at "In [%s]":\n' % str(exec_count)
-        message += strip_color("\n".join(traceback))
+        message += 'Exception encountered at "In [%s]":\n' % str(self.exec_count)
+        message += strip_color("\n".join(self.traceback))
         message += "\n"
-
-        super(PapermillExecutionError, self).__init__(message)
+        return message
 
 
 class PapermillRateLimitException(PapermillException):

papermill/tests/test_exceptions.py

@@ -1 +1,43 @@
+import pickle
+import tempfile
+
 import pytest  # noqa
+
+from .. import exceptions
+
+
+@pytest.fixture
+def temp_file():
+    f = tempfile.NamedTemporaryFile()
+    yield f
+    f.close()
+
+
+@pytest.mark.parametrize(
+    "exc,args",
+    [
+        (
+            exceptions.PapermillExecutionError,
+            (1, 2, "TestSource", "Exception", Exception(), ["Traceback", "Message"]),
+        ),
+        (exceptions.PapermillMissingParameterException, ("PapermillMissingParameterException",)),
+        (exceptions.AwsError, ("AwsError",)),
+        (exceptions.FileExistsError, ("FileExistsError",)),
+        (exceptions.PapermillException, ("PapermillException",)),
+        (exceptions.PapermillRateLimitException, ("PapermillRateLimitException",)),
+        (
+            exceptions.PapermillOptionalDependencyException,
+            ("PapermillOptionalDependencyException",),
+        ),
+    ],
+)
+def test_exceptions_are_unpickleable(temp_file, exc, args):
+    """Ensure exceptions can be unpickled"""
+    err = exc(*args)
+    with open(temp_file.name, 'wb') as fd:
+        pickle.dump(err, fd)
+    # Read the Pickled File
+    temp_file.seek(0)
+    data = temp_file.read()
+    pickled_err = pickle.loads(data)
+    assert str(pickled_err) == str(err)