fix duration pickle for all protocols and all supported python versions and variants. (closes #10)

2 files changed, 28 insertions, 4 deletions

diff --git a/src/isodate/duration.py b/src/isodate/duration.py
index 62b14ca..1aa623d 100644
--- a/src/isodate/duration.py
+++ b/src/isodate/duration.py
@@ -101,6 +101,12 @@ class Duration(object):
         self.tdelta = timedelta(days, seconds, microseconds, milliseconds,
                                 minutes, hours, weeks)
 
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
     def __getattr__(self, name):
         '''
         Provide direct access to attributes of included timedelta instance.
diff --git a/src/isodate/tests/test_pickle.py b/src/isodate/tests/test_pickle.py
index 867bc40..b52f8cb 100644
--- a/src/isodate/tests/test_pickle.py
+++ b/src/isodate/tests/test_pickle.py
@@ -9,14 +9,32 @@ class TestPickle(unittest.TestCase):
     datetime object.
     '''
 
-    def test_pickle(self):
+    def test_pickle_datetime(self):
         '''
         Parse an ISO datetime string and compare it to the expected value.
         '''
         dti = isodate.parse_datetime('2012-10-26T09:33+00:00')
-        pikl = pickle.dumps(dti, 2)
-        dto = pickle.loads(pikl)
-        self.assertEqual(dti, dto)
+        for proto in range(0, pickle.HIGHEST_PROTOCOL + 1):
+            pikl = pickle.dumps(dti, proto)
+            self.assertEqual(dti, pickle.loads(pikl),
+                             "pickle proto %d failed" % proto)
+
+    def test_pickle_duration(self):
+        '''
+        Pickle / unpickle duration objects.
+        '''
+        from isodate.duration import Duration
+        dur = Duration()
+        failed = []
+        for proto in range(0, pickle.HIGHEST_PROTOCOL + 1):
+            try:
+                pikl = pickle.dumps(dur, proto)
+                if dur != pickle.loads(pikl):
+                    raise Exception("not equal")
+            except Exception, e:
+                failed.append("pickle proto %d failed (%s)" % (proto, repr(e)))
+        self.assertEqual(len(failed), 0, "pickle protos failed: %s" %
+                         str(failed))
 
 
 def test_suite():