Repository: nteract/papermill

Issue:
PapermillExecutionError cannot be loaded from picklefile
## 🐛 Bug

As part of my team's workflow, we extensively use Papermill to run notebooks as part of our Apache Airflow pipelines. When upgrading Airflow to version 2.X.X we found that our default error handling tools were no longer providing detailed errors from our PapermIll tasks. The cause for this is that Airflow changed the way it handles errors in it's tasks. It temporarily pickles the error message, and then at a later point loads the pickled exception. 

Because PapermillExecutionError is called with multiple arguments, but then only passes up a message argument to its superclass, pickle is unable to load it. See [this stack overflow question/answer](https://stackoverflow.com/questions/41808912/cannot-unpickle-exception-subclass) for a general case of this issue.

This error can be recreated with this simple test:
```{python}
import tempfile
import pickle

from papermill.exceptions import PapermillExecutionError

def test_exceptions_are_unpickleable():
    """Ensure exceptions can be unpickled"""
    temp_file = NamedTemporaryFile()
    err = PapermillExecutionError(1,2, "TestSource", "Exception", Exception(), ["Traceback", "message"])

    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)
```

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract) Report
> Merging [#629](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract) (df523ba) into [main](https://codecov.io/gh/nteract/papermill/commit/3f6225ed51e2d41904424059c95306b65be5fb81?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract) (3f6225e) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/nteract/papermill/pull/629/graphs/tree.svg?width=650&height=150&src=pr&token=X0WnER07g1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract)](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract)

```diff
@@            Coverage Diff             @@
##             main     #629      +/-   ##
==========================================
+ Coverage   91.63%   91.65%   +0.01%     
==========================================
  Files          17       17              
  Lines        1555     1558       +3     
==========================================
+ Hits         1425     1428       +3     
  Misses        130      130              
```


| [Impacted Files](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract) | Coverage Δ | |
|---|---|---|
| [papermill/exceptions.py](https://codecov.io/gh/nteract/papermill/pull/629/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract#diff-cGFwZXJtaWxsL2V4Y2VwdGlvbnMucHk=) | `87.87% <100.00%> (+1.21%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract). Last update [3f6225e...df523ba](https://codecov.io/gh/nteract/papermill/pull/629?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=nteract).


Patch:
diff --git a/papermill/exceptions.py b/papermill/exceptions.py
index ac3c49f..6a45a8c 100644
--- a/papermill/exceptions.py
+++ b/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):


Test patch:
diff --git a/papermill/tests/test_exceptions.py b/papermill/tests/test_exceptions.py
index 45027c7..2c6968c 100644
--- a/papermill/tests/test_exceptions.py
+++ b/papermill/tests/test_exceptions.py
@@ -1,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)