change ObjectId __getstate__ and __setstate__ to just use __id also added unit test for pickling/unpickling ObjectId

shearic · behackett · commit f202c535d552 · 2011-04-02T11:10:26.000-07:00
diff --git a/bson/objectid.py b/bson/objectid.py
@@ -191,12 +191,12 @@ def __getstate__(self):
         """return value of object for pickling.
         needed explicitly because __slots__() defined.
         """
-        return self.__str__()
+        return self.__id
 
-    def __setstate__(self, oid):
+    def __setstate__(self, value):
         """explicit state set from pickling
         """
-        self.__validate(oid)
+        self.__id = value
 
     def __str__(self):
         return self.__id.encode("hex")
diff --git a/test/test_objectid.py b/test/test_objectid.py
@@ -15,6 +15,7 @@
 """Tests for the objectid module."""
 
 import datetime
+import pickle
 import warnings
 import unittest
 import sys
@@ -124,5 +125,9 @@ def test_from_datetime(self):
         oid = ObjectId.from_datetime(aware)
         self.assertEqual(as_utc, oid.generation_time)
 
+    def test_pickling(self):
+        orig = ObjectId()
+        self.assertEqual(orig, pickle.loads(pickle.dumps(orig)))
+
 if __name__ == "__main__":
     unittest.main()
