From af62f4c6747b561e6d0ccd0305021d8415f400d7 Mon Sep 17 00:00:00 2001
From: Joshua Harlow <harlowja@yahoo-inc.com>
Date: Fri, 12 Sep 2014 18:49:29 -0700
Subject: [PATCH] Ensure that failures can be pickled

When a failure happens in a subprocess it needs to be
pickleable (and not contain a traceback, since those can
not be pickled) when being sent across the process boundary
so that the receiving process can read it and unpickle
the corresponding object.

Part of blueprint process-executor

Change-Id: I2f26faa4e02da6acf4f0840239d0b17143de8d76
---
 taskflow/types/failure.py | 41 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/taskflow/types/failure.py b/taskflow/types/failure.py
index a0b3a17e..b9d7a399 100644
--- a/taskflow/types/failure.py
+++ b/taskflow/types/failure.py
@@ -34,6 +34,23 @@ def _copy_exc_info(exc_info):
     return (exc_type, copy.copy(exc_value), tb)
 
 
+def _fill_iter(it, desired_len, filler=None):
+    """Iterates over a provided iterator up to the desired length.
+
+    If the source iterator does not have enough values then the filler
+    value is yielded until the desired length is reached.
+    """
+    count = 0
+    for value in it:
+        if count >= desired_len:
+            return
+        yield value
+        count += 1
+    while count < desired_len:
+        yield filler
+        count += 1
+
+
 def _are_equal_exc_info_tuples(ei1, ei2):
     if ei1 == ei2:
         return True
@@ -274,6 +291,30 @@ class Failure(object):
         for et in self._exc_type_names:
             yield et
 
+    def __getstate__(self):
+        dct = self.to_dict()
+        if self._exc_info:
+            # Avoids 'TypeError: can't pickle traceback objects'
+            dct['exc_info'] = self._exc_info[0:2]
+        return dct
+
+    def __setstate__(self, dct):
+        self._exception_str = dct['exception_str']
+        self._traceback_str = dct['traceback_str']
+        self._exc_type_names = dct['exc_type_names']
+        if 'exc_info' in dct:
+            # Tracebacks can't be serialized/deserialized, but since we
+            # provide a traceback string (and more) this should be
+            # acceptable...
+            #
+            # TODO(harlowja): in the future we could do something like
+            # what the twisted people have done, see for example
+            # twisted-13.0.0/twisted/python/failure.py#L89 for how they
+            # created a fake traceback object...
+            self._exc_info = tuple(_fill_iter(dct['exc_info'], 3))
+        else:
+            self._exc_info = None
+
     @classmethod
     def from_dict(cls, data):
         """Converts this from a dictionary to a object."""